1pub mod manifest;
10pub mod manifest_feature_flags;
11
12use arrow::array::Float32Array;
13use arrow::record_batch::RecordBatchIterator;
14use arrow_ipc::reader::StreamReader;
15use async_trait::async_trait;
16use bytes::Bytes;
17use futures::{StreamExt, TryStreamExt};
18use lance::dataset::builder::DatasetBuilder;
19use lance::dataset::refs::check_valid_branch;
20use lance::dataset::scanner::Scanner;
21use lance::dataset::statistics::DatasetStatisticsExt;
22use lance::dataset::transaction::{Operation, Transaction};
23use lance::dataset::{
24 Dataset, MergeInsertBuilder, UpdateBuilder, WhenMatched, WhenNotMatched,
25 WhenNotMatchedBySource, WriteMode, WriteParams,
26};
27use lance::index::{DatasetIndexExt, IndexParams, vector::VectorIndexParams};
28use lance::session::Session;
29use lance_index::scalar::{
30 BuiltinIndexType, FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams,
31};
32use lance_index::vector::{
33 bq::{RABIT_MAX_NUM_BITS, RABIT_MIN_NUM_BITS, RQBuildParams, validate_supported_rq_num_bits},
34 hnsw::builder::HnswBuildParams,
35 ivf::IvfBuildParams,
36 pq::PQBuildParams,
37 sq::builder::SQBuildParams,
38};
39use lance_index::{IndexType, is_system_index};
40use lance_io::object_store::throttle::is_throttle_error;
41use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry, ReadDirOptions};
42use lance_linalg::distance::MetricType;
43use lance_table::io::commit::{ManifestNamingScheme, VERSIONS_DIR};
44use object_store::ObjectStoreExt;
45use object_store::path::Path;
46use object_store::{
47 Error as ObjectStoreError, ObjectMeta, ObjectStore as OSObjectStore, PutMode, PutOptions,
48};
49use std::collections::HashMap;
50use std::io::Cursor;
51use std::sync::{Arc, Mutex};
52use tokio::sync::OnceCell;
53
54use crate::context::DynamicContextProvider;
55use crate::merge_insert_on_columns;
56use lance_namespace::models::{
57 AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest,
58 AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse,
59 AlterTransactionRequest, AlterTransactionResponse, AnalyzeTableQueryPlanRequest,
60 BatchDeleteTableVersionsRequest, BatchDeleteTableVersionsResponse,
61 BranchContents as ModelBranchContents, CountTableRowsRequest, CreateNamespaceRequest,
62 CreateNamespaceResponse, CreateTableBranchRequest, CreateTableBranchResponse,
63 CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest, CreateTableResponse,
64 CreateTableScalarIndexResponse, CreateTableTagRequest, CreateTableTagResponse,
65 CreateTableVersionRequest, CreateTableVersionResponse, DeclareTableRequest,
66 DeclareTableResponse, DeleteFromTableRequest, DeleteFromTableResponse,
67 DeleteTableBranchRequest, DeleteTableBranchResponse, DeleteTableTagRequest,
68 DeleteTableTagResponse, DescribeNamespaceRequest, DescribeNamespaceResponse,
69 DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, DescribeTableRequest,
70 DescribeTableResponse, DescribeTableVersionRequest, DescribeTableVersionResponse,
71 DescribeTransactionRequest, DescribeTransactionResponse, DropNamespaceRequest,
72 DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, DropTableRequest,
73 DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary,
74 GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest,
75 GetTableTagVersionResponse, Identity, IndexContent, InsertIntoTableRequest,
76 InsertIntoTableResponse, ListNamespacesRequest, ListNamespacesResponse,
77 ListTableBranchesRequest, ListTableBranchesResponse, ListTableIndicesRequest,
78 ListTableIndicesResponse, ListTableTagsRequest, ListTableTagsResponse,
79 ListTableVersionsRequest, ListTableVersionsResponse, ListTablesRequest, ListTablesResponse,
80 MergeInsertIntoTableRequest, MergeInsertIntoTableResponse, NamespaceExistsRequest,
81 QueryTableRequest, QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest,
82 RestoreTableResponse, TableExistsRequest, TableVersion, TagContents as ModelTagContents,
83 UpdateTableRequest, UpdateTableResponse, UpdateTableSchemaMetadataRequest,
84 UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, UpdateTableTagResponse,
85};
86
87use lance_core::utils::parse::str_to_bool;
88use lance_core::{Error, Result, box_error};
89use lance_index::scalar::inverted::query::{
90 BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery,
91};
92use lance_namespace::LanceNamespace;
93use lance_namespace::error::NamespaceError;
94use lance_namespace::schema::arrow_schema_to_json;
95
96use crate::credentials::{
97 CredentialVendor, create_credential_vendor_for_location, has_credential_vendor_config,
98};
99
100#[derive(Debug, Default)]
105pub struct OpsMetrics {
106 counters: Mutex<HashMap<String, u64>>,
107}
108
109impl OpsMetrics {
110 pub fn increment(&self, operation: &str) {
112 if let Ok(mut counters) = self.counters.lock() {
113 *counters.entry(operation.to_string()).or_insert(0) += 1;
114 }
115 }
116
117 pub fn retrieve(&self) -> HashMap<String, u64> {
119 self.counters.lock().map(|c| c.clone()).unwrap_or_default()
120 }
121
122 pub fn reset(&self) {
124 if let Ok(mut counters) = self.counters.lock() {
125 counters.clear();
126 }
127 }
128}
129
130pub(crate) fn build_sql_expressions(
133 new_columns: &[lance_namespace::models::AddColumnsEntry],
134) -> Result<Vec<(String, String)>> {
135 new_columns
136 .iter()
137 .map(|col| {
138 let expression = col.expression.clone().and_then(|opt| opt).ok_or_else(|| {
141 Error::invalid_input(format!(
142 "Expression is required for new column '{}'",
143 col.name
144 ))
145 })?;
146 Ok((col.name.clone(), expression))
147 })
148 .collect()
149}
150
151pub(crate) fn build_column_alterations(
154 alterations: &[lance_namespace::models::AlterColumnsEntry],
155) -> Result<Vec<lance::dataset::ColumnAlteration>> {
156 alterations
157 .iter()
158 .map(|entry| {
159 let mut alteration = lance::dataset::ColumnAlteration::new(entry.path.clone());
160 if let Some(Some(rename)) = &entry.rename {
162 alteration = alteration.rename(rename.clone());
163 }
164 if let Some(Some(nullable)) = entry.nullable {
166 alteration = alteration.set_nullable(nullable);
167 }
168 if let Some(data_type) = &entry.data_type
170 && !data_type.is_null()
171 {
172 let type_str = data_type.as_str().ok_or_else(|| {
173 Error::invalid_input(format!(
174 "data_type for column '{}' must be a JSON string, got: {}",
175 entry.path, data_type
176 ))
177 })?;
178 let json_type =
179 lance_namespace::models::JsonArrowDataType::new(type_str.to_string());
180 let dt =
181 lance_namespace::schema::convert_json_arrow_type(&json_type).map_err(|e| {
182 Error::invalid_input(format!(
183 "Failed to parse data_type '{}' for column '{}': {}",
184 type_str, entry.path, e
185 ))
186 })?;
187 alteration = alteration.cast_to(dt);
188 }
189 Ok(alteration)
190 })
191 .collect()
192}
193
194pub(crate) struct TableStatus {
199 pub(crate) exists: bool,
201 pub(crate) is_deregistered: bool,
203 pub(crate) has_reserved_file: bool,
205}
206
207enum DirectoryIndexParams {
208 Scalar {
209 index_type: IndexType,
210 params: ScalarIndexParams,
211 },
212 Inverted(InvertedIndexParams),
213 Vector {
214 index_type: IndexType,
215 params: VectorIndexParams,
216 },
217}
218
219impl DirectoryIndexParams {
220 fn index_type(&self) -> IndexType {
221 match self {
222 Self::Scalar { index_type, .. } | Self::Vector { index_type, .. } => *index_type,
223 Self::Inverted(_) => IndexType::Inverted,
224 }
225 }
226
227 fn params(&self) -> &dyn IndexParams {
228 match self {
229 Self::Scalar { params, .. } => params,
230 Self::Inverted(params) => params,
231 Self::Vector { params, .. } => params,
232 }
233 }
234}
235
236#[derive(Clone)]
271pub struct DirectoryNamespaceBuilder {
272 root: String,
273 storage_options: Option<HashMap<String, String>>,
274 session: Option<Arc<Session>>,
275 manifest_enabled: bool,
276 dir_listing_enabled: bool,
277 inline_optimization_enabled: bool,
278 table_version_tracking_enabled: bool,
279 dir_listing_to_manifest_migration_enabled: bool,
284 credential_vendor_properties: HashMap<String, String>,
285 context_provider: Option<Arc<dyn DynamicContextProvider>>,
286 commit_retries: Option<u32>,
287 vend_input_storage_options: bool,
290 vend_input_storage_options_refresh_interval_millis: Option<u64>,
295 ops_metrics_enabled: bool,
297}
298
299impl std::fmt::Debug for DirectoryNamespaceBuilder {
300 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
301 f.debug_struct("DirectoryNamespaceBuilder")
302 .field("root", &self.root)
303 .field("storage_options", &self.storage_options)
304 .field("manifest_enabled", &self.manifest_enabled)
305 .field("dir_listing_enabled", &self.dir_listing_enabled)
306 .field(
307 "inline_optimization_enabled",
308 &self.inline_optimization_enabled,
309 )
310 .field(
311 "table_version_tracking_enabled",
312 &self.table_version_tracking_enabled,
313 )
314 .field(
315 "dir_listing_to_manifest_migration_enabled",
316 &self.dir_listing_to_manifest_migration_enabled,
317 )
318 .field(
319 "context_provider",
320 &self.context_provider.as_ref().map(|_| "Some(...)"),
321 )
322 .field(
323 "vend_input_storage_options",
324 &self.vend_input_storage_options,
325 )
326 .field(
327 "vend_input_storage_options_refresh_interval_millis",
328 &self.vend_input_storage_options_refresh_interval_millis,
329 )
330 .field("ops_metrics_enabled", &self.ops_metrics_enabled)
331 .finish()
332 }
333}
334
335impl DirectoryNamespaceBuilder {
336 pub fn new(root: impl Into<String>) -> Self {
342 Self {
343 root: root.into().trim_end_matches('/').to_string(),
344 storage_options: None,
345 session: None,
346 manifest_enabled: true,
347 dir_listing_enabled: true, inline_optimization_enabled: false,
349 table_version_tracking_enabled: false, dir_listing_to_manifest_migration_enabled: false, credential_vendor_properties: HashMap::new(),
352 context_provider: None,
353 commit_retries: None,
354 vend_input_storage_options: false,
355 vend_input_storage_options_refresh_interval_millis: None,
356 ops_metrics_enabled: false,
357 }
358 }
359
360 pub fn manifest_enabled(mut self, enabled: bool) -> Self {
365 self.manifest_enabled = enabled;
366 self
367 }
368
369 pub fn dir_listing_enabled(mut self, enabled: bool) -> Self {
374 self.dir_listing_enabled = enabled;
375 self
376 }
377
378 pub fn dir_listing_to_manifest_migration_enabled(mut self, enabled: bool) -> Self {
385 self.dir_listing_to_manifest_migration_enabled = enabled;
386 self
387 }
388
389 pub fn inline_optimization_enabled(mut self, enabled: bool) -> Self {
394 self.inline_optimization_enabled = enabled;
395 self
396 }
397
398 pub fn table_version_tracking_enabled(mut self, enabled: bool) -> Self {
406 self.table_version_tracking_enabled = enabled;
407 self
408 }
409
410 pub fn from_properties(
478 properties: HashMap<String, String>,
479 session: Option<Arc<Session>>,
480 ) -> Result<Self> {
481 let root = properties.get("root").cloned().ok_or_else(|| {
483 lance_core::Error::from(NamespaceError::InvalidInput {
484 message: "Missing required property 'root' for directory namespace".to_string(),
485 })
486 })?;
487
488 let storage_options: HashMap<String, String> = properties
490 .iter()
491 .filter_map(|(k, v)| {
492 k.strip_prefix("storage.")
493 .map(|key| (key.to_string(), v.clone()))
494 })
495 .collect();
496
497 let storage_options = if storage_options.is_empty() {
498 None
499 } else {
500 Some(storage_options)
501 };
502
503 let manifest_enabled = properties
505 .get("manifest_enabled")
506 .and_then(|v| str_to_bool(v))
507 .unwrap_or(true);
508
509 let dir_listing_enabled = properties
511 .get("dir_listing_enabled")
512 .and_then(|v| str_to_bool(v))
513 .unwrap_or(true);
514
515 let inline_optimization_enabled = properties
517 .get("inline_optimization_enabled")
518 .and_then(|v| str_to_bool(v))
519 .unwrap_or(false);
520
521 let table_version_tracking_enabled = properties
523 .get("table_version_tracking_enabled")
524 .and_then(|v| str_to_bool(v))
525 .unwrap_or(false);
526
527 let dir_listing_to_manifest_migration_enabled = properties
529 .get("dir_listing_to_manifest_migration_enabled")
530 .and_then(|v| str_to_bool(v))
531 .unwrap_or(false);
532
533 let credential_vendor_properties: HashMap<String, String> = properties
537 .iter()
538 .filter_map(|(k, v)| {
539 k.strip_prefix("credential_vendor.")
540 .map(|key| (key.to_string(), v.clone()))
541 })
542 .collect();
543
544 let commit_retries = properties
545 .get("commit_retries")
546 .and_then(|v| v.parse::<u32>().ok());
547
548 let vend_input_storage_options = properties
550 .get("vend_input_storage_options")
551 .and_then(|v| str_to_bool(v))
552 .unwrap_or(false);
553
554 let vend_input_storage_options_refresh_interval_millis = properties
556 .get("vend_input_storage_options_refresh_interval_millis")
557 .and_then(|v| v.parse::<u64>().ok());
558
559 let ops_metrics_enabled = properties
561 .get("ops_metrics_enabled")
562 .and_then(|v| str_to_bool(v))
563 .unwrap_or(false);
564
565 Ok(Self {
566 root: root.trim_end_matches('/').to_string(),
567 storage_options,
568 session,
569 manifest_enabled,
570 dir_listing_enabled,
571 inline_optimization_enabled,
572 table_version_tracking_enabled,
573 dir_listing_to_manifest_migration_enabled,
574 credential_vendor_properties,
575 context_provider: None,
576 commit_retries,
577 vend_input_storage_options,
578 vend_input_storage_options_refresh_interval_millis,
579 ops_metrics_enabled,
580 })
581 }
582
583 pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
590 self.storage_options
591 .get_or_insert_with(HashMap::new)
592 .insert(key.into(), value.into());
593 self
594 }
595
596 pub fn storage_options(mut self, options: HashMap<String, String>) -> Self {
602 self.storage_options
603 .get_or_insert_with(HashMap::new)
604 .extend(options);
605 self
606 }
607
608 pub fn session(mut self, session: Arc<Session>) -> Self {
618 self.session = Some(session);
619 self
620 }
621
622 pub fn commit_retries(mut self, retries: u32) -> Self {
625 self.commit_retries = Some(retries);
626 self
627 }
628
629 pub fn credential_vendor_property(
657 mut self,
658 key: impl Into<String>,
659 value: impl Into<String>,
660 ) -> Self {
661 self.credential_vendor_properties
662 .insert(key.into(), value.into());
663 self
664 }
665
666 pub fn credential_vendor_properties(mut self, properties: HashMap<String, String>) -> Self {
674 self.credential_vendor_properties.extend(properties);
675 self
676 }
677
678 pub fn context_provider(mut self, provider: Arc<dyn DynamicContextProvider>) -> Self {
688 self.context_provider = Some(provider);
689 self
690 }
691
692 pub fn vend_input_storage_options(mut self, enabled: bool) -> Self {
701 self.vend_input_storage_options = enabled;
702 self
703 }
704
705 pub fn vend_input_storage_options_refresh_interval_millis(
717 mut self,
718 interval_millis: u64,
719 ) -> Self {
720 self.vend_input_storage_options_refresh_interval_millis = Some(interval_millis);
721 self
722 }
723
724 pub fn ops_metrics_enabled(mut self, enabled: bool) -> Self {
732 self.ops_metrics_enabled = enabled;
733 self
734 }
735
736 pub async fn build(self) -> Result<DirectoryNamespace> {
749 let (object_store, base_path) =
750 Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?;
751
752 let manifest_ns = if self.manifest_enabled {
753 match manifest::ManifestNamespace::open_from_directory(
754 self.root.clone(),
755 self.storage_options.clone(),
756 self.session.clone(),
757 object_store.clone(),
758 base_path.clone(),
759 self.dir_listing_enabled,
760 self.inline_optimization_enabled,
761 self.commit_retries,
762 )
763 .await
764 {
765 Ok(ns) => Some(Arc::new(ns)),
766 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
767 return Err(e);
771 }
772 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => {
773 log::debug!("Manifest namespace does not exist yet: {}", e);
774 None
775 }
776 Err(e) => return Err(e),
777 }
778 } else {
779 None
780 };
781 let manifest_cell = OnceCell::new();
782 if let Some(manifest_ns) = manifest_ns {
783 let _ = manifest_cell.set(manifest_ns);
784 }
785
786 let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties)
788 {
789 create_credential_vendor_for_location(&self.root, &self.credential_vendor_properties)
790 .await?
791 .map(Arc::from)
792 } else {
793 None
794 };
795
796 let ops_metrics = if self.ops_metrics_enabled {
797 Some(Arc::new(OpsMetrics::default()))
798 } else {
799 None
800 };
801
802 Ok(DirectoryNamespace {
803 root: self.root,
804 storage_options: self.storage_options,
805 session: self.session,
806 object_store,
807 base_path,
808 manifest_ns: manifest_cell,
809 write_manifest_ns: OnceCell::new(),
810 manifest_enabled: self.manifest_enabled,
811 dir_listing_enabled: self.dir_listing_enabled,
812 inline_optimization_enabled: self.inline_optimization_enabled,
813 commit_retries: self.commit_retries,
814 dir_listing_to_manifest_migration_enabled: self
815 .dir_listing_to_manifest_migration_enabled,
816 table_version_tracking_enabled: self.table_version_tracking_enabled,
817 credential_vendor,
818 context_provider: self.context_provider,
819 vend_input_storage_options: self.vend_input_storage_options,
820 vend_input_storage_options_refresh_interval_millis: self
821 .vend_input_storage_options_refresh_interval_millis,
822 ops_metrics,
823 })
824 }
825
826 async fn initialize_object_store(
828 root: &str,
829 storage_options: &Option<HashMap<String, String>>,
830 session: &Option<Arc<Session>>,
831 ) -> Result<(Arc<ObjectStore>, Path)> {
832 let accessor = storage_options.clone().map(|opts| {
834 Arc::new(lance_io::object_store::StorageOptionsAccessor::with_static_options(opts))
835 });
836 let params = ObjectStoreParams {
837 storage_options_accessor: accessor,
838 ..Default::default()
839 };
840
841 let registry = if let Some(session) = session {
843 session.store_registry()
844 } else {
845 Arc::new(ObjectStoreRegistry::default())
846 };
847
848 let (object_store, base_path) = ObjectStore::from_uri_and_params(registry, root, ¶ms)
850 .await
851 .map_err(|e| {
852 lance_core::Error::from(NamespaceError::Internal {
853 message: format!("Failed to create object store: {:?}", e),
854 })
855 })?;
856
857 Ok((object_store, base_path))
858 }
859}
860
861pub struct DirectoryNamespace {
885 root: String,
886 storage_options: Option<HashMap<String, String>>,
887 session: Option<Arc<Session>>,
888 object_store: Arc<ObjectStore>,
889 base_path: Path,
890 manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
891 write_manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
892 manifest_enabled: bool,
893 dir_listing_enabled: bool,
894 inline_optimization_enabled: bool,
895 commit_retries: Option<u32>,
896 dir_listing_to_manifest_migration_enabled: bool,
900 table_version_tracking_enabled: bool,
903 credential_vendor: Option<Arc<dyn CredentialVendor>>,
906 #[allow(dead_code)]
909 context_provider: Option<Arc<dyn DynamicContextProvider>>,
910 vend_input_storage_options: bool,
912 vend_input_storage_options_refresh_interval_millis: Option<u64>,
915 ops_metrics: Option<Arc<OpsMetrics>>,
917}
918
919impl std::fmt::Debug for DirectoryNamespace {
920 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
921 write!(f, "{}", self.namespace_id())
922 }
923}
924
925impl std::fmt::Display for DirectoryNamespace {
926 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
927 write!(f, "{}", self.namespace_id())
928 }
929}
930
931struct ExistingTableVersionResolve<'a> {
933 staging_path: &'a Path,
934 final_path: &'a Path,
935 version: u64,
936 table_uri: &'a str,
937 final_meta: &'a ObjectMeta,
938 request_manifest_size: Option<i64>,
939}
940
941struct TableDeleteEntry {
944 table_id: Option<Vec<String>>,
945 ranges: Vec<(i64, i64)>,
946}
947
948#[derive(Debug, Clone, Default)]
958struct TransactionAlteration {
959 status: Option<String>,
961 properties: HashMap<String, String>,
965 removed_properties: std::collections::HashSet<String>,
969}
970
971impl TransactionAlteration {
972 const F_STATUS: &'static str = "status";
974 const F_PROPERTIES: &'static str = "properties";
975 const F_REMOVED_PROPERTIES: &'static str = "removed_properties";
976
977 fn to_json_bytes(&self) -> serde_json::Result<Vec<u8>> {
984 serde_json::to_vec(&serde_json::json!({
985 Self::F_STATUS: self.status,
986 Self::F_PROPERTIES: self.properties,
987 Self::F_REMOVED_PROPERTIES: self.removed_properties,
988 }))
989 }
990
991 fn from_json_slice(bytes: &[u8]) -> serde_json::Result<Self> {
996 let mut obj: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes)?;
997 Ok(Self {
998 status: serde_json::from_value(
999 obj.remove(Self::F_STATUS)
1000 .unwrap_or(serde_json::Value::Null),
1001 )?,
1002 properties: serde_json::from_value(
1003 obj.remove(Self::F_PROPERTIES)
1004 .unwrap_or(serde_json::Value::Null),
1005 )
1006 .unwrap_or_default(),
1007 removed_properties: serde_json::from_value(
1008 obj.remove(Self::F_REMOVED_PROPERTIES)
1009 .unwrap_or(serde_json::Value::Null),
1010 )
1011 .unwrap_or_default(),
1012 })
1013 }
1014}
1015
1016impl DirectoryNamespace {
1017 fn manifest_ns_for_read(&self) -> Option<&Arc<manifest::ManifestNamespace>> {
1018 self.write_manifest_ns
1019 .get()
1020 .or_else(|| self.manifest_ns.get())
1021 }
1022
1023 async fn manifest_ns_for_write(&self) -> Result<Option<Arc<manifest::ManifestNamespace>>> {
1024 if !self.manifest_enabled {
1025 return Ok(None);
1026 }
1027
1028 let manifest_ns = self
1029 .write_manifest_ns
1030 .get_or_try_init(|| async {
1031 manifest::ManifestNamespace::from_directory(
1032 self.root.clone(),
1033 self.storage_options.clone(),
1034 self.session.clone(),
1035 self.object_store.clone(),
1036 self.base_path.clone(),
1037 self.dir_listing_enabled,
1038 self.inline_optimization_enabled,
1039 self.commit_retries,
1040 )
1041 .await
1042 .map(Arc::new)
1043 })
1044 .await?;
1045 Ok(Some(manifest_ns.clone()))
1046 }
1047
1048 async fn ensure_read_manifest(&self) -> Result<()> {
1060 if !self.manifest_enabled
1061 || self.manifest_ns.get().is_some()
1062 || self.write_manifest_ns.get().is_some()
1063 {
1064 return Ok(());
1065 }
1066 match self
1067 .manifest_ns
1068 .get_or_try_init(|| async {
1069 manifest::ManifestNamespace::open_from_directory(
1070 self.root.clone(),
1071 self.storage_options.clone(),
1072 self.session.clone(),
1073 self.object_store.clone(),
1074 self.base_path.clone(),
1075 self.dir_listing_enabled,
1076 self.inline_optimization_enabled,
1077 self.commit_retries,
1078 )
1079 .await
1080 .map(Arc::new)
1081 })
1082 .await
1083 {
1084 Ok(_) => Ok(()),
1085 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(()),
1090 Err(e) => Err(e),
1091 }
1092 }
1093
1094 fn child_namespace_requires_manifest_error(&self) -> Error {
1095 if self.manifest_enabled {
1096 NamespaceError::NamespaceNotFound {
1097 message: "Child namespace reads require an existing __manifest dataset".to_string(),
1098 }
1099 .into()
1100 } else {
1101 NamespaceError::Unsupported {
1102 message: "Child namespaces are only supported when manifest mode is enabled"
1103 .to_string(),
1104 }
1105 .into()
1106 }
1107 }
1108
1109 fn apply_pagination(
1122 names: &mut Vec<String>,
1123 page_token: Option<String>,
1124 limit: Option<i32>,
1125 ) -> Option<String> {
1126 names.sort();
1128
1129 if let Some(start_after) = page_token {
1131 if let Some(index) = names
1132 .iter()
1133 .position(|name| name.as_str() > start_after.as_str())
1134 {
1135 names.drain(0..index);
1136 } else {
1137 names.clear();
1138 }
1139 }
1140
1141 if let Some(limit) = limit
1143 && limit >= 0
1144 {
1145 let limit = limit as usize;
1146 if names.len() > limit {
1147 let next_page_token = if limit > 0 {
1148 Some(names[limit - 1].clone())
1149 } else {
1150 None
1151 };
1152 names.truncate(limit);
1153 return next_page_token;
1154 }
1155 }
1156
1157 None
1158 }
1159
1160 const LIST_DIRECTORY_PAGE_SIZE: usize = 1000;
1164
1165 async fn list_directory_tables(&self) -> Result<Vec<String>> {
1167 let mut tables = Vec::new();
1168 let mut page_token = None;
1169
1170 loop {
1171 let page = self
1172 .object_store
1173 .read_dir_page(
1174 self.base_path.clone(),
1175 ReadDirOptions {
1176 page_token,
1177 limit: Some(Self::LIST_DIRECTORY_PAGE_SIZE),
1182 },
1183 )
1184 .await
1185 .map_err(|e| {
1186 lance_core::Error::from(NamespaceError::Internal {
1187 message: format!("Failed to list directory: {:?}", e),
1188 })
1189 })?;
1190
1191 let candidates: Vec<String> = page
1192 .result
1193 .common_prefixes
1194 .iter()
1195 .chain(page.result.objects.iter().map(|o| &o.location))
1196 .filter_map(|p| {
1197 p.filename()?
1198 .trim_end_matches('/')
1199 .strip_suffix(".lance")
1200 .map(|name| name.to_string())
1201 })
1202 .collect();
1203
1204 let mut stream =
1208 futures::stream::iter(candidates.into_iter().map(|table_name| async move {
1209 let status = self.check_table_status(&table_name).await?;
1210 Ok::<Option<String>, Error>((!status.is_deregistered).then_some(table_name))
1211 }))
1212 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1213
1214 while let Some(result) = stream.next().await {
1215 if let Some(table_name) = result? {
1216 tables.push(table_name);
1217 }
1218 }
1219
1220 page_token = page.page_token;
1225 if page_token.is_none() {
1226 break;
1227 }
1228 }
1229
1230 Ok(tables)
1231 }
1232
1233 fn validate_root_namespace_id(id: &Option<Vec<String>>) -> Result<()> {
1235 if let Some(id) = id
1236 && !id.is_empty()
1237 {
1238 return Err(NamespaceError::Unsupported {
1239 message: format!(
1240 "Directory namespace only supports root namespace operations, but got namespace ID: {:?}. Expected empty ID.",
1241 id
1242 ),
1243 }
1244 .into());
1245 }
1246 Ok(())
1247 }
1248
1249 fn table_name_from_id(id: &Option<Vec<String>>) -> Result<String> {
1251 let id = id.as_ref().ok_or_else(|| {
1252 lance_core::Error::from(NamespaceError::InvalidInput {
1253 message: "Directory namespace table ID cannot be empty".to_string(),
1254 })
1255 })?;
1256
1257 if id.len() != 1 {
1258 return Err(NamespaceError::Unsupported {
1259 message: format!(
1260 "Multi-level table IDs are only supported when manifest mode is enabled, but got: {:?}",
1261 id
1262 ),
1263 }
1264 .into());
1265 }
1266
1267 Ok(id[0].clone())
1268 }
1269
1270 fn format_table_id(table_id: &[String]) -> String {
1271 format!(
1272 "table id '{}'",
1273 manifest::ManifestNamespace::str_object_id(table_id)
1274 )
1275 }
1276
1277 fn format_table_id_from_request(id: &Option<Vec<String>>) -> String {
1278 id.as_ref()
1279 .map(|table_id| Self::format_table_id(table_id))
1280 .unwrap_or_else(|| "table id '<unknown>'".to_string())
1281 }
1282
1283 async fn resolve_table_location(&self, id: &Option<Vec<String>>) -> Result<String> {
1284 let mut describe_req = DescribeTableRequest::new();
1285 describe_req.id = id.clone();
1286 describe_req.load_detailed_metadata = Some(false);
1287
1288 let describe_resp = self.describe_table_impl(describe_req).await?;
1290
1291 describe_resp.location.ok_or_else(|| {
1292 lance_core::Error::from(NamespaceError::TableNotFound {
1293 message: format!("Table location not found for: {:?}", id),
1294 })
1295 })
1296 }
1297
1298 fn map_tag_error(err: lance_core::Error, tag: &str, table_uri: &str) -> lance_core::Error {
1301 match err {
1302 lance_core::Error::RefNotFound { .. } => NamespaceError::TableTagNotFound {
1303 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1304 }
1305 .into(),
1306 lance_core::Error::RefConflict { .. } => NamespaceError::TableTagAlreadyExists {
1307 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1308 }
1309 .into(),
1310 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1311 message: format!("invalid tag '{}': {}", tag, message),
1312 }
1313 .into(),
1314 lance_core::Error::VersionNotFound { message } => {
1315 NamespaceError::TableVersionNotFound {
1316 message: format!(
1317 "version referenced by tag '{}' not found for table at '{}': {}",
1318 tag, table_uri, message
1319 ),
1320 }
1321 .into()
1322 }
1323 other => NamespaceError::Internal {
1324 message: format!(
1325 "tag operation failed for tag '{}' on table at '{}': {}",
1326 tag, table_uri, other
1327 ),
1328 }
1329 .into(),
1330 }
1331 }
1332
1333 fn map_branch_error(
1339 err: lance_core::Error,
1340 branch: &str,
1341 table_uri: &str,
1342 ) -> lance_core::Error {
1343 match err {
1344 lance_core::Error::RefNotFound { .. } => NamespaceError::TableBranchNotFound {
1345 message: format!("branch '{}' for table at '{}'", branch, table_uri),
1346 }
1347 .into(),
1348 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1349 message: format!("invalid branch '{}': {}", branch, message),
1350 }
1351 .into(),
1352 lance_core::Error::VersionNotFound { message } => {
1353 NamespaceError::TableVersionNotFound {
1354 message: format!(
1355 "source version for branch '{}' not found for table at '{}': {}",
1356 branch, table_uri, message
1357 ),
1358 }
1359 .into()
1360 }
1361 other => NamespaceError::Internal {
1362 message: format!(
1363 "branch operation failed for branch '{}' on table at '{}': {}",
1364 branch, table_uri, other
1365 ),
1366 }
1367 .into(),
1368 }
1369 }
1370
1371 fn map_mutation_error(
1383 err: lance_core::Error,
1384 operation: &str,
1385 table_uri: &str,
1386 ) -> lance_core::Error {
1387 let detail = err.to_string();
1388 let ns_err = match &err {
1389 lance_core::Error::InvalidInput { .. }
1390 | lance_core::Error::Unprocessable { .. }
1391 | lance_core::Error::InvalidRef { .. } => NamespaceError::InvalidInput {
1392 message: format!(
1393 "Invalid input for {} on table at '{}': {}",
1394 operation, table_uri, detail
1395 ),
1396 },
1397 lance_core::Error::NotFound { .. } | lance_core::Error::DatasetNotFound { .. } => {
1398 NamespaceError::TableNotFound {
1399 message: format!(
1400 "Table at '{}' not found while running {}: {}",
1401 table_uri, operation, detail
1402 ),
1403 }
1404 }
1405 lance_core::Error::SchemaMismatch { .. } | lance_core::Error::Schema { .. } => {
1406 NamespaceError::TableSchemaValidationError {
1407 message: format!(
1408 "Schema validation failed for {} on table at '{}': {}",
1409 operation, table_uri, detail
1410 ),
1411 }
1412 }
1413 lance_core::Error::CommitConflict { .. } => NamespaceError::Throttling {
1417 message: format!(
1418 "Too many concurrent writes for {} on table at '{}', please retry later: {}",
1419 operation, table_uri, detail
1420 ),
1421 },
1422 lance_core::Error::TooMuchWriteContention { .. }
1426 | lance_core::Error::RetryableCommitConflict { .. }
1427 | lance_core::Error::IncompatibleTransaction { .. }
1428 | lance_core::Error::VersionConflict { .. } => NamespaceError::ConcurrentModification {
1429 message: format!(
1430 "Concurrent modification detected for {} on table at '{}': {}",
1431 operation, table_uri, detail
1432 ),
1433 },
1434 lance_core::Error::NotSupported { .. } => NamespaceError::Unsupported {
1435 message: format!(
1436 "{} is not supported on table at '{}': {}",
1437 operation, table_uri, detail
1438 ),
1439 },
1440 _ => NamespaceError::Internal {
1441 message: format!(
1442 "Failed to run {} on table at '{}': {}",
1443 operation, table_uri, detail
1444 ),
1445 },
1446 };
1447 ns_err.into()
1448 }
1449
1450 async fn table_has_actual_manifests(&self, table_name: &str) -> Result<bool> {
1451 manifest::ManifestNamespace::path_has_actual_manifests(
1452 &self.object_store,
1453 &self.table_path(table_name),
1454 )
1455 .await
1456 }
1457
1458 async fn filter_declared_tables(
1459 &self,
1460 tables: Vec<String>,
1461 include_declared: bool,
1462 ) -> Result<Vec<String>> {
1463 if include_declared {
1464 return Ok(tables);
1465 }
1466
1467 let mut stream = futures::stream::iter(tables.into_iter().map(|table_name| async move {
1468 if self.table_has_actual_manifests(&table_name).await? {
1472 Ok::<Option<String>, Error>(Some(table_name))
1473 } else {
1474 Ok::<Option<String>, Error>(None)
1475 }
1476 }))
1477 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1478
1479 let mut filtered = Vec::new();
1480 while let Some(result) = stream.next().await {
1481 if let Some(table_name) = result? {
1482 filtered.push(table_name);
1483 }
1484 }
1485 Ok(filtered)
1486 }
1487
1488 fn ipc_reader_from_request_data(
1489 request_data: &Bytes,
1490 operation: &str,
1491 ) -> Result<(
1492 Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1493 usize,
1494 )> {
1495 if request_data.is_empty() {
1496 return Err(NamespaceError::InvalidInput {
1497 message: format!(
1498 "Request data (Arrow IPC stream) is required for {}",
1499 operation
1500 ),
1501 }
1502 .into());
1503 }
1504
1505 let cursor = Cursor::new(request_data.as_ref());
1506 let stream_reader =
1507 StreamReader::try_new(cursor, None).map_err(|e| NamespaceError::InvalidInput {
1508 message: format!("Invalid Arrow IPC stream: {}", e),
1509 })?;
1510 let arrow_schema = stream_reader.schema();
1511
1512 let mut num_rows = 0usize;
1513 let mut batches = Vec::new();
1514 for batch_result in stream_reader {
1515 let batch = batch_result.map_err(|e| NamespaceError::Internal {
1516 message: format!("Failed to read batch from IPC stream: {}", e),
1517 })?;
1518 num_rows += batch.num_rows();
1519 batches.push(batch);
1520 }
1521
1522 let reader: Box<dyn arrow::record_batch::RecordBatchReader + Send> = if batches.is_empty() {
1523 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
1524 Box::new(RecordBatchIterator::new(vec![Ok(batch)], arrow_schema))
1525 } else {
1526 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
1527 Box::new(RecordBatchIterator::new(batch_results, arrow_schema))
1528 };
1529
1530 Ok((reader, num_rows))
1531 }
1532
1533 async fn table_uri_has_actual_manifests(&self, table_uri: &str) -> Result<bool> {
1534 let table_path = self.object_store_path_from_uri(table_uri)?;
1535 manifest::ManifestNamespace::path_has_actual_manifests(&self.object_store, &table_path)
1536 .await
1537 }
1538
1539 fn object_store_path_from_uri(&self, uri: &str) -> Result<Path> {
1540 let registry = self
1541 .session
1542 .as_ref()
1543 .map(|session| session.store_registry())
1544 .unwrap_or_else(|| Arc::new(ObjectStoreRegistry::default()));
1545 ObjectStore::extract_path_from_uri(registry, uri)
1546 }
1547
1548 fn normalized_branch(branch: Option<&str>) -> Result<Option<&str>> {
1553 match branch.filter(|b| !b.is_empty() && *b != "main") {
1554 Some(branch) => {
1555 check_valid_branch(branch).map_err(|e| {
1556 lance_core::Error::from(NamespaceError::InvalidInput {
1557 message: format!("invalid branch name '{}': {}", branch, e),
1558 })
1559 })?;
1560 Ok(Some(branch))
1561 }
1562 None => Ok(None),
1563 }
1564 }
1565
1566 async fn open_validated_branch(&self, table_uri: &str, branch: &str) -> Result<Dataset> {
1567 let dataset = self
1568 .configured_builder(table_uri)
1569 .with_branch(branch, None)
1570 .load()
1571 .await
1572 .map_err(|e| {
1573 let message = format!(
1574 "branch '{}' not found for table at '{}': {}",
1575 branch, table_uri, e
1576 );
1577 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1578 })?;
1579 dataset.branches().get(branch).await.map_err(|e| {
1580 Self::map_open_error(
1581 e,
1582 NamespaceError::TableNotFound {
1583 message: format!("branch '{}' not found for table at '{}'", branch, table_uri),
1584 },
1585 )
1586 })?;
1587 Ok(dataset)
1588 }
1589
1590 async fn resolve_branch_location(&self, table_uri: &str, branch: &str) -> Result<String> {
1591 Ok(self
1592 .open_validated_branch(table_uri, branch)
1593 .await?
1594 .branch_location()
1595 .uri)
1596 }
1597
1598 async fn resolve_branch_for_commit(
1608 &self,
1609 table_uri: &str,
1610 branch: &str,
1611 ) -> Result<(String, Path, Option<u64>)> {
1612 let main = self
1613 .configured_builder(table_uri)
1614 .load()
1615 .await
1616 .map_err(|e| {
1617 let message = format!("table at '{}' not found: {}", table_uri, e);
1618 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1619 })?;
1620 let branch_location = main.branch_location().find_branch(Some(branch))?;
1621 match main.branches().get(branch).await {
1622 Ok(contents) => Ok((
1623 branch_location.uri,
1624 branch_location.path,
1625 Some(contents.parent_version),
1626 )),
1627 Err(lance_core::Error::RefNotFound { .. }) => {
1628 if self
1629 .branch_has_committed_versions(&branch_location.path)
1630 .await?
1631 {
1632 return Err(NamespaceError::TableNotFound {
1633 message: format!(
1634 "branch '{}' not found for table at '{}'",
1635 branch, table_uri
1636 ),
1637 }
1638 .into());
1639 }
1640 Ok((branch_location.uri, branch_location.path, None))
1641 }
1642 Err(e) => Err(e),
1643 }
1644 }
1645
1646 async fn branch_has_committed_versions(&self, branch_path: &Path) -> Result<bool> {
1647 Ok(!self
1648 .list_versions_under(branch_path, false, Some(1))
1649 .await?
1650 .is_empty())
1651 }
1652
1653 fn validate_dir_only_properties(
1654 properties: Option<&HashMap<String, String>>,
1655 operation: &str,
1656 ) -> Result<()> {
1657 if properties.is_some_and(|properties| !properties.is_empty()) {
1661 return Err(NamespaceError::Unsupported {
1662 message: format!(
1663 "{} with non-empty table properties requires manifest_enabled=true",
1664 operation
1665 ),
1666 }
1667 .into());
1668 }
1669 Ok(())
1670 }
1671
1672 async fn write_reader_to_table(
1673 &self,
1674 table_uri: &str,
1675 reader: Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1676 mode: WriteMode,
1677 extra_storage_options: Option<HashMap<String, String>>,
1678 ) -> Result<Dataset> {
1679 let mut merged_storage_options = self.storage_options.clone().unwrap_or_default();
1682 if let Some(extra_storage_options) = extra_storage_options {
1683 merged_storage_options.extend(extra_storage_options);
1684 }
1685 let store_params = (!merged_storage_options.is_empty()).then(|| ObjectStoreParams {
1686 storage_options_accessor: Some(Arc::new(
1687 lance_io::object_store::StorageOptionsAccessor::with_static_options(
1688 merged_storage_options,
1689 ),
1690 )),
1691 ..Default::default()
1692 });
1693
1694 let write_params = WriteParams {
1695 mode,
1696 store_params,
1697 session: self.session.clone(),
1698 ..Default::default()
1699 };
1700
1701 let dataset = Dataset::write(reader, table_uri, Some(write_params))
1702 .await
1703 .map_err(|e| NamespaceError::Internal {
1704 message: format!("Failed to write table at '{}': {}", table_uri, e),
1705 })?;
1706
1707 Ok(dataset)
1708 }
1709
1710 fn manifest_version_from_filename(filename: &str) -> Option<u64> {
1714 ManifestNamingScheme::detect_scheme(filename)?.parse_version(filename)
1715 }
1716
1717 fn create_table_version_response(
1719 version: u64,
1720 final_path: &Path,
1721 final_meta: &ObjectMeta,
1722 ) -> CreateTableVersionResponse {
1723 CreateTableVersionResponse {
1724 transaction_id: None,
1725 version: Some(Box::new(TableVersion {
1726 version: version as i64,
1727 manifest_path: final_path.to_string(),
1728 manifest_size: Some(final_meta.size as i64),
1729 e_tag: final_meta.e_tag.clone(),
1730 timestamp_millis: None,
1731 metadata: None,
1732 })),
1733 ..Default::default()
1734 }
1735 }
1736
1737 async fn staging_matches_final_manifest(
1745 &self,
1746 staging_path: &Path,
1747 final_path: &Path,
1748 final_meta: &ObjectMeta,
1749 request_manifest_size: Option<i64>,
1750 ) -> Result<bool> {
1751 if let Some(size) = request_manifest_size
1752 && size != final_meta.size as i64
1753 {
1754 return Ok(false);
1755 }
1756
1757 let staging_bytes = match self.object_store.inner.get(staging_path).await {
1758 Ok(r) => r.bytes().await.map_err(|e| {
1759 lance_core::Error::from(NamespaceError::Internal {
1760 message: format!(
1761 "Failed to read staging manifest at '{}': {}",
1762 staging_path, e
1763 ),
1764 })
1765 })?,
1766 Err(ObjectStoreError::NotFound { .. }) => return Ok(false),
1767 Err(e) => {
1768 return Err(lance_core::Error::from(NamespaceError::Internal {
1769 message: format!(
1770 "Failed to read staging manifest at '{}': {}",
1771 staging_path, e
1772 ),
1773 }));
1774 }
1775 };
1776
1777 let final_bytes = self
1778 .object_store
1779 .inner
1780 .get(final_path)
1781 .await
1782 .map_err(|e| {
1783 lance_core::Error::from(NamespaceError::Internal {
1784 message: format!(
1785 "Failed to read existing version manifest at '{}': {}",
1786 final_path, e
1787 ),
1788 })
1789 })?
1790 .bytes()
1791 .await
1792 .map_err(|e| {
1793 lance_core::Error::from(NamespaceError::Internal {
1794 message: format!(
1795 "Failed to read existing version manifest bytes at '{}': {}",
1796 final_path, e
1797 ),
1798 })
1799 })?;
1800
1801 Ok(staging_bytes.as_ref() == final_bytes.as_ref())
1802 }
1803
1804 async fn resolve_existing_table_version(
1806 &self,
1807 args: ExistingTableVersionResolve<'_>,
1808 ) -> Result<CreateTableVersionResponse> {
1809 if self
1810 .staging_matches_final_manifest(
1811 args.staging_path,
1812 args.final_path,
1813 args.final_meta,
1814 args.request_manifest_size,
1815 )
1816 .await?
1817 {
1818 if let Err(e) = self.object_store.inner.delete(args.staging_path).await {
1820 log::warn!(
1821 "Failed to delete staging manifest at '{}': {:?}",
1822 args.staging_path,
1823 e
1824 );
1825 }
1826 return Ok(Self::create_table_version_response(
1827 args.version,
1828 args.final_path,
1829 args.final_meta,
1830 ));
1831 }
1832
1833 Err(lance_core::Error::from(
1834 NamespaceError::ConcurrentModification {
1835 message: format!(
1836 "Version {} already exists for table at '{}' with different content",
1837 args.version, args.table_uri
1838 ),
1839 },
1840 ))
1841 }
1842
1843 async fn enforce_create_table_version_cas(
1852 &self,
1853 table_path: &Path,
1854 version: u64,
1855 table_uri: &str,
1856 is_branch: bool,
1857 branch_parent_version: Option<u64>,
1858 ) -> Result<()> {
1859 let latest = self.list_versions_under(table_path, true, Some(1)).await?;
1860 let expected = match latest.first() {
1861 Some(v) => (v.version as u64).checked_add(1).ok_or_else(|| {
1862 lance_core::Error::from(NamespaceError::ConcurrentModification {
1863 message: format!(
1864 "Version overflow computing next version for table at '{}': \
1865 latest version {} cannot advance",
1866 table_uri, v.version
1867 ),
1868 })
1869 })?,
1870 None => {
1871 if is_branch {
1872 match branch_parent_version {
1875 Some(parent_version) => parent_version,
1876 None => version,
1877 }
1878 } else {
1879 1
1880 }
1881 }
1882 };
1883 if version != expected {
1884 let latest_display = latest
1885 .first()
1886 .map(|v| v.version.to_string())
1887 .unwrap_or_else(|| "none".to_string());
1888 return Err(lance_core::Error::from(
1889 NamespaceError::ConcurrentModification {
1890 message: format!(
1891 "Version CAS failed for table at '{}': requested {}, expected {} (latest {})",
1892 table_uri, version, expected, latest_display
1893 ),
1894 },
1895 ));
1896 }
1897 Ok(())
1898 }
1899
1900 async fn materialize_version_manifest_create(
1902 &self,
1903 staging_path: &Path,
1904 final_path: &Path,
1905 staging_manifest_path: &str,
1906 ) -> std::result::Result<(), ObjectStoreError> {
1907 match self
1908 .object_store
1909 .inner
1910 .copy_if_not_exists(staging_path, final_path)
1911 .await
1912 {
1913 Ok(()) => Ok(()),
1914 Err(ObjectStoreError::NotImplemented { .. })
1915 | Err(ObjectStoreError::NotSupported { .. }) => {
1916 let manifest_data = self
1917 .object_store
1918 .inner
1919 .get(staging_path)
1920 .await?
1921 .bytes()
1922 .await
1923 .map_err(|e| ObjectStoreError::Generic {
1924 store: "DirectoryNamespace",
1925 source: Box::new(std::io::Error::other(format!(
1926 "Failed to read staging manifest bytes at '{}': {}",
1927 staging_manifest_path, e
1928 ))),
1929 })?;
1930 self.object_store
1931 .inner
1932 .put_opts(
1933 final_path,
1934 manifest_data.into(),
1935 PutOptions {
1936 mode: PutMode::Create,
1937 ..Default::default()
1938 },
1939 )
1940 .await
1941 .map(|_| ())
1942 }
1943 Err(e) => Err(e),
1944 }
1945 }
1946
1947 async fn list_table_versions_from_storage(
1948 &self,
1949 table_uri: &str,
1950 descending: bool,
1951 limit: Option<i32>,
1952 ) -> Result<Vec<TableVersion>> {
1953 let table_path = self.object_store_path_from_uri(table_uri)?;
1954 self.list_versions_under(&table_path, descending, limit)
1955 .await
1956 }
1957
1958 async fn list_versions_under(
1962 &self,
1963 table_path: &Path,
1964 descending: bool,
1965 limit: Option<i32>,
1966 ) -> Result<Vec<TableVersion>> {
1967 let versions_dir = table_path.clone().join(VERSIONS_DIR);
1968 let mut stream = self.object_store.read_dir_all(&versions_dir, None);
1969 let list_err = |e: lance_core::Error| {
1970 lance_core::Error::from(NamespaceError::Internal {
1971 message: format!(
1972 "Failed to list manifest files under '{}': {}",
1973 versions_dir, e
1974 ),
1975 })
1976 };
1977
1978 let limit = limit
1979 .filter(|limit| *limit >= 0)
1980 .map(|limit| limit as usize);
1981
1982 let mut table_versions: Vec<TableVersion> = Vec::new();
1983 let push_meta = |meta: ObjectMeta, out: &mut Vec<TableVersion>| -> bool {
1984 let Some(filename) = meta.location.filename() else {
1985 return false;
1986 };
1987 let Some(actual_version) = Self::manifest_version_from_filename(filename) else {
1988 return false;
1989 };
1990 out.push(TableVersion {
1991 version: actual_version as i64,
1992 manifest_path: meta.location.to_string(),
1993 manifest_size: Some(meta.size as i64),
1994 e_tag: meta.e_tag,
1995 timestamp_millis: Some(meta.last_modified.timestamp_millis()),
1996 metadata: None,
1997 });
1998 true
1999 };
2000
2001 let mut first_manifest_filename_len = None;
2006 while first_manifest_filename_len.is_none() {
2007 match stream.try_next().await.map_err(list_err)? {
2008 Some(meta) => {
2009 let filename_len = meta.location.filename().map(|f| f.len());
2010 if push_meta(meta, &mut table_versions) {
2011 first_manifest_filename_len = filename_len;
2012 }
2013 }
2014 None => break,
2015 }
2016 }
2017 let is_v2_naming = first_manifest_filename_len == Some(29);
2018
2019 let list_is_ordered = self.object_store.list_is_lexically_ordered;
2024 let stream_matches_request = list_is_ordered
2025 && if is_v2_naming {
2026 descending
2027 } else {
2028 !descending
2029 };
2030 let early_stop_at = limit.filter(|_| stream_matches_request);
2031
2032 while early_stop_at.is_none_or(|n| table_versions.len() < n) {
2033 match stream.try_next().await.map_err(list_err)? {
2034 Some(meta) => {
2035 push_meta(meta, &mut table_versions);
2036 }
2037 None => break,
2038 }
2039 }
2040
2041 if let Some(n) = early_stop_at {
2044 table_versions.truncate(n);
2045 } else {
2046 if descending {
2047 table_versions.sort_by_key(|v| std::cmp::Reverse(v.version));
2048 } else {
2049 table_versions.sort_by_key(|v| v.version);
2050 }
2051 if let Some(limit) = limit {
2052 table_versions.truncate(limit);
2053 }
2054 }
2055
2056 Ok(table_versions)
2057 }
2058
2059 async fn describe_table_impl(
2063 &self,
2064 request: DescribeTableRequest,
2065 ) -> Result<DescribeTableResponse> {
2066 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
2067 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
2068 let skip_manifest_for_root = self.dir_listing_enabled
2069 && is_root_level
2070 && !self.dir_listing_to_manifest_migration_enabled;
2071 if is_child_table
2079 || !self.dir_listing_enabled
2080 || self.dir_listing_to_manifest_migration_enabled
2081 {
2082 self.ensure_read_manifest().await?;
2083 }
2084 if let Some(manifest_ns) = self.manifest_ns_for_read()
2085 && !skip_manifest_for_root
2086 {
2087 match manifest_ns.describe_table(request.clone()).await {
2088 Ok(mut response) => {
2089 if let Some(ref table_uri) = response.table_uri {
2090 let vend = request.vend_credentials.unwrap_or(true);
2092 let identity = request.identity.as_deref();
2093 response.storage_options = self
2094 .get_storage_options_for_table(table_uri, vend, identity)
2095 .await?;
2096 }
2097 if self.table_version_tracking_enabled {
2099 response.managed_versioning = Some(true);
2100 }
2101 return Ok(response);
2102 }
2103 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
2104 return Err(e);
2107 }
2108 Err(e) if self.dir_listing_enabled && is_root_level => {
2109 if !Self::is_manifest_table_absent_error(&e) {
2113 return Err(Self::classify_storage_error(e));
2114 }
2115 }
2116 Err(e) => return Err(e),
2117 }
2118 }
2119 if is_child_table {
2120 return Err(self.child_namespace_requires_manifest_error());
2121 }
2122
2123 let table_name = Self::table_name_from_id(&request.id)?;
2124 let table_id = Self::format_table_id_from_request(&request.id);
2125 if !self.dir_listing_enabled {
2126 return Err(NamespaceError::TableNotFound { message: table_id }.into());
2127 }
2128
2129 let table_uri = self.table_full_uri(&table_name);
2130
2131 let status = self.check_table_status(&table_name).await?;
2133
2134 if !status.exists {
2135 return Err(NamespaceError::TableNotFound {
2136 message: table_id.clone(),
2137 }
2138 .into());
2139 }
2140
2141 if status.is_deregistered {
2142 return Err(NamespaceError::TableNotFound {
2143 message: format!("Table is deregistered: {}", table_id),
2144 }
2145 .into());
2146 }
2147
2148 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
2149 let should_check_declared =
2150 load_detailed_metadata || request.check_declared.unwrap_or(false);
2151 let vend_credentials = request.vend_credentials.unwrap_or(true);
2153 let identity = request.identity.as_deref();
2154 let is_only_declared = if should_check_declared {
2155 if status.has_reserved_file {
2156 Some(!self.table_has_actual_manifests(&table_name).await?)
2157 } else {
2158 Some(false)
2159 }
2160 } else {
2161 None
2162 };
2163
2164 if !load_detailed_metadata {
2165 let storage_options = self
2166 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2167 .await?;
2168 return Ok(DescribeTableResponse {
2169 table: Some(table_name),
2170 namespace: request.id.as_ref().map(|id| {
2171 if id.len() > 1 {
2172 id[..id.len() - 1].to_vec()
2173 } else {
2174 vec![]
2175 }
2176 }),
2177 location: Some(table_uri.clone()),
2178 table_uri: Some(table_uri),
2179 storage_options,
2180 is_only_declared,
2181 managed_versioning: if self.table_version_tracking_enabled {
2182 Some(true)
2183 } else {
2184 None
2185 },
2186 ..Default::default()
2187 });
2188 }
2189
2190 if is_only_declared == Some(true) {
2191 let storage_options = self
2192 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2193 .await?;
2194 return Ok(DescribeTableResponse {
2195 table: Some(table_name),
2196 namespace: request.id.as_ref().map(|id| {
2197 if id.len() > 1 {
2198 id[..id.len() - 1].to_vec()
2199 } else {
2200 vec![]
2201 }
2202 }),
2203 location: Some(table_uri.clone()),
2204 table_uri: Some(table_uri),
2205 storage_options,
2206 is_only_declared,
2207 managed_versioning: if self.table_version_tracking_enabled {
2208 Some(true)
2209 } else {
2210 None
2211 },
2212 ..Default::default()
2213 });
2214 }
2215
2216 let mut builder = DatasetBuilder::from_uri(&table_uri);
2219 if let Some(opts) = &self.storage_options {
2220 builder = builder.with_storage_options(opts.clone());
2221 }
2222 if let Some(sess) = &self.session {
2223 builder = builder.with_session(sess.clone());
2224 }
2225 match builder.load().await {
2226 Ok(mut dataset) => {
2227 if let Some(requested_version) = request.version {
2229 dataset = dataset
2230 .checkout_version(requested_version as u64)
2231 .await
2232 .map_err(|e| {
2233 let message = format!(
2234 "Version {} not found for table '{}': {}",
2235 requested_version, table_name, e
2236 );
2237 Self::map_open_error(
2238 e,
2239 NamespaceError::TableVersionNotFound { message },
2240 )
2241 })?;
2242 }
2243
2244 let version_info = dataset.version();
2245 let lance_schema = dataset.schema();
2246 let arrow_schema: arrow_schema::Schema = lance_schema.into();
2247 let json_schema = arrow_schema_to_json(&arrow_schema)?;
2248 let storage_options = self
2249 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2250 .await?;
2251
2252 let metadata: std::collections::HashMap<String, String> =
2254 version_info.metadata.into_iter().collect();
2255
2256 Ok(DescribeTableResponse {
2257 table: Some(table_name),
2258 namespace: request.id.as_ref().map(|id| {
2259 if id.len() > 1 {
2260 id[..id.len() - 1].to_vec()
2261 } else {
2262 vec![]
2263 }
2264 }),
2265 version: Some(version_info.version as i64),
2266 location: Some(table_uri.clone()),
2267 table_uri: Some(table_uri),
2268 schema: Some(Box::new(json_schema)),
2269 storage_options,
2270 metadata: Some(metadata),
2271 is_only_declared,
2272 managed_versioning: if self.table_version_tracking_enabled {
2273 Some(true)
2274 } else {
2275 None
2276 },
2277 ..Default::default()
2278 })
2279 }
2280 Err(err) => {
2281 if manifest::ManifestNamespace::is_not_found_load_error(&err)
2282 && is_only_declared == Some(true)
2283 {
2284 let storage_options = self
2285 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2286 .await?;
2287 Ok(DescribeTableResponse {
2288 table: Some(table_name),
2289 namespace: request.id.as_ref().map(|id| {
2290 if id.len() > 1 {
2291 id[..id.len() - 1].to_vec()
2292 } else {
2293 vec![]
2294 }
2295 }),
2296 location: Some(table_uri.clone()),
2297 table_uri: Some(table_uri),
2298 storage_options,
2299 is_only_declared,
2300 managed_versioning: if self.table_version_tracking_enabled {
2301 Some(true)
2302 } else {
2303 None
2304 },
2305 ..Default::default()
2306 })
2307 } else {
2308 Err(NamespaceError::Internal {
2309 message: format!(
2310 "Table directory exists but cannot load dataset {}: {:?}",
2311 table_name, err
2312 ),
2313 }
2314 .into())
2315 }
2316 }
2317 }
2318 }
2319
2320 fn configured_builder(&self, table_uri: &str) -> DatasetBuilder {
2323 let mut builder = DatasetBuilder::from_uri(table_uri);
2324 if let Some(opts) = &self.storage_options {
2325 builder = builder.with_storage_options(opts.clone());
2326 }
2327 if let Some(sess) = &self.session {
2328 builder = builder.with_session(sess.clone());
2329 }
2330 builder
2331 }
2332
2333 async fn load_dataset(
2334 &self,
2335 table_uri: &str,
2336 version: Option<i64>,
2337 operation: &str,
2338 ) -> Result<Dataset> {
2339 if let Some(version) = version
2340 && version < 0
2341 {
2342 return Err(NamespaceError::InvalidInput {
2343 message: format!(
2344 "Table version for {} must be non-negative, got {}",
2345 operation, version
2346 ),
2347 }
2348 .into());
2349 }
2350
2351 let builder = self.configured_builder(table_uri);
2352
2353 let dataset = builder.load().await.map_err(|e| {
2354 let message = format!(
2355 "Failed to open table at '{}' for {}: {}",
2356 table_uri, operation, e
2357 );
2358 Self::map_open_error(e, NamespaceError::TableNotFound { message })
2359 })?;
2360
2361 if let Some(version) = version {
2362 return dataset.checkout_version(version as u64).await.map_err(|e| {
2363 let message = format!(
2364 "Failed to checkout version {} for table at '{}' during {}: {}",
2365 version, table_uri, operation, e
2366 );
2367 Self::map_open_error(e, NamespaceError::TableVersionNotFound { message })
2368 });
2369 }
2370
2371 Ok(dataset)
2372 }
2373
2374 fn parse_index_type(index_type: &str) -> Result<IndexType> {
2375 match index_type.trim().to_ascii_uppercase().as_str() {
2376 "SCALAR" | "BTREE" => Ok(IndexType::BTree),
2377 "BITMAP" => Ok(IndexType::Bitmap),
2378 "LABEL_LIST" | "LABELLIST" => Ok(IndexType::LabelList),
2379 "INVERTED" | "FTS" => Ok(IndexType::Inverted),
2380 "NGRAM" => Ok(IndexType::NGram),
2381 "ZONEMAP" | "ZONE_MAP" => Ok(IndexType::ZoneMap),
2382 "BLOOMFILTER" | "BLOOM_FILTER" => Ok(IndexType::BloomFilter),
2383 "RTREE" | "R_TREE" => Ok(IndexType::RTree),
2384 "VECTOR" | "IVF_PQ" => Ok(IndexType::IvfPq),
2385 "IVF_FLAT" => Ok(IndexType::IvfFlat),
2386 "IVF_SQ" => Ok(IndexType::IvfSq),
2387 "IVF_RQ" => Ok(IndexType::IvfRq),
2388 "IVF_HNSW_FLAT" => Ok(IndexType::IvfHnswFlat),
2389 "IVF_HNSW_SQ" => Ok(IndexType::IvfHnswSq),
2390 "IVF_HNSW_PQ" => Ok(IndexType::IvfHnswPq),
2391 other => Err(NamespaceError::InvalidInput {
2392 message: format!("Unsupported index_type '{}'", other),
2393 }
2394 .into()),
2395 }
2396 }
2397
2398 fn parse_metric_type(distance_type: Option<&str>) -> Result<MetricType> {
2399 let distance_type = distance_type.unwrap_or("l2");
2400 MetricType::try_from(distance_type).map_err(|e| {
2401 lance_core::Error::from(NamespaceError::InvalidInput {
2402 message: format!(
2403 "Unsupported distance_type '{}' for vector index: {}",
2404 distance_type, e
2405 ),
2406 })
2407 })
2408 }
2409
2410 fn build_index_params(request: &CreateTableIndexRequest) -> Result<DirectoryIndexParams> {
2411 let index_type = Self::parse_index_type(&request.index_type)?;
2412 Ok(match index_type {
2413 IndexType::BTree => DirectoryIndexParams::Scalar {
2414 index_type,
2415 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
2416 },
2417 IndexType::Bitmap => DirectoryIndexParams::Scalar {
2418 index_type,
2419 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
2420 },
2421 IndexType::LabelList => DirectoryIndexParams::Scalar {
2422 index_type,
2423 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
2424 },
2425 IndexType::NGram => DirectoryIndexParams::Scalar {
2426 index_type,
2427 params: ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
2428 },
2429 IndexType::ZoneMap => DirectoryIndexParams::Scalar {
2430 index_type,
2431 params: ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
2432 },
2433 IndexType::BloomFilter => DirectoryIndexParams::Scalar {
2434 index_type,
2435 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter),
2436 },
2437 IndexType::RTree => DirectoryIndexParams::Scalar {
2438 index_type,
2439 params: ScalarIndexParams::for_builtin(BuiltinIndexType::RTree),
2440 },
2441 IndexType::Inverted => {
2442 let mut params = InvertedIndexParams::default();
2443 if let Some(with_position) = request.with_position {
2444 params = params.with_position(with_position);
2445 }
2446 if let Some(base_tokenizer) = &request.base_tokenizer {
2447 params = params.base_tokenizer(base_tokenizer.clone());
2448 }
2449 if let Some(language) = &request.language {
2450 params = params.language(language)?;
2451 }
2452 if let Some(max_token_length) = request.max_token_length {
2453 if max_token_length < 0 {
2454 return Err(NamespaceError::InvalidInput {
2455 message: format!(
2456 "FTS max_token_length must be non-negative, got {}",
2457 max_token_length
2458 ),
2459 }
2460 .into());
2461 }
2462 params = params.max_token_length(Some(max_token_length as usize));
2463 }
2464 if let Some(lower_case) = request.lower_case {
2465 params = params.lower_case(lower_case);
2466 }
2467 if let Some(stem) = request.stem {
2468 params = params.stem(stem);
2469 }
2470 if let Some(remove_stop_words) = request.remove_stop_words {
2471 params = params.remove_stop_words(remove_stop_words);
2472 }
2473 if let Some(ascii_folding) = request.ascii_folding {
2474 params = params.ascii_folding(ascii_folding);
2475 }
2476 DirectoryIndexParams::Inverted(params)
2477 }
2478 IndexType::IvfFlat => DirectoryIndexParams::Vector {
2479 index_type,
2480 params: VectorIndexParams::with_ivf_flat_params(
2481 Self::parse_metric_type(request.distance_type.as_deref())?,
2482 IvfBuildParams::default(),
2483 ),
2484 },
2485 IndexType::IvfPq => DirectoryIndexParams::Vector {
2486 index_type,
2487 params: VectorIndexParams::with_ivf_pq_params(
2488 Self::parse_metric_type(request.distance_type.as_deref())?,
2489 IvfBuildParams::default(),
2490 PQBuildParams::default(),
2491 ),
2492 },
2493 IndexType::IvfSq => DirectoryIndexParams::Vector {
2494 index_type,
2495 params: VectorIndexParams::with_ivf_sq_params(
2496 Self::parse_metric_type(request.distance_type.as_deref())?,
2497 IvfBuildParams::default(),
2498 SQBuildParams::default(),
2499 ),
2500 },
2501 IndexType::IvfRq => {
2502 let rq_params = if let Some(requested_num_bits) = request.num_bits {
2503 let invalid_num_bits = || NamespaceError::InvalidInput {
2504 message: format!(
2505 "IVF_RQ num_bits must be in {}..={}, got {}",
2506 RABIT_MIN_NUM_BITS, RABIT_MAX_NUM_BITS, requested_num_bits
2507 ),
2508 };
2509 let num_bits =
2510 u8::try_from(requested_num_bits).map_err(|_| invalid_num_bits())?;
2511 validate_supported_rq_num_bits(num_bits).map_err(|_| invalid_num_bits())?;
2512 RQBuildParams::new(num_bits)
2513 } else {
2514 RQBuildParams::default()
2515 };
2516 DirectoryIndexParams::Vector {
2517 index_type,
2518 params: VectorIndexParams::with_ivf_rq_params(
2519 Self::parse_metric_type(request.distance_type.as_deref())?,
2520 IvfBuildParams::default(),
2521 rq_params,
2522 ),
2523 }
2524 }
2525 IndexType::IvfHnswFlat => DirectoryIndexParams::Vector {
2526 index_type,
2527 params: VectorIndexParams::ivf_hnsw(
2528 Self::parse_metric_type(request.distance_type.as_deref())?,
2529 IvfBuildParams::default(),
2530 HnswBuildParams::default(),
2531 ),
2532 },
2533 IndexType::IvfHnswSq => DirectoryIndexParams::Vector {
2534 index_type,
2535 params: VectorIndexParams::with_ivf_hnsw_sq_params(
2536 Self::parse_metric_type(request.distance_type.as_deref())?,
2537 IvfBuildParams::default(),
2538 HnswBuildParams::default(),
2539 SQBuildParams::default(),
2540 ),
2541 },
2542 IndexType::IvfHnswPq => DirectoryIndexParams::Vector {
2543 index_type,
2544 params: VectorIndexParams::with_ivf_hnsw_pq_params(
2545 Self::parse_metric_type(request.distance_type.as_deref())?,
2546 IvfBuildParams::default(),
2547 HnswBuildParams::default(),
2548 PQBuildParams::default(),
2549 ),
2550 },
2551 other => {
2552 return Err(NamespaceError::InvalidInput {
2553 message: format!("Unsupported index type for namespace API: {}", other),
2554 }
2555 .into());
2556 }
2557 })
2558 }
2559
2560 fn paginate_indices(
2561 indices: &mut Vec<IndexContent>,
2562 page_token: Option<String>,
2563 limit: Option<i32>,
2564 ) -> Option<String> {
2565 indices.sort_by(|a, b| a.index_name.cmp(&b.index_name));
2566
2567 if let Some(start_after) = page_token {
2568 if let Some(index) = indices
2569 .iter()
2570 .position(|index| index.index_name.as_str() > start_after.as_str())
2571 {
2572 indices.drain(0..index);
2573 } else {
2574 indices.clear();
2575 }
2576 }
2577
2578 let mut next_page_token = None;
2579 if let Some(limit) = limit
2580 && limit >= 0
2581 {
2582 let limit = limit as usize;
2583 if limit > 0 && indices.len() > limit {
2584 next_page_token = Some(indices[limit - 1].index_name.clone());
2585 }
2586 indices.truncate(limit);
2587 }
2588 if indices.is_empty() {
2589 None
2590 } else {
2591 next_page_token
2592 }
2593 }
2594
2595 fn transaction_operation_name(transaction: &Transaction) -> String {
2596 match &transaction.operation {
2597 Operation::CreateIndex {
2598 new_indices,
2599 removed_indices,
2600 ..
2601 } if new_indices.is_empty() && !removed_indices.is_empty() => "DropIndex".to_string(),
2602 _ => transaction.operation.to_string(),
2603 }
2604 }
2605
2606 fn transaction_response(
2607 version: u64,
2608 transaction: &Transaction,
2609 alteration: Option<TransactionAlteration>,
2610 ) -> DescribeTransactionResponse {
2611 let mut properties = transaction
2612 .transaction_properties
2613 .as_ref()
2614 .map(|properties| (**properties).clone())
2615 .unwrap_or_default();
2616
2617 let mut effective_status = "SUCCEEDED".to_string();
2620 if let Some(alteration) = alteration {
2621 for key in &alteration.removed_properties {
2622 properties.remove(key);
2623 }
2624 for (key, value) in alteration.properties {
2625 properties.insert(key, value);
2626 }
2627 if let Some(status) = alteration.status {
2628 effective_status = status;
2629 }
2630 }
2631
2632 properties.insert("uuid".to_string(), transaction.uuid.clone());
2633 properties.insert("version".to_string(), version.to_string());
2634 properties.insert(
2635 "read_version".to_string(),
2636 transaction.read_version.to_string(),
2637 );
2638 properties.insert(
2639 "operation".to_string(),
2640 Self::transaction_operation_name(transaction),
2641 );
2642 if let Some(tag) = &transaction.tag {
2643 properties.insert("tag".to_string(), tag.clone());
2644 }
2645
2646 DescribeTransactionResponse {
2647 status: effective_status,
2648 properties: Some(properties),
2649 ..Default::default()
2650 }
2651 }
2652
2653 fn describe_table_index_stats_response(
2654 stats: &serde_json::Value,
2655 ) -> DescribeTableIndexStatsResponse {
2656 let get_i64 = |key: &str| {
2657 stats.get(key).and_then(|value| {
2658 value
2659 .as_i64()
2660 .or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
2661 })
2662 };
2663
2664 DescribeTableIndexStatsResponse {
2665 distance_type: stats
2666 .get("distance_type")
2667 .and_then(|value| value.as_str())
2668 .map(str::to_string),
2669 index_type: stats
2670 .get("index_type")
2671 .and_then(|value| value.as_str())
2672 .map(str::to_string),
2673 num_indexed_rows: get_i64("num_indexed_rows"),
2674 num_unindexed_rows: get_i64("num_unindexed_rows"),
2675 num_indices: get_i64("num_indices").and_then(|value| i32::try_from(value).ok()),
2676 ..Default::default()
2677 }
2678 }
2679
2680 async fn find_transaction(&self, dataset: &Dataset, id: &str) -> Result<(u64, Transaction)> {
2685 if let Ok(version) = id.parse::<u64>() {
2686 let transaction = dataset
2687 .read_transaction_by_version(version)
2688 .await
2689 .map_err(|e| {
2690 lance_core::Error::from(NamespaceError::TransactionNotFound {
2691 message: format!(
2692 "Failed to read transaction for version {}: {}",
2693 version, e
2694 ),
2695 })
2696 })?
2697 .ok_or_else(|| {
2698 lance_core::Error::from(NamespaceError::TransactionNotFound {
2699 message: format!("version {}", version),
2700 })
2701 })?;
2702 return Ok((version, transaction));
2703 }
2704
2705 let versions = dataset.versions().await.map_err(|e| {
2706 lance_core::Error::from(NamespaceError::Internal {
2707 message: format!(
2708 "Failed to list table versions while resolving transaction '{}': {}",
2709 id, e
2710 ),
2711 })
2712 })?;
2713
2714 for version in versions.into_iter().rev() {
2715 if let Some(transaction) = dataset
2716 .read_transaction_by_version(version.version)
2717 .await
2718 .map_err(|e| {
2719 lance_core::Error::from(NamespaceError::Internal {
2720 message: format!(
2721 "Failed to read transaction for version {} while resolving '{}': {}",
2722 version.version, id, e
2723 ),
2724 })
2725 })?
2726 && transaction.uuid == id
2727 {
2728 return Ok((version.version, transaction));
2729 }
2730 }
2731
2732 Err(NamespaceError::TransactionNotFound {
2733 message: id.to_string(),
2734 }
2735 .into())
2736 }
2737
2738 const TRANSACTION_ALTERATIONS_DIR: &'static str = "_alter_transactions";
2742
2743 fn transaction_alteration_path(&self, table_uri: &str, txn_uuid: &str) -> Result<Path> {
2744 let table_path = self.object_store_path_from_uri(table_uri)?;
2745 Ok(table_path
2746 .join(Self::TRANSACTION_ALTERATIONS_DIR)
2747 .join(format!("{}.json", txn_uuid).as_str()))
2748 }
2749
2750 async fn load_transaction_alteration(
2751 &self,
2752 table_uri: &str,
2753 txn_uuid: &str,
2754 ) -> Result<Option<TransactionAlteration>> {
2755 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2756 match self.object_store.inner.get(&path).await {
2757 Ok(get_result) => {
2758 let bytes = get_result.bytes().await.map_err(|e| {
2759 lance_core::Error::from(NamespaceError::Internal {
2760 message: format!(
2761 "Failed to read alter_transaction sidecar for '{}': {}",
2762 txn_uuid, e
2763 ),
2764 })
2765 })?;
2766 let alteration = TransactionAlteration::from_json_slice(&bytes).map_err(|e| {
2767 lance_core::Error::from(NamespaceError::Internal {
2768 message: format!(
2769 "Failed to parse alter_transaction sidecar for '{}': {}",
2770 txn_uuid, e
2771 ),
2772 })
2773 })?;
2774 Ok(Some(alteration))
2775 }
2776 Err(ObjectStoreError::NotFound { .. }) => Ok(None),
2777 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2778 message: format!(
2779 "Failed to load alter_transaction sidecar for '{}': {}",
2780 txn_uuid, e
2781 ),
2782 })),
2783 }
2784 }
2785
2786 async fn save_transaction_alteration(
2787 &self,
2788 table_uri: &str,
2789 txn_uuid: &str,
2790 alteration: &TransactionAlteration,
2791 ) -> Result<()> {
2792 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2793 let bytes = alteration.to_json_bytes().map_err(|e| {
2794 lance_core::Error::from(NamespaceError::Internal {
2795 message: format!(
2796 "Failed to serialize alter_transaction sidecar for '{}': {}",
2797 txn_uuid, e
2798 ),
2799 })
2800 })?;
2801 self.object_store
2802 .inner
2803 .put(&path, bytes.into())
2804 .await
2805 .map_err(|e| {
2806 lance_core::Error::from(NamespaceError::Internal {
2807 message: format!(
2808 "Failed to persist alter_transaction sidecar for '{}': {}",
2809 txn_uuid, e
2810 ),
2811 })
2812 })?;
2813 Ok(())
2814 }
2815
2816 fn table_full_uri(&self, table_name: &str) -> String {
2817 format!("{}/{}.lance", self.root, table_name)
2818 }
2819
2820 fn table_path(&self, table_name: &str) -> Path {
2822 self.base_path
2823 .clone()
2824 .join(format!("{}.lance", table_name).as_str())
2825 }
2826
2827 fn table_reserved_file_path(&self, table_name: &str) -> Path {
2829 self.base_path
2830 .clone()
2831 .join(format!("{}.lance", table_name).as_str())
2832 .join(".lance-reserved")
2833 }
2834
2835 fn table_deregistered_file_path(&self, table_name: &str) -> Path {
2837 self.base_path
2838 .clone()
2839 .join(format!("{}.lance", table_name).as_str())
2840 .join(".lance-deregistered")
2841 }
2842
2843 pub(crate) async fn check_table_status(&self, table_name: &str) -> Result<TableStatus> {
2849 let table_path = self.table_path(table_name);
2850 match self.object_store.read_dir(table_path).await {
2851 Ok(entries) => {
2852 let exists = !entries.is_empty();
2853 let is_deregistered = entries.iter().any(|e| e.ends_with(".lance-deregistered"));
2854 let has_reserved_file = entries.iter().any(|e| e.ends_with(".lance-reserved"));
2855 Ok(TableStatus {
2856 exists,
2857 is_deregistered,
2858 has_reserved_file,
2859 })
2860 }
2861 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(TableStatus {
2864 exists: false,
2865 is_deregistered: false,
2866 has_reserved_file: false,
2867 }),
2868 Err(e) => Err(Self::classify_storage_error(e)),
2872 }
2873 }
2874
2875 fn classify_storage_error(err: Error) -> Error {
2880 if matches!(&err, Error::Namespace { .. }) {
2881 return err;
2882 }
2883 let detail = err.to_string();
2884 if let Error::IO { source, .. } = &err
2885 && let Some(os_err) = source.downcast_ref::<ObjectStoreError>()
2886 {
2887 if is_throttle_error(os_err) {
2888 return NamespaceError::Throttling {
2889 message: format!(
2890 "Storage request was throttled while resolving table: {detail}"
2891 ),
2892 }
2893 .into();
2894 }
2895 if Self::is_service_unavailable_error(os_err) {
2896 return NamespaceError::ServiceUnavailable {
2897 message: format!("Storage service unavailable while resolving table: {detail}"),
2898 }
2899 .into();
2900 }
2901 }
2902 NamespaceError::Internal {
2903 message: format!("Storage error while resolving table: {detail}"),
2904 }
2905 .into()
2906 }
2907
2908 fn is_service_unavailable_error(err: &ObjectStoreError) -> bool {
2912 if let ObjectStoreError::Generic { source, .. } = err {
2913 let message = source.to_string().to_ascii_lowercase();
2914 message.contains("503 service unavailable")
2915 || message.contains("502 bad gateway")
2916 || message.contains("504 gateway timeout")
2917 } else {
2918 false
2919 }
2920 }
2921
2922 fn is_manifest_table_absent_error(err: &Error) -> bool {
2926 if manifest::ManifestNamespace::is_not_found_load_error(err) {
2927 return true;
2928 }
2929 if let Error::Namespace { source, .. } = err
2930 && let Some(ns_err) = source.downcast_ref::<NamespaceError>()
2931 {
2932 return matches!(ns_err, NamespaceError::TableNotFound { .. });
2933 }
2934 false
2935 }
2936
2937 fn map_open_error(err: Error, not_found: NamespaceError) -> Error {
2941 if matches!(&err, Error::IO { .. })
2942 && !manifest::ManifestNamespace::is_not_found_load_error(&err)
2943 {
2944 return Self::classify_storage_error(err);
2945 }
2946 not_found.into()
2947 }
2948
2949 async fn get_storage_options_for_table(
2969 &self,
2970 table_uri: &str,
2971 vend_credentials: bool,
2972 identity: Option<&Identity>,
2973 ) -> Result<Option<HashMap<String, String>>> {
2974 if vend_credentials && let Some(ref vendor) = self.credential_vendor {
2975 let vended = vendor.vend_credentials(table_uri, identity).await?;
2976 return Ok(Some(vended.storage_options));
2977 }
2978 if self.vend_input_storage_options {
2981 let mut options = self.storage_options.clone().unwrap_or_default();
2982 if let Some(refresh_interval_millis) =
2984 self.vend_input_storage_options_refresh_interval_millis
2985 {
2986 let now_millis = std::time::SystemTime::now()
2987 .duration_since(std::time::UNIX_EPOCH)
2988 .unwrap()
2989 .as_millis() as u64;
2990 let expires_at_millis = now_millis + refresh_interval_millis;
2991 options.insert(
2992 "expires_at_millis".to_string(),
2993 expires_at_millis.to_string(),
2994 );
2995 }
2996 return Ok(Some(options));
2997 }
2998 Ok(None)
3001 }
3002
3003 pub async fn migrate(&self) -> Result<usize> {
3057 let Some(manifest_ns) = self.manifest_ns_for_write().await? else {
3059 return Ok(0); };
3061
3062 let manifest_locations = manifest_ns.list_manifest_table_locations().await?;
3064
3065 let dir_tables = self
3068 .filter_declared_tables(self.list_directory_tables().await?, false)
3069 .await?;
3070
3071 let mut migrated_count = 0;
3076 for table_name in dir_tables {
3077 let dir_name = format!("{}.lance", table_name);
3079 if !manifest_locations.contains(&dir_name) {
3080 manifest_ns.register_table(&table_name, dir_name).await?;
3081 migrated_count += 1;
3082 }
3083 }
3084
3085 Ok(migrated_count)
3086 }
3087
3088 async fn delete_physical_version_files(
3096 &self,
3097 table_entries: &[TableDeleteEntry],
3098 branch: Option<&str>,
3099 ) -> Result<i64> {
3100 let mut deleted_count = 0i64;
3101 for te in table_entries {
3102 let table_uri = self.resolve_table_location(&te.table_id).await?;
3103 let table_uri = match branch {
3104 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3105 None => table_uri,
3106 };
3107 let table_path = self.object_store_path_from_uri(&table_uri)?;
3108 let versions_dir_path = table_path.clone().join(VERSIONS_DIR);
3109
3110 let manifest_metas: Vec<_> = self
3112 .object_store
3113 .read_dir_all(&versions_dir_path, None)
3114 .try_collect()
3115 .await
3116 .map_err(|e| {
3117 lance_core::Error::from(NamespaceError::Internal {
3118 message: format!(
3119 "Failed to list manifest files for table at '{}': {}",
3120 table_uri, e
3121 ),
3122 })
3123 })?;
3124 let location_by_version: HashMap<u64, Path> = manifest_metas
3125 .into_iter()
3126 .filter_map(|meta| {
3127 let version = Self::manifest_version_from_filename(meta.location.filename()?)?;
3128 Some((version, meta.location))
3129 })
3130 .collect();
3131
3132 for (&v, version_path) in &location_by_version {
3133 let vi = v as i64;
3134 if !te.ranges.iter().any(|&(s, e)| vi >= s && (e < 0 || vi < e)) {
3135 continue;
3136 }
3137 match self.object_store.inner.delete(version_path).await {
3138 Ok(_) => {
3139 deleted_count += 1;
3140 }
3141 Err(object_store::Error::NotFound { .. }) => {}
3142 Err(e) => {
3143 return Err(NamespaceError::Internal {
3144 message: format!(
3145 "Failed to delete version {} for table at '{}': {}",
3146 v, table_uri, e
3147 ),
3148 }
3149 .into());
3150 }
3151 }
3152 }
3153 }
3154 Ok(deleted_count)
3155 }
3156
3157 #[allow(clippy::too_many_arguments)]
3162 fn apply_query_params_to_scanner(
3163 scanner: &mut Scanner,
3164 filter: Option<&str>,
3165 columns: Option<&QueryTableRequestColumns>,
3166 vector_column: Option<&str>,
3167 vector: &QueryTableRequestVector,
3168 k: i32,
3169 offset: Option<i32>,
3170 prefilter: Option<bool>,
3171 bypass_vector_index: Option<bool>,
3172 nprobes: Option<i32>,
3173 ef: Option<i32>,
3174 refine_factor: Option<i32>,
3175 distance_type: Option<&str>,
3176 fast_search_flag: Option<bool>,
3177 with_row_id: Option<bool>,
3178 lower_bound: Option<f32>,
3179 upper_bound: Option<f32>,
3180 operation: &str,
3181 ) -> Result<()> {
3182 if let Some(pf) = prefilter {
3184 scanner.prefilter(pf);
3185 }
3186
3187 if let Some(filter) = filter {
3188 scanner.filter(filter).map_err(|e| {
3189 Error::invalid_input_source(
3190 format!("Invalid filter expression for {}: {}", operation, e).into(),
3191 )
3192 })?;
3193 }
3194
3195 if let Some(cols) = columns {
3196 if let Some(ref names) = cols.column_names {
3197 scanner.project(names.as_slice()).map_err(|e| {
3198 Error::invalid_input_source(
3199 format!("Invalid column projection for {}: {}", operation, e).into(),
3200 )
3201 })?;
3202 } else if let Some(ref aliases) = cols.column_aliases {
3203 let pairs: Vec<(&str, &str)> = aliases
3205 .iter()
3206 .map(|(alias, src)| (alias.as_str(), src.as_str()))
3207 .collect();
3208 scanner.project_with_transform(&pairs).map_err(|e| {
3209 Error::invalid_input_source(
3210 format!("Invalid column aliases for {}: {}", operation, e).into(),
3211 )
3212 })?;
3213 }
3214 }
3215
3216 let query_vec: Option<Vec<f32>> = vector
3218 .single_vector
3219 .as_ref()
3220 .filter(|v| !v.is_empty())
3221 .cloned()
3222 .or_else(|| {
3223 vector
3224 .multi_vector
3225 .as_ref()
3226 .and_then(|mv| mv.first())
3227 .filter(|v| !v.is_empty())
3228 .cloned()
3229 });
3230
3231 if let Some(q_vec) = query_vec {
3232 let col = vector_column.unwrap_or("vector");
3233 let q = Arc::new(Float32Array::from(q_vec));
3234 scanner
3235 .nearest(col, q.as_ref(), k.max(1) as usize)
3236 .map_err(|e| {
3237 Error::invalid_input_source(
3238 format!("Invalid vector query for {}: {}", operation, e).into(),
3239 )
3240 })?;
3241
3242 if let Some(n) = nprobes {
3244 scanner.nprobes(n.max(1) as usize);
3245 }
3246 if let Some(e) = ef {
3247 scanner.ef(e.max(1) as usize);
3248 }
3249 if let Some(rf) = refine_factor {
3250 scanner.refine(rf.max(0) as u32);
3251 }
3252 if let Some(true) = bypass_vector_index {
3254 scanner.use_index(false);
3255 }
3256 if let Some(true) = fast_search_flag {
3257 scanner.fast_search();
3258 }
3259 if lower_bound.is_some() || upper_bound.is_some() {
3260 scanner.distance_range(lower_bound, upper_bound);
3261 }
3262 if let Some(dt) = distance_type {
3263 let metric = Self::parse_metric_type(Some(dt))?;
3264 scanner.distance_metric(metric);
3265 }
3266 if let Some(off) = offset.filter(|&o| o > 0) {
3268 scanner.limit(None, Some(off as i64)).map_err(|e| {
3269 Error::invalid_input_source(
3270 format!("Invalid offset for {}: {}", operation, e).into(),
3271 )
3272 })?;
3273 }
3274 } else {
3275 let limit = if k > 0 { Some(k as i64) } else { None };
3277 scanner
3278 .limit(limit, offset.map(|o| o as i64))
3279 .map_err(|e| {
3280 Error::invalid_input_source(
3281 format!("Invalid limit/offset for {}: {}", operation, e).into(),
3282 )
3283 })?;
3284 }
3285
3286 if let Some(true) = with_row_id {
3287 scanner.with_row_id();
3288 }
3289
3290 Ok(())
3291 }
3292
3293 pub fn retrieve_ops_metrics(&self) -> HashMap<String, u64> {
3300 self.ops_metrics
3301 .as_ref()
3302 .map(|m| m.retrieve())
3303 .unwrap_or_default()
3304 }
3305
3306 pub fn reset_ops_metrics(&self) {
3310 if let Some(ref metrics) = self.ops_metrics {
3311 metrics.reset();
3312 }
3313 }
3314
3315 fn record_op(&self, operation: &str) {
3317 if let Some(ref metrics) = self.ops_metrics {
3318 metrics.increment(operation);
3319 }
3320 }
3321}
3322
3323#[async_trait]
3324impl LanceNamespace for DirectoryNamespace {
3325 async fn list_namespaces(
3326 &self,
3327 request: ListNamespacesRequest,
3328 ) -> Result<ListNamespacesResponse> {
3329 self.record_op("list_namespaces");
3330 self.ensure_read_manifest().await?;
3331 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3332 return manifest_ns.list_namespaces(request).await;
3333 }
3334
3335 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3336 return Err(self.child_namespace_requires_manifest_error());
3337 }
3338 Self::validate_root_namespace_id(&request.id)?;
3339 Ok(ListNamespacesResponse::new(vec![]))
3340 }
3341
3342 async fn describe_namespace(
3343 &self,
3344 request: DescribeNamespaceRequest,
3345 ) -> Result<DescribeNamespaceResponse> {
3346 self.record_op("describe_namespace");
3347 self.ensure_read_manifest().await?;
3348 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3349 return manifest_ns.describe_namespace(request).await;
3350 }
3351
3352 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3353 return Err(self.child_namespace_requires_manifest_error());
3354 }
3355 Self::validate_root_namespace_id(&request.id)?;
3356 #[allow(clippy::needless_update)]
3357 Ok(DescribeNamespaceResponse {
3358 properties: Some(HashMap::new()),
3359 ..Default::default()
3360 })
3361 }
3362
3363 async fn create_namespace(
3364 &self,
3365 request: CreateNamespaceRequest,
3366 ) -> Result<CreateNamespaceResponse> {
3367 self.record_op("create_namespace");
3368 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3369 return manifest_ns.create_namespace(request).await;
3370 }
3371
3372 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3373 return Err(NamespaceError::NamespaceAlreadyExists {
3374 message: "root namespace".to_string(),
3375 }
3376 .into());
3377 }
3378
3379 Err(NamespaceError::Unsupported {
3380 message: "Child namespaces are only supported when manifest mode is enabled"
3381 .to_string(),
3382 }
3383 .into())
3384 }
3385
3386 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
3387 self.record_op("drop_namespace");
3388 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3389 return manifest_ns.drop_namespace(request).await;
3390 }
3391
3392 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3393 return Err(NamespaceError::InvalidInput {
3394 message: "Root namespace cannot be dropped".to_string(),
3395 }
3396 .into());
3397 }
3398
3399 Err(NamespaceError::Unsupported {
3400 message: "Child namespaces are only supported when manifest mode is enabled"
3401 .to_string(),
3402 }
3403 .into())
3404 }
3405
3406 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
3407 self.record_op("namespace_exists");
3408 self.ensure_read_manifest().await?;
3409 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3410 return manifest_ns.namespace_exists(request).await;
3411 }
3412
3413 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3414 return Ok(());
3415 }
3416
3417 Err(self.child_namespace_requires_manifest_error())
3418 }
3419
3420 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
3421 self.record_op("list_tables");
3422 let namespace_id = request.id.as_ref().ok_or_else(|| {
3424 lance_core::Error::from(NamespaceError::InvalidInput {
3425 message: "Namespace ID is required".to_string(),
3426 })
3427 })?;
3428
3429 if !namespace_id.is_empty()
3435 || !self.dir_listing_enabled
3436 || self.dir_listing_to_manifest_migration_enabled
3437 {
3438 self.ensure_read_manifest().await?;
3439 }
3440
3441 if !namespace_id.is_empty() {
3443 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3444 return manifest_ns.list_tables(request).await;
3445 }
3446 return Err(self.child_namespace_requires_manifest_error());
3447 }
3448
3449 if let Some(manifest_ns) = self.manifest_ns_for_read()
3451 && !self.dir_listing_enabled
3452 {
3453 return manifest_ns.list_tables(request).await;
3454 }
3455 if !self.dir_listing_enabled {
3456 return Ok(ListTablesResponse::new(vec![]));
3457 }
3458
3459 let mut tables = if self.manifest_ns_for_read().is_some()
3462 && self.dir_listing_enabled
3463 && self.dir_listing_to_manifest_migration_enabled
3464 {
3465 let manifest_locations = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3467 manifest_ns.list_manifest_table_locations().await?
3468 } else {
3469 std::collections::HashSet::new()
3470 };
3471
3472 let mut manifest_request = request.clone();
3474 manifest_request.limit = None;
3475 manifest_request.page_token = None;
3476 let manifest_tables = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3477 let manifest_response = manifest_ns.list_tables(manifest_request).await?;
3478 manifest_response.tables
3479 } else {
3480 vec![]
3481 };
3482
3483 let mut all_tables: Vec<String> = manifest_tables;
3486 let dir_tables = self.list_directory_tables().await?;
3487 for table_name in dir_tables {
3488 let full_location = format!("{}/{}.lance", self.root, table_name);
3491 let relative_location = format!("{}.lance", table_name);
3492 if !manifest_locations.contains(&full_location)
3493 && !manifest_locations.contains(&relative_location)
3494 {
3495 all_tables.push(table_name);
3496 }
3497 }
3498
3499 all_tables
3500 } else {
3501 self.list_directory_tables().await?
3502 };
3503
3504 tables = self
3505 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
3506 .await?;
3507
3508 let next_page_token =
3510 Self::apply_pagination(&mut tables, request.page_token, request.limit);
3511 let mut response = ListTablesResponse::new(tables);
3512 response.page_token = next_page_token;
3513 Ok(response)
3514 }
3515
3516 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
3517 self.record_op("describe_table");
3518 self.describe_table_impl(request).await
3519 }
3520
3521 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
3522 self.record_op("table_exists");
3523 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
3524 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
3525 let skip_manifest_for_root = self.dir_listing_enabled
3526 && is_root_level
3527 && !self.dir_listing_to_manifest_migration_enabled;
3528 if is_child_table
3531 || !self.dir_listing_enabled
3532 || self.dir_listing_to_manifest_migration_enabled
3533 {
3534 self.ensure_read_manifest().await?;
3535 }
3536 if let Some(manifest_ns) = self.manifest_ns_for_read()
3537 && !skip_manifest_for_root
3538 {
3539 match manifest_ns.table_exists(request.clone()).await {
3540 Ok(()) => return Ok(()),
3541 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
3542 return Err(e);
3545 }
3546 Err(e) if self.dir_listing_enabled && is_root_level => {
3547 if !Self::is_manifest_table_absent_error(&e) {
3551 return Err(Self::classify_storage_error(e));
3552 }
3553 }
3554 Err(e) => return Err(e),
3555 }
3556 }
3557 if is_child_table {
3558 return Err(self.child_namespace_requires_manifest_error());
3559 }
3560
3561 let table_name = Self::table_name_from_id(&request.id)?;
3562 let table_id = Self::format_table_id_from_request(&request.id);
3563 if !self.dir_listing_enabled {
3564 return Err(NamespaceError::TableNotFound { message: table_id }.into());
3565 }
3566
3567 let status = self.check_table_status(&table_name).await?;
3569
3570 if !status.exists {
3571 return Err(NamespaceError::TableNotFound {
3572 message: table_id.clone(),
3573 }
3574 .into());
3575 }
3576
3577 if status.is_deregistered {
3578 return Err(NamespaceError::TableNotFound {
3579 message: format!("Table is deregistered: {}", table_id),
3580 }
3581 .into());
3582 }
3583
3584 Ok(())
3585 }
3586
3587 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3588 self.record_op("drop_table");
3589 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3590 return manifest_ns.drop_table(request).await;
3591 }
3592
3593 let table_name = Self::table_name_from_id(&request.id)?;
3594 let table_uri = self.table_full_uri(&table_name);
3595 let table_path = self.table_path(&table_name);
3596
3597 self.object_store
3598 .remove_dir_all(table_path)
3599 .await
3600 .map_err(|e| {
3601 lance_core::Error::from(NamespaceError::Internal {
3602 message: format!("Failed to drop table {}: {:?}", table_name, e),
3603 })
3604 })?;
3605
3606 Ok(DropTableResponse {
3607 id: request.id,
3608 location: Some(table_uri),
3609 ..Default::default()
3610 })
3611 }
3612
3613 async fn create_table(
3614 &self,
3615 request: CreateTableRequest,
3616 request_data: Bytes,
3617 ) -> Result<CreateTableResponse> {
3618 self.record_op("create_table");
3619 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3620 return manifest_ns.create_table(request, request_data).await;
3621 }
3622
3623 Self::validate_dir_only_properties(request.properties.as_ref(), "create_table")?;
3624
3625 let table_name = Self::table_name_from_id(&request.id)?;
3626 let table_uri = self.table_full_uri(&table_name);
3627 let status = self.check_table_status(&table_name).await?;
3628 let (reader, _num_rows) =
3629 Self::ipc_reader_from_request_data(&request_data, "create_table")?;
3630
3631 if status.exists && self.table_has_actual_manifests(&table_name).await? {
3632 return Err(NamespaceError::TableAlreadyExists {
3633 message: table_name,
3634 }
3635 .into());
3636 }
3637
3638 let write_result = self
3639 .write_reader_to_table(
3640 &table_uri,
3641 reader,
3642 WriteMode::Create,
3643 request.storage_options.clone(),
3644 )
3645 .await;
3646 if let Err(err) = write_result {
3647 if self.table_uri_has_actual_manifests(&table_uri).await? {
3648 return Err(NamespaceError::TableAlreadyExists {
3649 message: table_name,
3650 }
3651 .into());
3652 }
3653 return Err(err);
3654 }
3655 Ok(CreateTableResponse {
3656 version: Some(1),
3657 location: Some(table_uri),
3658 storage_options: self.storage_options.clone(),
3659 properties: request.properties,
3660 ..Default::default()
3661 })
3662 }
3663
3664 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3665 self.record_op("declare_table");
3666 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3667 let mut response = manifest_ns.declare_table(request.clone()).await?;
3668 if let Some(ref location) = response.location {
3669 let vend = request.vend_credentials.unwrap_or(true);
3671 let identity = request.identity.as_deref();
3672 response.storage_options = self
3673 .get_storage_options_for_table(location, vend, identity)
3674 .await?;
3675 }
3676 if self.table_version_tracking_enabled {
3678 response.managed_versioning = Some(true);
3679 }
3680 return Ok(response);
3681 }
3682
3683 Self::validate_dir_only_properties(request.properties.as_ref(), "declare_table")?;
3684
3685 let table_name = Self::table_name_from_id(&request.id)?;
3686 let table_uri = self.table_full_uri(&table_name);
3687
3688 if let Some(location) = &request.location {
3690 let location = location.trim_end_matches('/');
3691 if location != table_uri {
3692 return Err(NamespaceError::InvalidInput {
3693 message: format!(
3694 "Cannot declare table {} at location {}, must be at location {}",
3695 table_name, location, table_uri
3696 ),
3697 }
3698 .into());
3699 }
3700 }
3701
3702 let status = self.check_table_status(&table_name).await?;
3706 if status.exists && !status.has_reserved_file {
3707 return Err(NamespaceError::TableAlreadyExists {
3709 message: table_name.to_string(),
3710 }
3711 .into());
3712 }
3713
3714 let reserved_file_path = self.table_reserved_file_path(&table_name);
3718
3719 put_marker_file_atomic(
3720 &self.object_store,
3721 &reserved_file_path,
3722 &format!("table {}", table_name),
3723 )
3724 .await
3725 .map_err(|e| match e {
3726 MarkerFileError::AlreadyExists { .. } => {
3727 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3728 message: table_name.to_string(),
3729 })
3730 }
3731 MarkerFileError::Other { message } => {
3732 lance_core::Error::from(NamespaceError::Internal { message })
3733 }
3734 })?;
3735
3736 let vend_credentials = request.vend_credentials.unwrap_or(true);
3738 let identity = request.identity.as_deref();
3739 let storage_options = self
3740 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
3741 .await?;
3742
3743 Ok(DeclareTableResponse {
3744 location: Some(table_uri),
3745 storage_options,
3746 properties: request.properties,
3747 managed_versioning: if self.table_version_tracking_enabled {
3748 Some(true)
3749 } else {
3750 None
3751 },
3752 ..Default::default()
3753 })
3754 }
3755
3756 async fn register_table(
3757 &self,
3758 request: lance_namespace::models::RegisterTableRequest,
3759 ) -> Result<lance_namespace::models::RegisterTableResponse> {
3760 self.record_op("register_table");
3761 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3763 return LanceNamespace::register_table(manifest_ns.as_ref(), request).await;
3764 }
3765
3766 Err(NamespaceError::Unsupported {
3768 message: "register_table is only supported when manifest mode is enabled".to_string(),
3769 }
3770 .into())
3771 }
3772
3773 async fn deregister_table(
3774 &self,
3775 request: lance_namespace::models::DeregisterTableRequest,
3776 ) -> Result<lance_namespace::models::DeregisterTableResponse> {
3777 self.record_op("deregister_table");
3778 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3780 return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await;
3781 }
3782
3783 let table_name = Self::table_name_from_id(&request.id)?;
3785 let table_uri = self.table_full_uri(&table_name);
3786
3787 let status = self.check_table_status(&table_name).await?;
3790
3791 if !status.exists {
3792 return Err(NamespaceError::TableNotFound {
3793 message: table_name.to_string(),
3794 }
3795 .into());
3796 }
3797
3798 if status.is_deregistered {
3799 return Err(NamespaceError::TableNotFound {
3800 message: format!("Table is already deregistered: {}", table_name),
3801 }
3802 .into());
3803 }
3804
3805 let deregistered_path = self.table_deregistered_file_path(&table_name);
3811 put_marker_file_atomic(
3812 &self.object_store,
3813 &deregistered_path,
3814 &format!("deregistration marker for table {}", table_name),
3815 )
3816 .await
3817 .map_err(|e| match e {
3818 MarkerFileError::AlreadyExists { .. } => {
3819 lance_core::Error::from(NamespaceError::InvalidTableState {
3820 message: format!("Table is already deregistered: {}", table_name),
3821 })
3822 }
3823 MarkerFileError::Other { message } => {
3824 lance_core::Error::from(NamespaceError::Internal { message })
3825 }
3826 })?;
3827
3828 Ok(lance_namespace::models::DeregisterTableResponse {
3829 id: request.id,
3830 location: Some(table_uri),
3831 ..Default::default()
3832 })
3833 }
3834
3835 async fn alter_table_add_columns(
3836 &self,
3837 request: AlterTableAddColumnsRequest,
3838 ) -> Result<AlterTableAddColumnsResponse> {
3839 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3840 return manifest_ns.alter_table_add_columns(request).await;
3841 }
3842
3843 let table_name = Self::table_name_from_id(&request.id)?;
3845 let table_uri = self.table_full_uri(&table_name);
3846
3847 let status = self.check_table_status(&table_name).await?;
3849 if !status.exists {
3850 return Err(NamespaceError::TableNotFound {
3851 message: table_name,
3852 }
3853 .into());
3854 }
3855 if status.is_deregistered {
3856 return Err(NamespaceError::TableNotFound {
3857 message: format!("Table is deregistered: {}", table_name),
3858 }
3859 .into());
3860 }
3861
3862 let mut dataset = self
3863 .configured_builder(&table_uri)
3864 .load()
3865 .await
3866 .map_err(|e| {
3867 Error::io_source(box_error(std::io::Error::other(format!(
3868 "Failed to open dataset: {}",
3869 e
3870 ))))
3871 })?;
3872
3873 let sql_expressions = build_sql_expressions(&request.new_columns)?;
3874
3875 dataset
3876 .add_columns(
3877 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3878 None,
3879 None,
3880 )
3881 .await
3882 .map_err(|e| {
3883 Error::io_source(box_error(std::io::Error::other(format!(
3884 "Failed to add columns: {}",
3885 e
3886 ))))
3887 })?;
3888
3889 let version = dataset.version().version as i64;
3890 Ok(AlterTableAddColumnsResponse::new(version))
3891 }
3892
3893 async fn alter_table_alter_columns(
3894 &self,
3895 request: AlterTableAlterColumnsRequest,
3896 ) -> Result<AlterTableAlterColumnsResponse> {
3897 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3898 return manifest_ns.alter_table_alter_columns(request).await;
3899 }
3900
3901 let table_name = Self::table_name_from_id(&request.id)?;
3902 let table_uri = self.table_full_uri(&table_name);
3903
3904 let status = self.check_table_status(&table_name).await?;
3906 if !status.exists {
3907 return Err(NamespaceError::TableNotFound {
3908 message: table_name,
3909 }
3910 .into());
3911 }
3912 if status.is_deregistered {
3913 return Err(NamespaceError::TableNotFound {
3914 message: format!("Table is deregistered: {}", table_name),
3915 }
3916 .into());
3917 }
3918
3919 let mut dataset = self
3920 .configured_builder(&table_uri)
3921 .load()
3922 .await
3923 .map_err(|e| {
3924 Error::io_source(box_error(std::io::Error::other(format!(
3925 "Failed to open dataset: {}",
3926 e
3927 ))))
3928 })?;
3929
3930 let alterations = build_column_alterations(&request.alterations)?;
3931
3932 dataset.alter_columns(&alterations).await.map_err(|e| {
3933 Error::io_source(box_error(std::io::Error::other(format!(
3934 "Failed to alter columns: {}",
3935 e
3936 ))))
3937 })?;
3938
3939 let version = dataset.version().version as i64;
3940 Ok(AlterTableAlterColumnsResponse::new(version))
3941 }
3942
3943 async fn alter_table_drop_columns(
3944 &self,
3945 request: AlterTableDropColumnsRequest,
3946 ) -> Result<AlterTableDropColumnsResponse> {
3947 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3948 return manifest_ns.alter_table_drop_columns(request).await;
3949 }
3950
3951 let table_name = Self::table_name_from_id(&request.id)?;
3952 let table_uri = self.table_full_uri(&table_name);
3953
3954 let status = self.check_table_status(&table_name).await?;
3956 if !status.exists {
3957 return Err(NamespaceError::TableNotFound {
3958 message: table_name,
3959 }
3960 .into());
3961 }
3962 if status.is_deregistered {
3963 return Err(NamespaceError::TableNotFound {
3964 message: format!("Table is deregistered: {}", table_name),
3965 }
3966 .into());
3967 }
3968
3969 let mut dataset = self
3970 .configured_builder(&table_uri)
3971 .load()
3972 .await
3973 .map_err(|e| {
3974 Error::io_source(box_error(std::io::Error::other(format!(
3975 "Failed to open dataset: {}",
3976 e
3977 ))))
3978 })?;
3979
3980 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
3981 dataset.drop_columns(&columns).await.map_err(|e| {
3982 Error::io_source(box_error(std::io::Error::other(format!(
3983 "Failed to drop columns: {}",
3984 e
3985 ))))
3986 })?;
3987
3988 let version = dataset.version().version as i64;
3989 Ok(AlterTableDropColumnsResponse::new(version))
3990 }
3991
3992 async fn list_table_versions(
3993 &self,
3994 request: ListTableVersionsRequest,
3995 ) -> Result<ListTableVersionsResponse> {
3996 self.record_op("list_table_versions");
3997 let branch = Self::normalized_branch(request.branch.as_deref())?;
3998 let table_uri = self.resolve_table_location(&request.id).await?;
3999 let table_uri = match branch {
4000 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
4001 None => table_uri,
4002 };
4003 let want_descending = request.descending == Some(true);
4004 let table_versions = self
4005 .list_table_versions_from_storage(&table_uri, want_descending, request.limit)
4006 .await?;
4007
4008 Ok(ListTableVersionsResponse {
4009 versions: table_versions,
4010 page_token: None,
4011 ..Default::default()
4012 })
4013 }
4014
4015 async fn create_table_version(
4016 &self,
4017 request: CreateTableVersionRequest,
4018 ) -> Result<CreateTableVersionResponse> {
4019 self.record_op("create_table_version");
4020 let branch = Self::normalized_branch(request.branch.as_deref())?;
4021 let table_uri = self.resolve_table_location(&request.id).await?;
4022 let (table_uri, table_path, branch_parent_version) = match branch {
4023 Some(b) => self.resolve_branch_for_commit(&table_uri, b).await?,
4024 None => {
4025 let table_path = self.object_store_path_from_uri(&table_uri)?;
4026 (table_uri, table_path, None)
4027 }
4028 };
4029
4030 let staging_manifest_path = &request.manifest_path;
4031 let version = request.version as u64;
4032
4033 let naming_scheme = match request.naming_scheme.as_deref() {
4035 Some("V1") => ManifestNamingScheme::V1,
4036 _ => ManifestNamingScheme::V2,
4037 };
4038
4039 let final_path = naming_scheme.manifest_path(&table_path, version);
4041
4042 let staging_path = Path::parse(staging_manifest_path).map_err(|e| {
4043 lance_core::Error::from(NamespaceError::InvalidInput {
4044 message: format!(
4045 "Invalid staging manifest path '{}': {}",
4046 staging_manifest_path, e
4047 ),
4048 })
4049 })?;
4050
4051 match self.object_store.inner.head(&final_path).await {
4053 Ok(existing_meta) => {
4054 return self
4055 .resolve_existing_table_version(ExistingTableVersionResolve {
4056 staging_path: &staging_path,
4057 final_path: &final_path,
4058 version,
4059 table_uri: &table_uri,
4060 final_meta: &existing_meta,
4061 request_manifest_size: request.manifest_size,
4062 })
4063 .await;
4064 }
4065 Err(ObjectStoreError::NotFound { .. }) => {}
4066 Err(e) => {
4067 return Err(lance_core::Error::from(NamespaceError::Internal {
4068 message: format!(
4069 "Failed to stat version {} for table at '{}': {}",
4070 version, table_uri, e
4071 ),
4072 }));
4073 }
4074 }
4075
4076 let is_branch = branch.is_some();
4079 self.enforce_create_table_version_cas(
4080 &table_path,
4081 version,
4082 &table_uri,
4083 is_branch,
4084 branch_parent_version,
4085 )
4086 .await?;
4087
4088 let copy_result = self
4090 .materialize_version_manifest_create(&staging_path, &final_path, staging_manifest_path)
4091 .await;
4092
4093 match copy_result {
4094 Ok(()) => {}
4095 Err(ObjectStoreError::AlreadyExists { .. })
4096 | Err(ObjectStoreError::Precondition { .. }) => {
4097 let existing_meta = self.object_store.inner.head(&final_path).await.map_err(|e| {
4099 lance_core::Error::from(NamespaceError::Internal {
4100 message: format!(
4101 "Version {} conflict for table at '{}' but failed to stat winner: {}",
4102 version, table_uri, e
4103 ),
4104 })
4105 })?;
4106 return self
4107 .resolve_existing_table_version(ExistingTableVersionResolve {
4108 staging_path: &staging_path,
4109 final_path: &final_path,
4110 version,
4111 table_uri: &table_uri,
4112 final_meta: &existing_meta,
4113 request_manifest_size: request.manifest_size,
4114 })
4115 .await;
4116 }
4117 Err(ObjectStoreError::NotFound { .. }) => {
4118 return Err(lance_core::Error::from(NamespaceError::InvalidInput {
4119 message: format!(
4120 "Staging manifest not found at '{}' for version {} of table at '{}'",
4121 staging_manifest_path, version, table_uri
4122 ),
4123 }));
4124 }
4125 Err(e) => {
4126 return Err(lance_core::Error::from(NamespaceError::Internal {
4127 message: format!(
4128 "Failed to create version {} for table at '{}': {}",
4129 version, table_uri, e
4130 ),
4131 }));
4132 }
4133 }
4134
4135 let final_meta = self
4136 .object_store
4137 .inner
4138 .head(&final_path)
4139 .await
4140 .map_err(|e| {
4141 lance_core::Error::from(NamespaceError::Internal {
4142 message: format!(
4143 "Failed to stat created version {} for table at '{}': {}",
4144 version, table_uri, e
4145 ),
4146 })
4147 })?;
4148
4149 if let Err(e) = self.object_store.inner.delete(&staging_path).await {
4151 log::warn!(
4152 "Failed to delete staging manifest at '{}': {:?}",
4153 staging_path,
4154 e
4155 );
4156 }
4157
4158 Ok(Self::create_table_version_response(
4159 version,
4160 &final_path,
4161 &final_meta,
4162 ))
4163 }
4164
4165 async fn describe_table_version(
4166 &self,
4167 request: DescribeTableVersionRequest,
4168 ) -> Result<DescribeTableVersionResponse> {
4169 self.record_op("describe_table_version");
4170 let branch = Self::normalized_branch(request.branch.as_deref())?;
4171 let table_uri = self.resolve_table_location(&request.id).await?;
4172 let table_uri = match branch {
4173 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
4174 None => table_uri,
4175 };
4176 let versions = self
4177 .list_table_versions_from_storage(&table_uri, true, None)
4178 .await?;
4179 let table_version = if let Some(requested_version) = request.version {
4180 versions
4181 .into_iter()
4182 .find(|version| version.version == requested_version)
4183 .ok_or_else(|| {
4184 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4185 message: format!(
4186 "version {} for table {}",
4187 requested_version,
4188 Self::format_table_id_from_request(&request.id)
4189 ),
4190 })
4191 })?
4192 } else {
4193 versions.into_iter().next().ok_or_else(|| {
4194 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4195 message: format!(
4196 "latest version for table {}",
4197 Self::format_table_id_from_request(&request.id)
4198 ),
4199 })
4200 })?
4201 };
4202
4203 Ok(DescribeTableVersionResponse {
4204 version: Box::new(table_version),
4205 ..Default::default()
4206 })
4207 }
4208
4209 async fn batch_delete_table_versions(
4210 &self,
4211 request: BatchDeleteTableVersionsRequest,
4212 ) -> Result<BatchDeleteTableVersionsResponse> {
4213 self.record_op("batch_delete_table_versions");
4214 let branch = Self::normalized_branch(request.branch.as_deref())?;
4215 let ranges: Vec<(i64, i64)> = request
4218 .ranges
4219 .iter()
4220 .map(|r| (r.start_version, r.end_version))
4221 .collect();
4222
4223 const MAX_VERSIONS_PER_REQUEST: i128 = 1_000_000;
4227 let requested: i128 = ranges
4228 .iter()
4229 .map(|(s, e)| {
4230 if *e < 0 {
4231 0
4232 } else {
4233 (*e as i128 - *s as i128).max(0)
4234 }
4235 })
4236 .sum();
4237 if requested > MAX_VERSIONS_PER_REQUEST {
4238 return Err(NamespaceError::InvalidInput {
4239 message: format!(
4240 "batch_delete requested {} versions; limit is {}",
4241 requested, MAX_VERSIONS_PER_REQUEST
4242 ),
4243 }
4244 .into());
4245 }
4246
4247 let table_entries = vec![TableDeleteEntry {
4248 table_id: request.id.clone(),
4249 ranges,
4250 }];
4251
4252 let total_deleted_count = self
4253 .delete_physical_version_files(&table_entries, branch)
4254 .await?;
4255
4256 Ok(BatchDeleteTableVersionsResponse {
4257 deleted_count: Some(total_deleted_count),
4258 transaction_id: None,
4259 ..Default::default()
4260 })
4261 }
4262
4263 async fn create_table_index(
4264 &self,
4265 request: CreateTableIndexRequest,
4266 ) -> Result<CreateTableIndexResponse> {
4267 self.record_op("create_table_index");
4268 let table_uri = self.resolve_table_location(&request.id).await?;
4269 let mut dataset = self
4270 .load_dataset(&table_uri, None, "create_table_index")
4271 .await?;
4272 let index_request = Self::build_index_params(&request)?;
4273
4274 dataset
4275 .create_index(
4276 &[request.column.as_str()],
4277 index_request.index_type(),
4278 request.name.clone(),
4279 index_request.params(),
4280 false,
4281 )
4282 .await
4283 .map_err(|e| {
4284 let err_msg = format!("{}", e);
4285 let ns_err = if err_msg.contains("already exists") {
4286 NamespaceError::TableIndexAlreadyExists {
4287 message: format!(
4288 "Index '{}' already exists on table '{}': {:?}",
4289 request.name.as_deref().unwrap_or("<auto-generated>"),
4290 table_uri,
4291 e
4292 ),
4293 }
4294 } else if err_msg.contains("not found") || err_msg.contains("does not exist") {
4295 NamespaceError::TableColumnNotFound {
4296 message: format!(
4297 "Column '{}' not found for table '{}': {:?}",
4298 request.column, table_uri, e
4299 ),
4300 }
4301 } else {
4302 NamespaceError::Internal {
4303 message: format!(
4304 "Failed to create {} index '{}' on column '{}' for table '{}': {:?}",
4305 request.index_type,
4306 request.name.as_deref().unwrap_or("<auto-generated>"),
4307 request.column,
4308 table_uri,
4309 e
4310 ),
4311 }
4312 };
4313 lance_core::Error::from(ns_err)
4314 })?;
4315
4316 let transaction_id = dataset
4317 .read_transaction()
4318 .await
4319 .map_err(|e| {
4320 lance_core::Error::from(NamespaceError::Internal {
4321 message: format!(
4322 "Failed to read committed transaction after creating index on '{}': {}",
4323 table_uri, e
4324 ),
4325 })
4326 })?
4327 .map(|transaction| transaction.uuid);
4328
4329 Ok(CreateTableIndexResponse {
4330 transaction_id,
4331 ..Default::default()
4332 })
4333 }
4334
4335 async fn list_table_indices(
4336 &self,
4337 request: ListTableIndicesRequest,
4338 ) -> Result<ListTableIndicesResponse> {
4339 self.record_op("list_table_indices");
4340 let table_uri = self.resolve_table_location(&request.id).await?;
4341 let dataset = self
4342 .load_dataset(&table_uri, request.version, "list_table_indices")
4343 .await?;
4344 let total_rows = dataset.count_rows(None).await.map_err(|e| {
4345 lance_core::Error::from(NamespaceError::Internal {
4346 message: format!("Failed to count rows for table '{}': {:?}", table_uri, e),
4347 })
4348 })? as u64;
4349 let mut indices = dataset
4350 .describe_indices(None)
4351 .await
4352 .map_err(|e| {
4353 lance_core::Error::from(NamespaceError::Internal {
4354 message: format!("Failed to describe table indices for '{}': {:?}", table_uri, e),
4355 })
4356 })?
4357 .into_iter()
4358 .filter(|description| {
4359 description
4360 .metadata()
4361 .first()
4362 .map(|metadata| !is_system_index(metadata))
4363 .unwrap_or(false)
4364 })
4365 .map(|description| {
4366 let columns = description
4367 .field_ids()
4368 .iter()
4369 .map(|field_id| {
4370 dataset
4371 .schema()
4372 .field_path_minimal(i32::try_from(*field_id).map_err(|e| {
4373 lance_core::Error::from(NamespaceError::Internal {
4374 message: format!(
4375 "Field id {} does not fit in i32 for table '{}': {}",
4376 field_id, table_uri, e
4377 ),
4378 })
4379 })?)
4380 .map_err(|e| {
4381 lance_core::Error::from(NamespaceError::Internal {
4382 message: format!(
4383 "Failed to resolve field path for field_id {} in table '{}': {}",
4384 field_id, table_uri, e
4385 ),
4386 })
4387 })
4388 })
4389 .collect::<Result<Vec<_>>>()?;
4390
4391 let segments = description.segments();
4392 let created_at = segments
4393 .iter()
4394 .filter_map(|segment| segment.created_at)
4395 .min()
4396 .map(|ts| ts.to_rfc3339());
4397
4398 #[allow(clippy::needless_update)]
4401 let content = IndexContent {
4402 index_name: description.name().to_string(),
4403 index_uuid: description.metadata()[0].uuid.to_string(),
4404 columns,
4405 status: "SUCCEEDED".to_string(),
4406 index_type: Some(description.index_type().to_string()),
4407 type_url: Some(description.type_url().to_string()),
4408 num_indexed_rows: Some(description.rows_indexed() as i64),
4409 num_unindexed_rows: Some(
4410 total_rows.saturating_sub(description.rows_indexed()) as i64,
4411 ),
4412 size_bytes: description.total_size_bytes().map(|size| size as i64),
4413 num_segments: Some(segments.len() as i32),
4414 created_at,
4415 index_version: segments.first().map(|segment| segment.index_version),
4416 index_details: description.details().ok(),
4417 ..Default::default()
4418 };
4419 Ok(content)
4420 })
4421 .collect::<Result<Vec<_>>>()?;
4422
4423 let page_token = Self::paginate_indices(&mut indices, request.page_token, request.limit);
4424 Ok(ListTableIndicesResponse {
4425 indexes: indices,
4426 page_token,
4427 ..Default::default()
4428 })
4429 }
4430
4431 async fn describe_table_index_stats(
4432 &self,
4433 request: DescribeTableIndexStatsRequest,
4434 ) -> Result<DescribeTableIndexStatsResponse> {
4435 self.record_op("describe_table_index_stats");
4436 let table_uri = self.resolve_table_location(&request.id).await?;
4437 let dataset = self
4438 .load_dataset(&table_uri, request.version, "describe_table_index_stats")
4439 .await?;
4440 let index_name = request.index_name.as_deref().ok_or_else(|| {
4441 lance_core::Error::from(NamespaceError::InvalidInput {
4442 message: "Index name is required for describe_table_index_stats".to_string(),
4443 })
4444 })?;
4445 let metadatas = dataset
4446 .load_indices_by_name(index_name)
4447 .await
4448 .map_err(|e| {
4449 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4450 message: format!(
4451 "Failed to load index '{}' metadata for table '{}': {}",
4452 index_name, table_uri, e
4453 ),
4454 })
4455 })?;
4456 if metadatas.first().is_some_and(is_system_index) {
4457 return Err(NamespaceError::Unsupported {
4458 message: format!("System index '{}' is not exposed by this API", index_name),
4459 }
4460 .into());
4461 }
4462
4463 let stats = <Dataset as DatasetIndexExt>::index_statistics(&dataset, index_name)
4464 .await
4465 .map_err(|e| {
4466 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4467 message: format!(
4468 "Failed to describe index statistics for '{}' on table '{}': {}",
4469 index_name, table_uri, e
4470 ),
4471 })
4472 })?;
4473 let stats: serde_json::Value = serde_json::from_str(&stats).map_err(|e| {
4474 lance_core::Error::from(NamespaceError::Internal {
4475 message: format!(
4476 "Failed to parse index statistics for '{}' on table '{}': {}",
4477 index_name, table_uri, e
4478 ),
4479 })
4480 })?;
4481
4482 Ok(Self::describe_table_index_stats_response(&stats))
4483 }
4484
4485 async fn describe_transaction(
4486 &self,
4487 request: DescribeTransactionRequest,
4488 ) -> Result<DescribeTransactionResponse> {
4489 self.record_op("describe_transaction");
4490 let mut request_id = request.id.ok_or_else(|| {
4491 lance_core::Error::from(NamespaceError::InvalidInput {
4492 message: "Transaction id must include table id and transaction identifier"
4493 .to_string(),
4494 })
4495 })?;
4496 if request_id.len() < 2 {
4497 return Err(NamespaceError::InvalidInput {
4498 message: format!(
4499 "Transaction request id must include table id and transaction identifier, got {:?}",
4500 request_id
4501 ),
4502 }
4503 .into());
4504 }
4505
4506 let id = request_id.pop().expect("request_id len checked above");
4507 let table_id = Some(request_id);
4508 let table_uri = self.resolve_table_location(&table_id).await?;
4509 let dataset = self
4510 .load_dataset(&table_uri, None, "describe_transaction")
4511 .await?;
4512 let (version, transaction) = self.find_transaction(&dataset, &id).await?;
4513
4514 let sidecar = self
4517 .load_transaction_alteration(&table_uri, &transaction.uuid)
4518 .await?;
4519
4520 Ok(Self::transaction_response(version, &transaction, sidecar))
4521 }
4522
4523 async fn alter_transaction(
4524 &self,
4525 request: AlterTransactionRequest,
4526 ) -> Result<AlterTransactionResponse> {
4527 self.record_op("alter_transaction");
4528
4529 let mut request_id = request.id.ok_or_else(|| {
4531 lance_core::Error::from(NamespaceError::InvalidInput {
4532 message: "Transaction id must include table id and transaction identifier"
4533 .to_string(),
4534 })
4535 })?;
4536 if request_id.len() < 2 {
4537 return Err(NamespaceError::InvalidInput {
4538 message: format!(
4539 "Transaction request id must include table id and transaction identifier, got {:?}",
4540 request_id
4541 ),
4542 }
4543 .into());
4544 }
4545
4546 let txn_id = request_id.pop().expect("request_id len checked above");
4547 let table_id = Some(request_id);
4548 let table_uri = self.resolve_table_location(&table_id).await?;
4549 let dataset = self
4550 .load_dataset(&table_uri, None, "alter_transaction")
4551 .await?;
4552 let (version, transaction) = self.find_transaction(&dataset, &txn_id).await?;
4553
4554 const RESERVED_KEYS: &[&str] = &["uuid", "version", "read_version", "operation", "tag"];
4558 let is_reserved = |key: &str| RESERVED_KEYS.contains(&key);
4559
4560 let mut sidecar = self
4563 .load_transaction_alteration(&table_uri, &transaction.uuid)
4564 .await?
4565 .unwrap_or_default();
4566
4567 for action in &request.actions {
4568 if let Some(ref set_status) = action.set_status_action
4569 && let Some(ref status) = set_status.status
4570 {
4571 let normalized = status.to_lowercase().replace('_', "");
4573 match normalized.as_str() {
4574 "queued" | "running" | "succeeded" | "failed" | "canceled" => {
4575 sidecar.status = Some(status.clone());
4576 }
4577 _ => {
4578 return Err(NamespaceError::InvalidInput {
4579 message: format!(
4580 "Invalid transaction status '{}'. Valid values are: Queued, Running, Succeeded, Failed, Canceled",
4581 status
4582 ),
4583 }
4584 .into());
4585 }
4586 }
4587 }
4588
4589 if let Some(ref set_property) = action.set_property_action
4590 && let (Some(key), Some(value)) = (&set_property.key, &set_property.value)
4591 {
4592 if is_reserved(key) {
4593 return Err(NamespaceError::InvalidInput {
4594 message: format!("Property '{}' is reserved and cannot be modified", key),
4595 }
4596 .into());
4597 }
4598 let mode = set_property
4599 .mode
4600 .as_deref()
4601 .unwrap_or("Overwrite")
4602 .to_lowercase();
4603 match mode.as_str() {
4604 "overwrite" => {
4605 sidecar.properties.insert(key.clone(), value.clone());
4606 }
4607 "fail" => {
4608 let exists = sidecar.properties.contains_key(key)
4611 || transaction
4612 .transaction_properties
4613 .as_ref()
4614 .is_some_and(|props| props.contains_key(key));
4615 if exists {
4616 return Err(NamespaceError::ConcurrentModification {
4617 message: format!(
4618 "Property '{}' already exists and mode is 'Fail'",
4619 key
4620 ),
4621 }
4622 .into());
4623 }
4624 sidecar.properties.insert(key.clone(), value.clone());
4625 }
4626 "skip" => {
4627 let exists = sidecar.properties.contains_key(key)
4628 || transaction
4629 .transaction_properties
4630 .as_ref()
4631 .is_some_and(|props| props.contains_key(key));
4632 if !exists {
4633 sidecar.properties.insert(key.clone(), value.clone());
4634 }
4635 }
4636 _ => {
4637 return Err(NamespaceError::InvalidInput {
4638 message: format!(
4639 "Invalid set_property mode '{}'. Valid values are: Overwrite, Fail, Skip",
4640 mode
4641 ),
4642 }
4643 .into());
4644 }
4645 }
4646 }
4647
4648 if let Some(ref unset_property) = action.unset_property_action
4649 && let Some(ref key) = unset_property.key
4650 {
4651 if is_reserved(key) {
4652 return Err(NamespaceError::InvalidInput {
4653 message: format!("Property '{}' is reserved and cannot be modified", key),
4654 }
4655 .into());
4656 }
4657 let mode = unset_property
4658 .mode
4659 .as_deref()
4660 .unwrap_or("Skip")
4661 .to_lowercase();
4662 let exists_in_transaction = transaction
4663 .transaction_properties
4664 .as_ref()
4665 .is_some_and(|props| props.contains_key(key));
4666 match mode.as_str() {
4667 "skip" => {
4668 sidecar.properties.remove(key);
4669 if exists_in_transaction {
4670 sidecar.removed_properties.insert(key.clone());
4673 }
4674 }
4675 "fail" => {
4676 if !sidecar.properties.contains_key(key) && !exists_in_transaction {
4677 return Err(NamespaceError::InvalidInput {
4678 message: format!(
4679 "Property '{}' does not exist and mode is 'Fail'",
4680 key
4681 ),
4682 }
4683 .into());
4684 }
4685 sidecar.properties.remove(key);
4686 if exists_in_transaction {
4687 sidecar.removed_properties.insert(key.clone());
4688 }
4689 }
4690 _ => {
4691 return Err(NamespaceError::InvalidInput {
4692 message: format!(
4693 "Invalid unset_property mode '{}'. Valid values are: Skip, Fail",
4694 mode
4695 ),
4696 }
4697 .into());
4698 }
4699 }
4700 }
4701 }
4702
4703 self.save_transaction_alteration(&table_uri, &transaction.uuid, &sidecar)
4707 .await?;
4708
4709 let final_status = sidecar
4712 .status
4713 .clone()
4714 .unwrap_or_else(|| "SUCCEEDED".to_string());
4715 let response = Self::transaction_response(version, &transaction, Some(sidecar));
4716 Ok(AlterTransactionResponse {
4717 status: final_status,
4718 properties: response.properties,
4719 ..Default::default()
4720 })
4721 }
4722
4723 async fn create_table_scalar_index(
4724 &self,
4725 request: CreateTableIndexRequest,
4726 ) -> Result<CreateTableScalarIndexResponse> {
4727 self.record_op("create_table_scalar_index");
4728 let index_type = Self::parse_index_type(&request.index_type)?;
4729 if !index_type.is_scalar() {
4730 return Err(NamespaceError::InvalidInput {
4731 message: format!(
4732 "create_table_scalar_index only supports scalar index types, got {}",
4733 request.index_type
4734 ),
4735 }
4736 .into());
4737 }
4738
4739 let response = self.create_table_index(request).await?;
4740 Ok(CreateTableScalarIndexResponse {
4741 transaction_id: response.transaction_id,
4742 ..Default::default()
4743 })
4744 }
4745
4746 async fn drop_table_index(
4747 &self,
4748 request: DropTableIndexRequest,
4749 ) -> Result<DropTableIndexResponse> {
4750 self.record_op("drop_table_index");
4751 let table_uri = self.resolve_table_location(&request.id).await?;
4752 let index_name = request.index_name.as_deref().ok_or_else(|| {
4753 lance_core::Error::from(NamespaceError::InvalidInput {
4754 message: "Index name is required for drop_table_index".to_string(),
4755 })
4756 })?;
4757 let mut dataset = self
4758 .load_dataset(&table_uri, None, "drop_table_index")
4759 .await?;
4760 let metadatas = dataset
4761 .load_indices_by_name(index_name)
4762 .await
4763 .map_err(|e| {
4764 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4765 message: format!(
4766 "Failed to load index '{}' before dropping it from table '{}': {}",
4767 index_name, table_uri, e
4768 ),
4769 })
4770 })?;
4771 if metadatas.first().is_some_and(is_system_index) {
4772 return Err(NamespaceError::Unsupported {
4773 message: format!(
4774 "System index '{}' cannot be dropped via this API",
4775 index_name
4776 ),
4777 }
4778 .into());
4779 }
4780
4781 dataset.drop_index(index_name).await.map_err(|e| {
4782 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4783 message: format!(
4784 "Failed to drop index '{}' from table '{}': {}",
4785 index_name, table_uri, e
4786 ),
4787 })
4788 })?;
4789
4790 let transaction_id = dataset
4791 .read_transaction()
4792 .await
4793 .map_err(|e| {
4794 lance_core::Error::from(NamespaceError::Internal {
4795 message: format!(
4796 "Failed to read committed transaction after dropping index '{}' from '{}': {}",
4797 index_name, table_uri, e
4798 ),
4799 })
4800 })?
4801 .map(|transaction| transaction.uuid);
4802
4803 Ok(DropTableIndexResponse {
4804 transaction_id,
4805 ..Default::default()
4806 })
4807 }
4808
4809 async fn list_all_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
4810 let mut tables = self.list_directory_tables().await?;
4813 tables = self
4814 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
4815 .await?;
4816 Self::apply_pagination(&mut tables, request.page_token, request.limit);
4817 Ok(ListTablesResponse::new(tables))
4818 }
4819
4820 async fn restore_table(&self, request: RestoreTableRequest) -> Result<RestoreTableResponse> {
4821 let version = request.version;
4822 if version < 0 {
4823 return Err(Error::invalid_input_source(
4824 format!(
4825 "Table version for restore_table must be non-negative, got {}",
4826 version
4827 )
4828 .into(),
4829 ));
4830 }
4831
4832 let branch = Self::normalized_branch(request.branch.as_deref())?;
4833 let table_uri = self.resolve_table_location(&request.id).await?;
4834 let mut dataset = match branch {
4835 Some(branch) => self.open_validated_branch(&table_uri, branch).await?,
4836 None => self.load_dataset(&table_uri, None, "restore_table").await?,
4837 };
4838
4839 dataset = dataset
4840 .checkout_version(version as u64)
4841 .await
4842 .map_err(|e| {
4843 Error::namespace_source(
4844 format!(
4845 "Failed to checkout version {} for restore at '{}': {}",
4846 version, table_uri, e
4847 )
4848 .into(),
4849 )
4850 })?;
4851
4852 dataset.restore().await.map_err(|e| {
4853 Error::namespace_source(
4854 format!(
4855 "Failed to restore table at '{}' to version {}: {}",
4856 table_uri, version, e
4857 )
4858 .into(),
4859 )
4860 })?;
4861
4862 let transaction_id = dataset
4863 .read_transaction()
4864 .await
4865 .map_err(|e| {
4866 Error::namespace_source(
4867 format!(
4868 "Failed to read transaction after restoring '{}': {}",
4869 table_uri, e
4870 )
4871 .into(),
4872 )
4873 })?
4874 .map(|t| t.uuid);
4875
4876 Ok(RestoreTableResponse {
4877 transaction_id,
4878 ..Default::default()
4879 })
4880 }
4881
4882 async fn update_table_schema_metadata(
4883 &self,
4884 request: UpdateTableSchemaMetadataRequest,
4885 ) -> Result<UpdateTableSchemaMetadataResponse> {
4886 let table_uri = self.resolve_table_location(&request.id).await?;
4887 let mut dataset = self
4888 .load_dataset(&table_uri, None, "update_table_schema_metadata")
4889 .await?;
4890
4891 let new_metadata = request.metadata.unwrap_or_default();
4892 let updated_metadata = dataset
4893 .update_schema_metadata(new_metadata.iter().map(|(k, v)| (k.as_str(), v.as_str())))
4894 .await
4895 .map_err(|e| {
4896 Error::namespace_source(
4897 format!(
4898 "Failed to update schema metadata for table at '{}': {}",
4899 table_uri, e
4900 )
4901 .into(),
4902 )
4903 })?;
4904
4905 let transaction_id = dataset
4906 .read_transaction()
4907 .await
4908 .map_err(|e| {
4909 Error::namespace_source(
4910 format!(
4911 "Failed to read transaction after updating metadata for '{}': {}",
4912 table_uri, e
4913 )
4914 .into(),
4915 )
4916 })?
4917 .map(|t| t.uuid);
4918
4919 Ok(UpdateTableSchemaMetadataResponse {
4920 metadata: Some(updated_metadata),
4921 transaction_id,
4922 ..Default::default()
4923 })
4924 }
4925
4926 async fn get_table_stats(
4927 &self,
4928 request: GetTableStatsRequest,
4929 ) -> Result<GetTableStatsResponse> {
4930 let table_uri = self.resolve_table_location(&request.id).await?;
4931 let dataset = Arc::new(
4932 self.load_dataset(&table_uri, None, "get_table_stats")
4933 .await?,
4934 );
4935
4936 let data_stats = dataset.calculate_data_stats().await.map_err(|e| {
4938 Error::namespace_source(
4939 format!(
4940 "Failed to calculate data statistics for table at '{}': {}",
4941 table_uri, e
4942 )
4943 .into(),
4944 )
4945 })?;
4946 let total_bytes: i64 = data_stats
4947 .fields
4948 .iter()
4949 .map(|f| f.bytes_on_disk as i64)
4950 .sum();
4951
4952 let fragment_row_futures: Vec<_> = dataset
4954 .get_fragments()
4955 .into_iter()
4956 .map(|f| async move { f.physical_rows().await })
4957 .collect();
4958 let fragment_row_results = futures::future::join_all(fragment_row_futures).await;
4959 let mut fragment_row_counts: Vec<i64> = fragment_row_results
4960 .into_iter()
4961 .filter_map(|r| r.ok())
4962 .map(|r| r as i64)
4963 .collect();
4964
4965 let num_fragments = fragment_row_counts.len() as i64;
4966 let num_rows: i64 = fragment_row_counts.iter().sum();
4967
4968 const SMALL_FRAGMENT_THRESHOLD: i64 = 1024 * 1024;
4971 let num_small_fragments = fragment_row_counts
4972 .iter()
4973 .filter(|&&r| r < SMALL_FRAGMENT_THRESHOLD)
4974 .count() as i64;
4975
4976 fragment_row_counts.sort_unstable();
4978 let lengths = if fragment_row_counts.is_empty() {
4979 FragmentSummary::new(0, 0, 0, 0, 0, 0, 0)
4980 } else {
4981 let len = fragment_row_counts.len();
4982 let min = fragment_row_counts[0];
4983 let max = fragment_row_counts[len - 1];
4984 let mean = num_rows / num_fragments;
4985 let pct = |p: f64| fragment_row_counts[((len - 1) as f64 * p) as usize];
4986 FragmentSummary::new(min, max, mean, pct(0.25), pct(0.50), pct(0.75), pct(0.99))
4987 };
4988
4989 let indices = dataset.load_indices().await.map_err(|e| {
4991 Error::namespace_source(
4992 format!("Failed to load indices for table at '{}': {}", table_uri, e).into(),
4993 )
4994 })?;
4995 let num_indices = indices.iter().filter(|m| !is_system_index(m)).count() as i64;
4996
4997 let fragment_stats = FragmentStats::new(num_fragments, num_small_fragments, lengths);
4998 Ok(GetTableStatsResponse::new(
4999 total_bytes,
5000 num_rows,
5001 num_indices,
5002 fragment_stats,
5003 ))
5004 }
5005
5006 async fn explain_table_query_plan(
5007 &self,
5008 request: ExplainTableQueryPlanRequest,
5009 ) -> Result<String> {
5010 let table_uri = self.resolve_table_location(&request.id).await?;
5011 let dataset = self
5012 .load_dataset(
5013 &table_uri,
5014 request.query.version,
5015 "explain_table_query_plan",
5016 )
5017 .await?;
5018 let verbose = request.verbose.unwrap_or(false);
5019
5020 let mut scanner = dataset.scan();
5021 Self::apply_query_params_to_scanner(
5022 &mut scanner,
5023 request.query.filter.as_deref(),
5024 request.query.columns.as_deref(),
5025 request.query.vector_column.as_deref(),
5026 &request.query.vector,
5027 request.query.k,
5028 request.query.offset,
5029 request.query.prefilter,
5030 request.query.bypass_vector_index,
5031 request.query.nprobes,
5032 request.query.ef,
5033 request.query.refine_factor,
5034 request.query.distance_type.as_deref(),
5035 request.query.fast_search,
5036 request.query.with_row_id,
5037 request.query.lower_bound,
5038 request.query.upper_bound,
5039 "explain_table_query_plan",
5040 )?;
5041
5042 scanner.explain_plan(verbose).await.map_err(|e| {
5043 Error::namespace_source(
5044 format!(
5045 "Failed to explain query plan for table at '{}': {}",
5046 table_uri, e
5047 )
5048 .into(),
5049 )
5050 })
5051 }
5052
5053 async fn analyze_table_query_plan(
5054 &self,
5055 request: AnalyzeTableQueryPlanRequest,
5056 ) -> Result<String> {
5057 let table_uri = self.resolve_table_location(&request.id).await?;
5058 let dataset = self
5059 .load_dataset(&table_uri, request.version, "analyze_table_query_plan")
5060 .await?;
5061
5062 let mut scanner = dataset.scan();
5063 Self::apply_query_params_to_scanner(
5064 &mut scanner,
5065 request.filter.as_deref(),
5066 request.columns.as_deref(),
5067 request.vector_column.as_deref(),
5068 &request.vector,
5069 request.k,
5070 request.offset,
5071 request.prefilter,
5072 request.bypass_vector_index,
5073 request.nprobes,
5074 request.ef,
5075 request.refine_factor,
5076 request.distance_type.as_deref(),
5077 request.fast_search,
5078 request.with_row_id,
5079 request.lower_bound,
5080 request.upper_bound,
5081 "analyze_table_query_plan",
5082 )?;
5083
5084 scanner.analyze_plan().await.map_err(|e| {
5085 Error::namespace_source(
5086 format!(
5087 "Failed to analyze query plan for table at '{}': {}",
5088 table_uri, e
5089 )
5090 .into(),
5091 )
5092 })
5093 }
5094
5095 async fn count_table_rows(&self, request: CountTableRowsRequest) -> Result<i64> {
5096 self.record_op("count_table_rows");
5097 let table_uri = self.resolve_table_location(&request.id).await?;
5098 let dataset = self
5099 .load_dataset(&table_uri, request.version, "count_table_rows")
5100 .await?;
5101
5102 let count =
5103 dataset
5104 .count_rows(request.predicate)
5105 .await
5106 .map_err(|e| NamespaceError::Internal {
5107 message: format!("Failed to count rows for table at '{}': {:?}", table_uri, e),
5108 })?;
5109
5110 Ok(count as i64)
5111 }
5112
5113 async fn insert_into_table(
5114 &self,
5115 request: InsertIntoTableRequest,
5116 request_data: Bytes,
5117 ) -> Result<InsertIntoTableResponse> {
5118 self.record_op("insert_into_table");
5119 let table_uri = self.resolve_table_location(&request.id).await?;
5120 let (reader, _num_rows) =
5121 Self::ipc_reader_from_request_data(&request_data, "insert_into_table")?;
5122
5123 let mode = match request.mode.as_deref() {
5124 Some(m) if m.eq_ignore_ascii_case("overwrite") => WriteMode::Overwrite,
5125 Some(m) if m.eq_ignore_ascii_case("append") => WriteMode::Append,
5126 None => WriteMode::Append,
5127 Some(m) => {
5128 return Err(lance_namespace::error::NamespaceError::InvalidInput {
5129 message: format!(
5130 "Unsupported write mode '{}'. Supported modes are: 'append', 'overwrite'",
5131 m
5132 ),
5133 }
5134 .into());
5135 }
5136 };
5137
5138 if !self.table_uri_has_actual_manifests(&table_uri).await? {
5139 self.write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
5140 .await?;
5141 } else {
5142 self.write_reader_to_table(&table_uri, reader, mode, None)
5143 .await?;
5144 }
5145
5146 Ok(InsertIntoTableResponse {
5147 transaction_id: None,
5148 ..Default::default()
5149 })
5150 }
5151
5152 async fn merge_insert_into_table(
5153 &self,
5154 request: MergeInsertIntoTableRequest,
5155 request_data: Bytes,
5156 ) -> Result<MergeInsertIntoTableResponse> {
5157 self.record_op("merge_insert_into_table");
5158 let table_uri = self.resolve_table_location(&request.id).await?;
5159 let on = merge_insert_on_columns(request.on.as_deref(), "merge_insert_into_table")?;
5160
5161 let table_has_manifests = self.table_uri_has_actual_manifests(&table_uri).await?;
5162 let (reader, num_rows) =
5163 Self::ipc_reader_from_request_data(&request_data, "merge_insert_into_table")?;
5164
5165 if !table_has_manifests {
5166 let dataset = self
5167 .write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
5168 .await?;
5169 let version = dataset.version().version as i64;
5170 return Ok(MergeInsertIntoTableResponse {
5171 transaction_id: None,
5172 num_updated_rows: Some(0),
5173 num_inserted_rows: Some(num_rows as i64),
5174 num_deleted_rows: Some(0),
5175 version: Some(version),
5176 ..Default::default()
5177 });
5178 }
5179
5180 let dataset = Arc::new(
5181 self.load_dataset(&table_uri, None, "merge_insert_into_table")
5182 .await?,
5183 );
5184
5185 let mut merge_builder =
5186 MergeInsertBuilder::try_new(dataset.clone(), on.to_vec()).map_err(|e| {
5187 lance_core::Error::from(NamespaceError::InvalidInput {
5188 message: format!("Failed to create merge_insert_into_table builder: {}", e),
5189 })
5190 })?;
5191
5192 if let Some(filter) = request.when_matched_update_all_filt.as_deref() {
5193 let behavior = WhenMatched::update_if(dataset.as_ref(), filter).map_err(|e| {
5194 lance_core::Error::from(NamespaceError::InvalidInput {
5195 message: format!(
5196 "Invalid when_matched_update_all_filt for merge_insert_into_table: {}",
5197 e
5198 ),
5199 })
5200 })?;
5201 merge_builder.when_matched(behavior);
5202 } else if request.when_matched_update_all.unwrap_or(false) {
5203 merge_builder.when_matched(WhenMatched::UpdateAll);
5204 }
5205
5206 if matches!(request.when_not_matched_insert_all, Some(false)) {
5207 merge_builder.when_not_matched(WhenNotMatched::DoNothing);
5208 } else {
5209 merge_builder.when_not_matched(WhenNotMatched::InsertAll);
5210 }
5211
5212 if let Some(filter) = request.when_not_matched_by_source_delete_filt.as_deref() {
5213 let behavior = WhenNotMatchedBySource::delete_if(dataset.as_ref(), filter).map_err(|e| {
5214 lance_core::Error::from(NamespaceError::InvalidInput {
5215 message: format!(
5216 "Invalid when_not_matched_by_source_delete_filt for merge_insert_into_table: {}",
5217 e
5218 ),
5219 })
5220 })?;
5221 merge_builder.when_not_matched_by_source(behavior);
5222 } else if request.when_not_matched_by_source_delete.unwrap_or(false) {
5223 merge_builder.when_not_matched_by_source(WhenNotMatchedBySource::Delete);
5224 }
5225
5226 if let Some(use_index) = request.use_index {
5227 merge_builder.use_index(use_index);
5228 }
5229
5230 let (dataset, stats) = merge_builder
5231 .try_build()
5232 .map_err(|e| {
5233 lance_core::Error::from(NamespaceError::InvalidInput {
5234 message: format!("Failed to build merge_insert_into_table job: {}", e),
5235 })
5236 })?
5237 .execute_reader(reader)
5238 .await
5239 .map_err(|e| Self::map_mutation_error(e, "merge_insert_into_table", &table_uri))?;
5240
5241 Ok(MergeInsertIntoTableResponse {
5242 transaction_id: None,
5243 num_updated_rows: Some(stats.num_updated_rows as i64),
5244 num_inserted_rows: Some(stats.num_inserted_rows as i64),
5245 num_deleted_rows: Some(stats.num_deleted_rows as i64),
5246 version: Some(dataset.version().version as i64),
5247 ..Default::default()
5248 })
5249 }
5250
5251 async fn update_table(&self, request: UpdateTableRequest) -> Result<UpdateTableResponse> {
5252 self.record_op("update_table");
5253
5254 if request.updates.is_empty() {
5255 return Err(NamespaceError::InvalidInput {
5256 message: "update_table requires at least one [column, expression] pair".to_string(),
5257 }
5258 .into());
5259 }
5260
5261 let mut seen_columns: HashMap<String, ()> = HashMap::with_capacity(request.updates.len());
5264 for (idx, pair) in request.updates.iter().enumerate() {
5265 if pair.len() != 2 {
5266 return Err(NamespaceError::InvalidInput {
5267 message: format!(
5268 "update_table updates[{}] must be a [column, expression] pair, got {} elements",
5269 idx,
5270 pair.len()
5271 ),
5272 }
5273 .into());
5274 }
5275 let column = &pair[0];
5276 if column.trim().is_empty() {
5277 return Err(NamespaceError::InvalidInput {
5278 message: format!("update_table updates[{}] has an empty column name", idx),
5279 }
5280 .into());
5281 }
5282 if seen_columns.insert(column.clone(), ()).is_some() {
5283 return Err(NamespaceError::InvalidInput {
5284 message: format!(
5285 "update_table cannot update column '{}' more than once",
5286 column
5287 ),
5288 }
5289 .into());
5290 }
5291 }
5292
5293 let table_uri = self.resolve_table_location(&request.id).await?;
5294 let dataset = Arc::new(self.load_dataset(&table_uri, None, "update_table").await?);
5295
5296 let mut builder = UpdateBuilder::new(dataset);
5297 for pair in &request.updates {
5298 builder = builder.set(&pair[0], &pair[1]).map_err(|e| {
5300 lance_core::Error::from(NamespaceError::InvalidInput {
5301 message: format!("Invalid update expression for column '{}': {}", pair[0], e),
5302 })
5303 })?;
5304 }
5305 if let Some(predicate) = request.predicate.as_deref()
5306 && !predicate.trim().is_empty()
5307 {
5308 builder = builder.update_where(predicate).map_err(|e| {
5309 lance_core::Error::from(NamespaceError::InvalidInput {
5310 message: format!("Invalid update_table predicate '{}': {}", predicate, e),
5311 })
5312 })?;
5313 }
5314
5315 let job = builder.build().map_err(|e| {
5316 lance_core::Error::from(NamespaceError::InvalidInput {
5317 message: format!("Failed to build update_table job: {}", e),
5318 })
5319 })?;
5320
5321 let result = job
5322 .execute()
5323 .await
5324 .map_err(|e| Self::map_mutation_error(e, "update_table", &table_uri))?;
5325
5326 let version = result.new_dataset.version().version as i64;
5327 Ok(UpdateTableResponse {
5328 transaction_id: None,
5329 updated_rows: result.rows_updated as i64,
5330 version,
5331 properties: None,
5332 ..Default::default()
5333 })
5334 }
5335
5336 async fn delete_from_table(
5337 &self,
5338 request: DeleteFromTableRequest,
5339 ) -> Result<DeleteFromTableResponse> {
5340 self.record_op("delete_from_table");
5341
5342 if request.predicate.trim().is_empty() {
5343 return Err(NamespaceError::InvalidInput {
5344 message: "delete_from_table requires a non-empty predicate".to_string(),
5345 }
5346 .into());
5347 }
5348
5349 let table_uri = self.resolve_table_location(&request.id).await?;
5350 let mut dataset = self
5351 .load_dataset(&table_uri, None, "delete_from_table")
5352 .await?;
5353
5354 let result = dataset
5355 .delete(&request.predicate)
5356 .await
5357 .map_err(|e| Self::map_mutation_error(e, "delete_from_table", &table_uri))?;
5358
5359 Ok(DeleteFromTableResponse {
5360 transaction_id: None,
5361 version: Some(result.new_dataset.version().version as i64),
5362 ..Default::default()
5363 })
5364 }
5365
5366 async fn query_table(&self, request: QueryTableRequest) -> Result<Bytes> {
5367 use arrow::ipc::writer::FileWriter;
5368
5369 self.record_op("query_table");
5370 let table_uri = self.resolve_table_location(&request.id).await?;
5371 let dataset = self
5372 .load_dataset(&table_uri, request.version, "query_table")
5373 .await?;
5374
5375 let mut scanner = dataset.scan();
5377
5378 let has_vector_query = request
5381 .vector
5382 .single_vector
5383 .as_ref()
5384 .map(|sv| !sv.is_empty())
5385 .unwrap_or(false)
5386 || request
5387 .vector
5388 .multi_vector
5389 .as_ref()
5390 .map(|mv| !mv.is_empty())
5391 .unwrap_or(false);
5392
5393 if let Some(prefilter) = request.prefilter {
5395 scanner.prefilter(prefilter);
5396 }
5397
5398 if has_vector_query {
5400 let vector_column = request.vector_column.as_deref().unwrap_or("vector");
5401
5402 let query_vector: Vec<f32> = request
5404 .vector
5405 .single_vector
5406 .clone()
5407 .or_else(|| {
5408 request
5409 .vector
5410 .multi_vector
5411 .as_ref()
5412 .and_then(|mv| mv.first().cloned())
5413 })
5414 .unwrap_or_default();
5415
5416 if !query_vector.is_empty() {
5417 let k = if request.k > 0 {
5418 request.k as usize
5419 } else {
5420 10
5421 };
5422 let query_array = Float32Array::from(query_vector);
5423 scanner
5424 .nearest(vector_column, &query_array, k)
5425 .map_err(|e| NamespaceError::InvalidInput {
5426 message: format!("Invalid vector search: {:?}", e),
5427 })?;
5428
5429 if let Some(ref distance_type) = request.distance_type {
5431 let metric = match distance_type.to_lowercase().as_str() {
5432 "l2" | "euclidean" => MetricType::L2,
5433 "cosine" => MetricType::Cosine,
5434 "dot" | "inner_product" => MetricType::Dot,
5435 "hamming" => MetricType::Hamming,
5436 _ => {
5437 return Err(NamespaceError::InvalidInput {
5438 message: format!("Unknown distance type: {}", distance_type),
5439 }
5440 .into());
5441 }
5442 };
5443 scanner.distance_metric(metric);
5444 }
5445
5446 if let Some(nprobes) = request.nprobes {
5448 scanner.minimum_nprobes(nprobes as usize);
5449 }
5450
5451 if let Some(ef) = request.ef {
5453 scanner.ef(ef as usize);
5454 }
5455
5456 if let Some(refine_factor) = request.refine_factor {
5458 scanner.refine(refine_factor as u32);
5459 }
5460
5461 if request.lower_bound.is_some() || request.upper_bound.is_some() {
5463 scanner.distance_range(request.lower_bound, request.upper_bound);
5464 }
5465
5466 if let Some(bypass) = request.bypass_vector_index {
5468 scanner.use_index(!bypass);
5469 }
5470
5471 if request.fast_search == Some(true) {
5473 scanner.fast_search();
5474 }
5475 }
5476 }
5477
5478 if let Some(ref fts_query) = request.full_text_query {
5480 if let Some(ref string_query) = fts_query.string_query {
5482 let mut fts = FullTextSearchQuery::new(string_query.query.clone());
5483
5484 if let Some(ref columns) = string_query.columns
5486 && !columns.is_empty()
5487 {
5488 fts = fts
5489 .with_columns(columns)
5490 .map_err(|e| NamespaceError::InvalidInput {
5491 message: format!("Invalid FTS columns: {:?}", e),
5492 })?;
5493 }
5494
5495 scanner
5496 .full_text_search(fts)
5497 .map_err(|e| NamespaceError::InvalidInput {
5498 message: format!("Invalid full text search: {:?}", e),
5499 })?;
5500 } else if let Some(ref structured_query) = fts_query.structured_query {
5501 let engine_query = build_engine_fts_query(&structured_query.query)?;
5503 let fts = FullTextSearchQuery::new_query(engine_query);
5504 scanner
5505 .full_text_search(fts)
5506 .map_err(|e| NamespaceError::InvalidInput {
5507 message: format!("Invalid full text search: {:?}", e),
5508 })?;
5509 }
5510 }
5511
5512 if let Some(ref columns) = request.columns {
5514 if let Some(ref column_names) = columns.column_names
5515 && !column_names.is_empty()
5516 {
5517 scanner
5518 .project(column_names)
5519 .map_err(|e| NamespaceError::InvalidInput {
5520 message: format!("Invalid column projection: {:?}", e),
5521 })?;
5522 } else if let Some(ref column_aliases) = columns.column_aliases
5523 && !column_aliases.is_empty()
5524 {
5525 let transform_pairs: Vec<(String, String)> = column_aliases
5527 .iter()
5528 .map(|(alias, sql)| (alias.clone(), sql.clone()))
5529 .collect();
5530 scanner
5531 .project_with_transform(
5532 &transform_pairs
5533 .iter()
5534 .map(|(a, s)| (a.as_str(), s.as_str()))
5535 .collect::<Vec<_>>(),
5536 )
5537 .map_err(|e| NamespaceError::InvalidInput {
5538 message: format!("Invalid column alias expression: {:?}", e),
5539 })?;
5540 }
5541 }
5542
5543 if let Some(ref filter) = request.filter
5545 && !filter.is_empty()
5546 {
5547 scanner
5548 .filter(filter)
5549 .map_err(|e| NamespaceError::InvalidInput {
5550 message: format!("Invalid filter expression: {:?}", e),
5551 })?;
5552 }
5553
5554 if request.with_row_id == Some(true) {
5556 scanner.with_row_id();
5557 }
5558
5559 if !has_vector_query && request.k > 0 {
5563 let offset = request.offset.map(|o| o as i64);
5564 scanner.limit(Some(request.k as i64), offset).map_err(|e| {
5565 NamespaceError::InvalidInput {
5566 message: format!("Invalid limit/offset: {:?}", e),
5567 }
5568 })?;
5569 } else if has_vector_query && request.offset.is_some() {
5570 let offset = request.offset.map(|o| o as i64);
5572 scanner
5573 .limit(None, offset)
5574 .map_err(|e| NamespaceError::InvalidInput {
5575 message: format!("Invalid offset: {:?}", e),
5576 })?;
5577 }
5578
5579 let batch = scanner
5581 .try_into_batch()
5582 .await
5583 .map_err(|e| NamespaceError::Internal {
5584 message: format!("Failed to execute query: {:?}", e),
5585 })?;
5586
5587 let schema = batch.schema();
5589 let mut buffer = Vec::new();
5590 {
5591 let mut writer = FileWriter::try_new(&mut buffer, &schema).map_err(|e| {
5592 NamespaceError::Internal {
5593 message: format!("Failed to create IPC writer: {:?}", e),
5594 }
5595 })?;
5596 writer.write(&batch).map_err(|e| NamespaceError::Internal {
5597 message: format!("Failed to write batch to IPC: {:?}", e),
5598 })?;
5599 writer.finish().map_err(|e| NamespaceError::Internal {
5600 message: format!("Failed to finish IPC writer: {:?}", e),
5601 })?;
5602 }
5603
5604 Ok(Bytes::from(buffer))
5605 }
5606
5607 async fn list_table_tags(
5608 &self,
5609 request: ListTableTagsRequest,
5610 ) -> Result<ListTableTagsResponse> {
5611 self.record_op("list_table_tags");
5612 let table_uri = self.resolve_table_location(&request.id).await?;
5613 let dataset = self
5614 .load_dataset(&table_uri, None, "list_table_tags")
5615 .await?;
5616
5617 let raw_tags = dataset.tags().list().await.map_err(|e| {
5618 lance_core::Error::from(NamespaceError::Internal {
5619 message: format!("Failed to list tags for table at '{}': {}", table_uri, e),
5620 })
5621 })?;
5622
5623 let tags = raw_tags
5624 .into_iter()
5625 .map(|(name, contents)| {
5626 let mut tag_model =
5627 ModelTagContents::new(contents.version as i64, contents.manifest_size as i64);
5628 tag_model.branch = contents.branch;
5629 (name, tag_model)
5630 })
5631 .collect();
5632
5633 Ok(ListTableTagsResponse {
5634 tags,
5635 page_token: None,
5636 ..Default::default()
5637 })
5638 }
5639
5640 async fn get_table_tag_version(
5641 &self,
5642 request: GetTableTagVersionRequest,
5643 ) -> Result<GetTableTagVersionResponse> {
5644 self.record_op("get_table_tag_version");
5645 if request.tag.is_empty() {
5646 return Err(NamespaceError::InvalidInput {
5647 message: "tag name must not be empty for get_table_tag_version".to_string(),
5648 }
5649 .into());
5650 }
5651
5652 let table_uri = self.resolve_table_location(&request.id).await?;
5653 let dataset = self
5654 .load_dataset(&table_uri, None, "get_table_tag_version")
5655 .await?;
5656
5657 let contents = dataset
5658 .tags()
5659 .get(&request.tag)
5660 .await
5661 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5662
5663 Ok(GetTableTagVersionResponse {
5664 version: contents.version as i64,
5665 branch: contents.branch,
5666 ..Default::default()
5667 })
5668 }
5669
5670 async fn create_table_tag(
5671 &self,
5672 request: CreateTableTagRequest,
5673 ) -> Result<CreateTableTagResponse> {
5674 self.record_op("create_table_tag");
5675 if request.tag.is_empty() {
5676 return Err(NamespaceError::InvalidInput {
5677 message: "tag name must not be empty for create_table_tag".to_string(),
5678 }
5679 .into());
5680 }
5681 if request.version <= 0 {
5682 return Err(NamespaceError::InvalidInput {
5683 message: format!(
5684 "tag version must be a positive integer, got {} for create_table_tag",
5685 request.version
5686 ),
5687 }
5688 .into());
5689 }
5690
5691 let table_uri = self.resolve_table_location(&request.id).await?;
5692 let dataset = self
5693 .load_dataset(&table_uri, None, "create_table_tag")
5694 .await?;
5695
5696 dataset
5697 .tags()
5698 .create(&request.tag, request.version as u64)
5699 .await
5700 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5701
5702 Ok(CreateTableTagResponse {
5703 transaction_id: None,
5704 ..Default::default()
5705 })
5706 }
5707
5708 async fn delete_table_tag(
5709 &self,
5710 request: DeleteTableTagRequest,
5711 ) -> Result<DeleteTableTagResponse> {
5712 self.record_op("delete_table_tag");
5713 if request.tag.is_empty() {
5714 return Err(NamespaceError::InvalidInput {
5715 message: "tag name must not be empty for delete_table_tag".to_string(),
5716 }
5717 .into());
5718 }
5719
5720 let table_uri = self.resolve_table_location(&request.id).await?;
5721 let dataset = self
5722 .load_dataset(&table_uri, None, "delete_table_tag")
5723 .await?;
5724
5725 dataset
5726 .tags()
5727 .delete(&request.tag)
5728 .await
5729 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5730
5731 Ok(DeleteTableTagResponse {
5732 transaction_id: None,
5733 ..Default::default()
5734 })
5735 }
5736
5737 async fn update_table_tag(
5738 &self,
5739 request: UpdateTableTagRequest,
5740 ) -> Result<UpdateTableTagResponse> {
5741 self.record_op("update_table_tag");
5742 if request.tag.is_empty() {
5743 return Err(NamespaceError::InvalidInput {
5744 message: "tag name must not be empty for update_table_tag".to_string(),
5745 }
5746 .into());
5747 }
5748 if request.version <= 0 {
5749 return Err(NamespaceError::InvalidInput {
5750 message: format!(
5751 "tag version must be a positive integer, got {} for update_table_tag",
5752 request.version
5753 ),
5754 }
5755 .into());
5756 }
5757
5758 let table_uri = self.resolve_table_location(&request.id).await?;
5759 let dataset = self
5760 .load_dataset(&table_uri, None, "update_table_tag")
5761 .await?;
5762
5763 dataset
5764 .tags()
5765 .update(&request.tag, request.version as u64)
5766 .await
5767 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5768
5769 Ok(UpdateTableTagResponse {
5770 transaction_id: None,
5771 ..Default::default()
5772 })
5773 }
5774
5775 async fn create_table_branch(
5776 &self,
5777 request: CreateTableBranchRequest,
5778 ) -> Result<CreateTableBranchResponse> {
5779 self.record_op("create_table_branch");
5780 if request.name.is_empty() {
5781 return Err(NamespaceError::InvalidInput {
5782 message: "branch name must not be empty for create_table_branch".to_string(),
5783 }
5784 .into());
5785 }
5786 let from_version = match request.from_version {
5787 Some(v) if v <= 0 => {
5788 return Err(NamespaceError::InvalidInput {
5789 message: format!(
5790 "from_version must be a positive integer, got {} for create_table_branch",
5791 v
5792 ),
5793 }
5794 .into());
5795 }
5796 Some(v) => Some(v as u64),
5797 None => None,
5798 };
5799
5800 let table_uri = self.resolve_table_location(&request.id).await?;
5801 let mut dataset = self
5802 .load_dataset(&table_uri, None, "create_table_branch")
5803 .await?;
5804
5805 if dataset.branches().get(&request.name).await.is_ok() {
5810 return Err(NamespaceError::TableBranchAlreadyExists {
5811 message: format!("branch '{}' for table at '{}'", request.name, table_uri),
5812 }
5813 .into());
5814 }
5815
5816 dataset
5817 .create_branch(
5818 &request.name,
5819 (request.from_branch.as_deref(), from_version),
5820 None,
5821 )
5822 .await
5823 .map_err(|e| {
5824 if matches!(e, lance_core::Error::DatasetNotFound { .. }) {
5827 NamespaceError::InvalidInput {
5828 message: format!(
5829 "from_branch/from_version for branch '{}' refers to a source that does not exist: {}",
5830 request.name, e
5831 ),
5832 }
5833 .into()
5834 } else {
5835 Self::map_branch_error(e, &request.name, &table_uri)
5836 }
5837 })?;
5838
5839 Ok(CreateTableBranchResponse {
5840 transaction_id: None,
5841 ..Default::default()
5842 })
5843 }
5844
5845 async fn list_table_branches(
5846 &self,
5847 request: ListTableBranchesRequest,
5848 ) -> Result<ListTableBranchesResponse> {
5849 self.record_op("list_table_branches");
5850 let table_uri = self.resolve_table_location(&request.id).await?;
5851 let dataset = self
5852 .load_dataset(&table_uri, None, "list_table_branches")
5853 .await?;
5854
5855 let raw_branches = dataset.list_branches().await.map_err(|e| {
5856 lance_core::Error::from(NamespaceError::Internal {
5857 message: format!(
5858 "Failed to list branches for table at '{}': {}",
5859 table_uri, e
5860 ),
5861 })
5862 })?;
5863
5864 let branches = raw_branches
5865 .into_iter()
5866 .map(|(name, contents)| {
5867 let mut branch_model = ModelBranchContents::new(
5870 contents.parent_version as i64,
5871 contents.create_at as i64,
5872 contents.manifest_size as i64,
5873 );
5874 branch_model.parent_branch = contents.parent_branch;
5875 branch_model.metadata = if contents.metadata.is_empty() {
5876 None
5877 } else {
5878 Some(contents.metadata)
5879 };
5880 (name, branch_model)
5881 })
5882 .collect();
5883
5884 Ok(ListTableBranchesResponse {
5885 branches,
5886 page_token: None,
5887 ..Default::default()
5888 })
5889 }
5890
5891 async fn delete_table_branch(
5892 &self,
5893 request: DeleteTableBranchRequest,
5894 ) -> Result<DeleteTableBranchResponse> {
5895 self.record_op("delete_table_branch");
5896 if request.name.is_empty() {
5897 return Err(NamespaceError::InvalidInput {
5898 message: "branch name must not be empty for delete_table_branch".to_string(),
5899 }
5900 .into());
5901 }
5902
5903 let table_uri = self.resolve_table_location(&request.id).await?;
5904 let mut dataset = self
5905 .load_dataset(&table_uri, None, "delete_table_branch")
5906 .await?;
5907
5908 dataset
5909 .delete_branch(&request.name)
5910 .await
5911 .map_err(|e| match e {
5912 lance_core::Error::RefConflict { message } => NamespaceError::InvalidInput {
5913 message: format!(
5914 "branch '{}' for table at '{}': {}",
5915 request.name, table_uri, message
5916 ),
5917 }
5918 .into(),
5919 other => Self::map_branch_error(other, &request.name, &table_uri),
5920 })?;
5921
5922 Ok(DeleteTableBranchResponse {
5923 transaction_id: None,
5924 ..Default::default()
5925 })
5926 }
5927
5928 fn namespace_id(&self) -> String {
5929 format!("DirectoryNamespace {{ root: {:?} }}", self.root)
5930 }
5931}
5932
5933#[derive(Debug)]
5935pub(crate) enum MarkerFileError {
5936 AlreadyExists { description: String },
5938 Other { message: String },
5940}
5941
5942impl std::fmt::Display for MarkerFileError {
5943 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5944 match self {
5945 Self::AlreadyExists { description } => {
5946 write!(f, "{} already exists", description)
5947 }
5948 Self::Other { message } => write!(f, "{}", message),
5949 }
5950 }
5951}
5952
5953pub(crate) async fn put_marker_file_atomic(
5974 object_store: &ObjectStore,
5975 path: &Path,
5976 file_description: &str,
5977) -> std::result::Result<(), MarkerFileError> {
5978 let staging_name = format!("lance-marker.staging.{}", uuid::Uuid::new_v4().simple());
5979 let path_str = path.as_ref();
5980 let staging_path = match path_str.rfind('/') {
5981 Some(idx) => Path::from(format!("{}/{}", &path_str[..idx], staging_name)),
5982 None => Path::from(staging_name.as_str()),
5983 };
5984
5985 object_store
5986 .inner
5987 .put(&staging_path, bytes::Bytes::from_static(b"reserved").into())
5988 .await
5989 .map_err(|e| MarkerFileError::Other {
5990 message: format!("Failed to stage {}: {:?}", file_description, e),
5991 })?;
5992
5993 let mut staging_consumed = false;
5996 let publish_result = match object_store
5997 .inner
5998 .rename_if_not_exists(&staging_path, path)
5999 .await
6000 {
6001 Ok(()) => {
6002 staging_consumed = true;
6003 Ok(())
6004 }
6005 Err(ObjectStoreError::NotImplemented { .. })
6006 | Err(ObjectStoreError::NotSupported { .. }) => {
6007 match object_store
6008 .inner
6009 .copy_if_not_exists(&staging_path, path)
6010 .await
6011 {
6012 Ok(()) => Ok(()),
6013 Err(ObjectStoreError::NotImplemented { .. })
6014 | Err(ObjectStoreError::NotSupported { .. }) => object_store
6015 .inner
6016 .put_opts(
6017 path,
6018 bytes::Bytes::from_static(b"reserved").into(),
6019 PutOptions {
6020 mode: PutMode::Create,
6021 ..Default::default()
6022 },
6023 )
6024 .await
6025 .map(|_| ()),
6026 Err(e) => Err(e),
6027 }
6028 }
6029 Err(e) => Err(e),
6030 };
6031
6032 if !staging_consumed {
6033 delete_staging_marker_best_effort(object_store, &staging_path).await;
6034 }
6035
6036 match publish_result {
6037 Ok(()) => Ok(()),
6038 Err(ObjectStoreError::AlreadyExists { .. })
6039 | Err(ObjectStoreError::Precondition { .. }) => Err(MarkerFileError::AlreadyExists {
6040 description: file_description.to_string(),
6041 }),
6042 Err(e) => Err(MarkerFileError::Other {
6043 message: format!("Failed to create {}: {:?}", file_description, e),
6044 }),
6045 }
6046}
6047
6048async fn delete_staging_marker_best_effort(object_store: &ObjectStore, staging_path: &Path) {
6052 const MAX_ATTEMPTS: u32 = 3;
6053 const BACKOFF_MS: [u64; 2] = [20, 50];
6054
6055 let mut last_err: Option<ObjectStoreError> = None;
6056 for attempt in 0..MAX_ATTEMPTS {
6057 match object_store.inner.delete(staging_path).await {
6058 Ok(()) => return,
6059 Err(ObjectStoreError::NotFound { .. }) => return,
6060 Err(e) => {
6061 last_err = Some(e);
6062 if let Some(&delay_ms) = BACKOFF_MS.get(attempt as usize) {
6063 tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
6064 }
6065 }
6066 }
6067 }
6068 if let Some(del_err) = last_err {
6069 log::warn!(
6070 "Failed to delete staging marker at '{}' after {} attempts: {:?}",
6071 staging_path,
6072 MAX_ATTEMPTS,
6073 del_err
6074 );
6075 }
6076}
6077
6078fn build_engine_fts_query(
6082 query: &lance_namespace::models::FtsQuery,
6083) -> std::result::Result<FtsQuery, NamespaceError> {
6084 if let Some(ref m) = query.r#match {
6085 Ok(FtsQuery::Match(build_engine_match_query(m)?))
6086 } else if let Some(ref p) = query.phrase {
6087 let mut phrase = PhraseQuery::new(p.terms.clone());
6088 if let Some(ref column) = p.column {
6089 phrase = phrase.with_column(Some(column.clone()));
6090 }
6091 if let Some(slop) = p.slop {
6092 phrase = phrase.with_slop(slop as u32);
6093 }
6094 Ok(FtsQuery::Phrase(phrase))
6095 } else if let Some(ref mm) = query.multi_match {
6096 let match_queries = mm
6097 .match_queries
6098 .iter()
6099 .map(build_engine_match_query)
6100 .collect::<std::result::Result<Vec<_>, _>>()?;
6101 Ok(FtsQuery::MultiMatch(MultiMatchQuery { match_queries }))
6102 } else if let Some(ref b) = query.boolean {
6103 let mut clauses: Vec<(Occur, FtsQuery)> = Vec::new();
6104 for clause in &b.must {
6105 clauses.push((Occur::Must, build_engine_fts_query(clause)?));
6106 }
6107 for clause in &b.should {
6108 clauses.push((Occur::Should, build_engine_fts_query(clause)?));
6109 }
6110 for clause in &b.must_not {
6111 clauses.push((Occur::MustNot, build_engine_fts_query(clause)?));
6112 }
6113 Ok(FtsQuery::Boolean(BooleanQuery::new(clauses)))
6114 } else if let Some(ref boost) = query.boost {
6115 let positive = build_engine_fts_query(&boost.positive)?;
6116 let negative = build_engine_fts_query(&boost.negative)?;
6117 Ok(FtsQuery::Boost(BoostQuery::new(
6118 positive,
6119 negative,
6120 boost.negative_boost,
6121 )))
6122 } else {
6123 Err(NamespaceError::InvalidInput {
6124 message: "structured_query.query must set exactly one of match, phrase, multi_match, \
6125 boolean, or boost"
6126 .to_string(),
6127 })
6128 }
6129}
6130
6131fn build_engine_match_query(
6132 m: &lance_namespace::models::MatchQuery,
6133) -> std::result::Result<MatchQuery, NamespaceError> {
6134 let mut match_query = MatchQuery::new(m.terms.clone());
6135 if let Some(ref column) = m.column {
6136 match_query = match_query.with_column(Some(column.clone()));
6137 }
6138 if let Some(boost) = m.boost {
6139 match_query = match_query.with_boost(boost);
6140 }
6141 if let Some(fuzziness) = m.fuzziness {
6142 match_query = match_query.with_fuzziness(Some(fuzziness as u32));
6143 }
6144 if let Some(max_expansions) = m.max_expansions {
6145 match_query = match_query.with_max_expansions(max_expansions as usize);
6146 }
6147 if let Some(ref operator) = m.operator {
6148 let op =
6149 Operator::try_from(operator.as_str()).map_err(|e| NamespaceError::InvalidInput {
6150 message: format!("Invalid FTS operator: {:?}", e),
6151 })?;
6152 match_query = match_query.with_operator(op);
6153 }
6154 if let Some(prefix_length) = m.prefix_length {
6155 match_query = match_query.with_prefix_length(prefix_length as u32);
6156 }
6157 Ok(match_query)
6158}
6159
6160#[cfg(test)]
6161mod tests {
6162 use super::*;
6163 use arrow_ipc::reader::{FileReader, StreamReader};
6164 use lance::index::vector::StageParams;
6165 use rstest::rstest;
6166
6167 fn build_ivf_rq_num_bits(num_bits: Option<i32>) -> Result<u8> {
6168 let mut request = CreateTableIndexRequest::new("vector".to_string(), "IVF_RQ".to_string());
6169 request.num_bits = num_bits;
6170
6171 let DirectoryIndexParams::Vector {
6172 index_type: IndexType::IvfRq,
6173 params,
6174 } = DirectoryNamespace::build_index_params(&request)?
6175 else {
6176 panic!("expected IVF_RQ vector index params");
6177 };
6178 match params.stages.as_slice() {
6179 [StageParams::Ivf(_), StageParams::RQ(rq)] => Ok(rq.num_bits),
6180 stages => panic!("expected IVF and RQ stages, got {stages:?}"),
6181 }
6182 }
6183
6184 #[rstest]
6185 #[case::omitted(None, 5)]
6186 #[case::explicit_one(Some(1), 1)]
6187 #[case::explicit_max(Some(9), 9)]
6188 fn test_build_index_params_ivf_rq_num_bits(
6189 #[case] requested: Option<i32>,
6190 #[case] expected: u8,
6191 ) {
6192 assert_eq!(build_ivf_rq_num_bits(requested).unwrap(), expected);
6193 }
6194
6195 #[rstest]
6196 #[case::negative(-1)]
6197 #[case::zero(0)]
6198 #[case::above_max(10)]
6199 #[case::conversion_overflow(i32::MAX)]
6200 fn test_build_index_params_rejects_invalid_ivf_rq_num_bits(#[case] requested: i32) {
6201 let error = build_ivf_rq_num_bits(Some(requested))
6202 .expect_err("invalid IVF_RQ num_bits should fail");
6203 let message = error.to_string();
6204
6205 assert_eq!(mutation_error_code(error), ErrorCode::InvalidInput);
6206 assert!(
6207 message.contains(&format!(
6208 "IVF_RQ num_bits must be in 1..=9, got {requested}"
6209 )),
6210 "unexpected error message: {message}"
6211 );
6212 }
6213
6214 #[test]
6215 fn test_build_engine_fts_query_match() {
6216 let mut ns_match = lance_namespace::models::MatchQuery::new("hello world".to_string());
6217 ns_match.column = Some("body".to_string());
6218 ns_match.operator = Some("AND".to_string());
6219 ns_match.fuzziness = Some(1);
6220 ns_match.max_expansions = Some(30);
6221 ns_match.boost = Some(2.0);
6222 ns_match.prefix_length = Some(2);
6223
6224 let mut ns_query = lance_namespace::models::FtsQuery::new();
6225 ns_query.r#match = Some(Box::new(ns_match));
6226
6227 match build_engine_fts_query(&ns_query).unwrap() {
6228 FtsQuery::Match(m) => {
6229 assert_eq!(m.terms, "hello world");
6230 assert_eq!(m.column, Some("body".to_string()));
6231 assert_eq!(m.operator, Operator::And);
6232 assert_eq!(m.fuzziness, Some(1));
6233 assert_eq!(m.max_expansions, 30);
6234 assert_eq!(m.boost, 2.0);
6235 assert_eq!(m.prefix_length, 2);
6236 }
6237 other => panic!("expected Match, got {:?}", other),
6238 }
6239 }
6240
6241 fn ns_match_query(terms: &str, column: &str) -> lance_namespace::models::FtsQuery {
6244 let mut m = lance_namespace::models::MatchQuery::new(terms.to_string());
6245 m.column = Some(column.to_string());
6246 let mut q = lance_namespace::models::FtsQuery::new();
6247 q.r#match = Some(Box::new(m));
6248 q
6249 }
6250
6251 #[test]
6252 fn test_build_engine_fts_query_phrase() {
6253 let mut ns_phrase = lance_namespace::models::PhraseQuery::new("hello world".to_string());
6254 ns_phrase.column = Some("body".to_string());
6255 ns_phrase.slop = Some(2);
6256
6257 let mut ns_query = lance_namespace::models::FtsQuery::new();
6258 ns_query.phrase = Some(Box::new(ns_phrase));
6259
6260 match build_engine_fts_query(&ns_query).unwrap() {
6261 FtsQuery::Phrase(p) => {
6262 assert_eq!(p.terms, "hello world");
6263 assert_eq!(p.column, Some("body".to_string()));
6264 assert_eq!(p.slop, 2);
6265 }
6266 other => panic!("expected Phrase, got {:?}", other),
6267 }
6268 }
6269
6270 #[test]
6271 fn test_build_engine_fts_query_multi_match() {
6272 let mut m1 = lance_namespace::models::MatchQuery::new("hello".to_string());
6273 m1.column = Some("title".to_string());
6274 let mut m2 = lance_namespace::models::MatchQuery::new("hello".to_string());
6275 m2.column = Some("body".to_string());
6276 m2.boost = Some(2.0);
6277
6278 let ns_multi = lance_namespace::models::MultiMatchQuery::new(vec![m1, m2]);
6279 let mut ns_query = lance_namespace::models::FtsQuery::new();
6280 ns_query.multi_match = Some(Box::new(ns_multi));
6281
6282 match build_engine_fts_query(&ns_query).unwrap() {
6283 FtsQuery::MultiMatch(mm) => {
6284 assert_eq!(mm.match_queries.len(), 2);
6285 assert_eq!(mm.match_queries[0].terms, "hello");
6286 assert_eq!(mm.match_queries[0].column, Some("title".to_string()));
6287 assert_eq!(mm.match_queries[1].column, Some("body".to_string()));
6288 assert_eq!(mm.match_queries[1].boost, 2.0);
6289 }
6290 other => panic!("expected MultiMatch, got {:?}", other),
6291 }
6292 }
6293
6294 #[test]
6295 fn test_build_engine_fts_query_boolean() {
6296 let ns_boolean = lance_namespace::models::BooleanQuery::new(
6298 vec![ns_match_query("must-term", "body")],
6299 vec![ns_match_query("must-not-term", "body")],
6300 vec![ns_match_query("should-term", "body")],
6301 );
6302 let mut ns_query = lance_namespace::models::FtsQuery::new();
6303 ns_query.boolean = Some(Box::new(ns_boolean));
6304
6305 match build_engine_fts_query(&ns_query).unwrap() {
6306 FtsQuery::Boolean(b) => {
6307 assert!(matches!(&b.must[..], [FtsQuery::Match(m)] if m.terms == "must-term"));
6308 assert!(
6309 matches!(&b.must_not[..], [FtsQuery::Match(m)] if m.terms == "must-not-term")
6310 );
6311 assert!(matches!(&b.should[..], [FtsQuery::Match(m)] if m.terms == "should-term"));
6312 }
6313 other => panic!("expected Boolean, got {:?}", other),
6314 }
6315 }
6316
6317 #[test]
6318 fn test_build_engine_fts_query_boost() {
6319 let mut ns_boost = lance_namespace::models::BoostQuery::new(
6320 ns_match_query("positive-term", "body"),
6321 ns_match_query("negative-term", "body"),
6322 );
6323 ns_boost.negative_boost = Some(0.25);
6324
6325 let mut ns_query = lance_namespace::models::FtsQuery::new();
6326 ns_query.boost = Some(Box::new(ns_boost));
6327
6328 match build_engine_fts_query(&ns_query).unwrap() {
6329 FtsQuery::Boost(b) => {
6330 assert!(
6331 matches!(b.positive.as_ref(), FtsQuery::Match(m) if m.terms == "positive-term")
6332 );
6333 assert!(
6334 matches!(b.negative.as_ref(), FtsQuery::Match(m) if m.terms == "negative-term")
6335 );
6336 assert_eq!(b.negative_boost, 0.25);
6337 }
6338 other => panic!("expected Boost, got {:?}", other),
6339 }
6340 }
6341
6342 #[test]
6343 fn test_build_engine_fts_query_requires_a_variant() {
6344 let empty = lance_namespace::models::FtsQuery::new();
6346 assert!(build_engine_fts_query(&empty).is_err());
6347 }
6348 use lance::dataset::Dataset;
6349 use lance::index::DatasetIndexExt;
6350 use lance_core::utils::tempfile::{TempStdDir, TempStrDir};
6351 use lance_core::utils::testing::CountingObjectStore;
6352 use lance_io::object_store::{providers::local::FileStoreProvider, uri_to_url};
6353 use lance_namespace::error::ErrorCode;
6354 use lance_namespace::models::{
6355 CreateTableRequest, JsonArrowDataType, JsonArrowField, JsonArrowSchema, ListTablesRequest,
6356 QueryTableRequestColumns,
6357 };
6358 use lance_namespace::schema::convert_json_arrow_schema;
6359 use std::io::Cursor;
6360 use std::sync::{
6361 Arc,
6362 atomic::{AtomicUsize, Ordering},
6363 };
6364 use url::Url;
6365
6366 fn assert_plan_contains_all(plan: &str, expected_fragments: &[&str], context: &str) {
6367 for expected_fragment in expected_fragments {
6368 assert!(
6369 plan.contains(expected_fragment),
6370 "{}. Missing fragment: '{}'. Plan:\n{}",
6371 context,
6372 expected_fragment,
6373 plan
6374 );
6375 }
6376 }
6377
6378 fn mutation_error_code(err: lance_core::Error) -> ErrorCode {
6379 match err {
6380 lance_core::Error::Namespace { source, .. } => source
6381 .downcast_ref::<NamespaceError>()
6382 .expect("mutation error should wrap a NamespaceError")
6383 .code(),
6384 other => panic!("expected Namespace error, got: {other:?}"),
6385 }
6386 }
6387
6388 #[test]
6393 fn test_map_mutation_error_commit_conflict_alignment() {
6394 let boxed = || -> Box<dyn std::error::Error + Send + Sync + 'static> {
6395 Box::<dyn std::error::Error + Send + Sync>::from("inner conflict")
6396 };
6397
6398 let throttling_cases = vec![lance_core::Error::commit_conflict_source(1, boxed())];
6399 for err in throttling_cases {
6400 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6401 err,
6402 "update",
6403 "memory://t",
6404 ));
6405 assert_eq!(code, ErrorCode::Throttling);
6406 }
6407
6408 let concurrent_cases = vec![
6409 lance_core::Error::too_much_write_contention("contention"),
6410 lance_core::Error::retryable_commit_conflict_source(1, boxed()),
6411 lance_core::Error::incompatible_transaction_source(boxed()),
6412 lance_core::Error::version_conflict("conflict", 0, 3),
6413 ];
6414 for err in concurrent_cases {
6415 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6416 err,
6417 "update",
6418 "memory://t",
6419 ));
6420 assert_eq!(code, ErrorCode::ConcurrentModification);
6421 }
6422 }
6423
6424 async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) {
6426 let temp_dir = TempStdDir::default();
6427
6428 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
6429 .build()
6430 .await
6431 .unwrap();
6432 (namespace, temp_dir)
6433 }
6434
6435 #[tokio::test]
6438 async fn test_list_versions_under_ordering_and_limit() {
6439 use lance_table::io::commit::ManifestNamingScheme;
6440
6441 async fn seed_and_check(ns: &DirectoryNamespace) {
6442 let table_path = ns.base_path.clone().join("lv_test.lance");
6443 for v in 1..=7u64 {
6444 let p = ManifestNamingScheme::V2.manifest_path(&table_path, v);
6445 ns.object_store.put(&p, b"m".as_slice()).await.unwrap();
6446 }
6447 let staging = Path::parse(format!(
6451 "{}-cee4fbbb-eb19-4ea3-8ca7-54f5ec33dedc",
6452 ManifestNamingScheme::V2.manifest_path(&table_path, 8)
6453 ))
6454 .unwrap();
6455 ns.object_store
6456 .put(&staging, b"s".as_slice())
6457 .await
6458 .unwrap();
6459 let detached = table_path.clone().join(VERSIONS_DIR).join("d123.manifest");
6460 ns.object_store
6461 .put(&detached, b"d".as_slice())
6462 .await
6463 .unwrap();
6464 fn versions(r: &[TableVersion]) -> Vec<i64> {
6465 r.iter().map(|t| t.version).collect()
6466 }
6467
6468 let got = ns
6469 .list_versions_under(&table_path, true, Some(1))
6470 .await
6471 .unwrap();
6472 assert_eq!(versions(&got), vec![7]);
6473 let got = ns
6474 .list_versions_under(&table_path, true, Some(3))
6475 .await
6476 .unwrap();
6477 assert_eq!(versions(&got), vec![7, 6, 5]);
6478
6479 let got = ns
6480 .list_versions_under(&table_path, false, Some(2))
6481 .await
6482 .unwrap();
6483 assert_eq!(versions(&got), vec![1, 2]);
6484
6485 let got = ns
6486 .list_versions_under(&table_path, true, None)
6487 .await
6488 .unwrap();
6489 assert_eq!(versions(&got), vec![7, 6, 5, 4, 3, 2, 1]);
6490 let got = ns
6491 .list_versions_under(&table_path, false, None)
6492 .await
6493 .unwrap();
6494 assert_eq!(versions(&got), vec![1, 2, 3, 4, 5, 6, 7]);
6495
6496 let got = ns
6497 .list_versions_under(&table_path, true, Some(0))
6498 .await
6499 .unwrap();
6500 assert!(got.is_empty());
6501 let got = ns
6502 .list_versions_under(&table_path, true, Some(100))
6503 .await
6504 .unwrap();
6505 assert_eq!(versions(&got), vec![7, 6, 5, 4, 3, 2, 1]);
6506
6507 let got = ns
6509 .list_versions_under(&table_path, true, Some(-1))
6510 .await
6511 .unwrap();
6512 assert_eq!(versions(&got), vec![7, 6, 5, 4, 3, 2, 1]);
6513 }
6514
6515 let ns_mem = DirectoryNamespaceBuilder::new("memory://lv-test")
6516 .build()
6517 .await
6518 .unwrap();
6519 assert!(ns_mem.object_store.list_is_lexically_ordered);
6520 seed_and_check(&ns_mem).await;
6521
6522 let (ns_fs, _tmp) = create_test_namespace().await;
6523 assert!(!ns_fs.object_store.list_is_lexically_ordered);
6524 seed_and_check(&ns_fs).await;
6525 }
6526
6527 #[tokio::test]
6531 async fn test_list_versions_under_early_stop_bounded_consumption() {
6532 use lance_io::object_store::providers::memory::MemoryStoreProvider;
6533 use lance_table::io::commit::ManifestNamingScheme;
6534
6535 #[derive(Debug)]
6536 struct EntryCountingStore {
6537 target: Arc<dyn OSObjectStore>,
6538 entries_listed: Arc<AtomicUsize>,
6539 }
6540
6541 impl std::fmt::Display for EntryCountingStore {
6542 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6543 write!(f, "EntryCountingStore({})", self.target)
6544 }
6545 }
6546
6547 #[async_trait]
6548 impl OSObjectStore for EntryCountingStore {
6549 async fn put_opts(
6550 &self,
6551 location: &Path,
6552 bytes: PutPayload,
6553 opts: PutOptions,
6554 ) -> OSResult<PutResult> {
6555 self.target.put_opts(location, bytes, opts).await
6556 }
6557
6558 async fn put_multipart_opts(
6559 &self,
6560 location: &Path,
6561 opts: PutMultipartOptions,
6562 ) -> OSResult<Box<dyn MultipartUpload>> {
6563 self.target.put_multipart_opts(location, opts).await
6564 }
6565
6566 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
6567 self.target.get_opts(location, options).await
6568 }
6569
6570 fn delete_stream(
6571 &self,
6572 locations: BoxStream<'static, OSResult<Path>>,
6573 ) -> BoxStream<'static, OSResult<Path>> {
6574 self.target.delete_stream(locations)
6575 }
6576
6577 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
6578 let entries_listed = self.entries_listed.clone();
6579 self.target
6580 .list(prefix)
6581 .inspect(move |_| {
6582 entries_listed.fetch_add(1, Ordering::SeqCst);
6583 })
6584 .boxed()
6585 }
6586
6587 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
6588 self.target.list_with_delimiter(prefix).await
6589 }
6590
6591 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
6592 self.target.copy_opts(from, to, opts).await
6593 }
6594 }
6595
6596 #[derive(Debug)]
6597 struct EntryCountingMemoryProvider {
6598 entries_listed: Arc<AtomicUsize>,
6599 }
6600
6601 #[async_trait]
6602 impl lance_io::object_store::ObjectStoreProvider for EntryCountingMemoryProvider {
6603 async fn new_store(
6604 &self,
6605 base_path: Url,
6606 params: &ObjectStoreParams,
6607 ) -> Result<ObjectStore> {
6608 let mut store = MemoryStoreProvider.new_store(base_path, params).await?;
6609 store.inner = Arc::new(EntryCountingStore {
6610 target: store.inner.clone(),
6611 entries_listed: self.entries_listed.clone(),
6612 });
6613 Ok(store)
6614 }
6615
6616 fn extract_path(&self, url: &Url) -> Result<Path> {
6617 MemoryStoreProvider.extract_path(url)
6618 }
6619
6620 fn calculate_object_store_prefix(
6621 &self,
6622 url: &Url,
6623 storage_options: Option<&HashMap<String, String>>,
6624 ) -> Result<String> {
6625 MemoryStoreProvider.calculate_object_store_prefix(url, storage_options)
6626 }
6627 }
6628
6629 let entries_listed = Arc::new(AtomicUsize::new(0));
6630 let registry = Arc::new(ObjectStoreRegistry::default());
6631 registry.insert(
6632 "memory-object-store",
6633 Arc::new(EntryCountingMemoryProvider {
6634 entries_listed: entries_listed.clone(),
6635 }),
6636 );
6637 let session = Arc::new(Session::new(0, 0, registry));
6638 let ns = DirectoryNamespaceBuilder::new("memory-object-store://lv-count")
6639 .session(session)
6640 .build()
6641 .await
6642 .unwrap();
6643 assert!(ns.object_store.list_is_lexically_ordered);
6644
6645 let table_path = ns.base_path.clone().join("lv_count.lance");
6646 for v in 1..=100u64 {
6647 let p = ManifestNamingScheme::V2.manifest_path(&table_path, v);
6648 ns.object_store.put(&p, b"m".as_slice()).await.unwrap();
6649 }
6650 let staging = Path::parse(format!(
6652 "{}-cee4fbbb-eb19-4ea3-8ca7-54f5ec33dedc",
6653 ManifestNamingScheme::V2.manifest_path(&table_path, 101)
6654 ))
6655 .unwrap();
6656 ns.object_store
6657 .put(&staging, b"s".as_slice())
6658 .await
6659 .unwrap();
6660
6661 let consumed_before = entries_listed.load(Ordering::SeqCst);
6662 let got = ns
6663 .list_versions_under(&table_path, true, Some(1))
6664 .await
6665 .unwrap();
6666 let consumed = entries_listed.load(Ordering::SeqCst) - consumed_before;
6667
6668 assert_eq!(got.len(), 1);
6669 assert_eq!(got[0].version, 100);
6670 assert_eq!(
6671 consumed, 2,
6672 "latest-version query must consume only the staging entry plus the \
6673 first committed manifest, not the whole directory (consumed {} of \
6674 101 entries)",
6675 consumed
6676 );
6677 }
6678
6679 #[derive(Debug)]
6680 #[allow(dead_code)]
6681 struct CountingFileStoreProvider {
6682 listing_count: Arc<AtomicUsize>,
6683 }
6684
6685 #[async_trait]
6686 impl lance_io::object_store::ObjectStoreProvider for CountingFileStoreProvider {
6687 async fn new_store(
6688 &self,
6689 base_path: Url,
6690 params: &ObjectStoreParams,
6691 ) -> Result<ObjectStore> {
6692 let provider = FileStoreProvider;
6693 let mut store = provider.new_store(base_path, params).await?;
6694 store.inner = Arc::new(CountingObjectStore::new(
6695 store.inner.clone(),
6696 self.listing_count.clone(),
6697 ));
6698 Ok(store)
6699 }
6700
6701 fn extract_path(&self, url: &Url) -> Result<Path> {
6702 let provider = FileStoreProvider;
6703 provider.extract_path(url)
6704 }
6705
6706 fn calculate_object_store_prefix(
6707 &self,
6708 url: &Url,
6709 storage_options: Option<&HashMap<String, String>>,
6710 ) -> Result<String> {
6711 let provider = FileStoreProvider;
6712 provider.calculate_object_store_prefix(url, storage_options)
6713 }
6714 }
6715
6716 #[allow(dead_code)]
6717 fn file_object_store_uri(path: &str) -> String {
6718 let file_url = uri_to_url(path).unwrap();
6719 let mut url = Url::parse("file-object-store:///").unwrap();
6720 url.set_path(file_url.path());
6721 url.to_string()
6722 }
6723
6724 #[allow(dead_code)]
6725 fn build_listing_counting_session(listing_count: Arc<AtomicUsize>) -> Arc<Session> {
6726 let registry = Arc::new(ObjectStoreRegistry::default());
6727 registry.insert(
6728 "file-object-store",
6729 Arc::new(CountingFileStoreProvider { listing_count }),
6730 );
6731 Arc::new(Session::new(0, 0, registry))
6732 }
6733
6734 use futures::stream::BoxStream;
6738 use object_store::{
6739 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
6740 PutMultipartOptions, PutPayload, PutResult, Result as OSResult,
6741 };
6742 use std::ops::Range;
6743
6744 #[derive(Debug, Clone, Copy)]
6745 enum ListBehavior {
6746 Throttle,
6747 ServiceUnavailable,
6748 Internal,
6749 NotFound,
6750 EmptyListing,
6751 }
6752
6753 #[derive(Debug)]
6754 struct FailingListStore {
6755 target: Arc<dyn OSObjectStore>,
6756 behavior: Arc<Mutex<Option<ListBehavior>>>,
6757 }
6758
6759 impl std::fmt::Display for FailingListStore {
6760 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6761 write!(f, "FailingListStore({})", self.target)
6762 }
6763 }
6764
6765 #[async_trait]
6766 impl OSObjectStore for FailingListStore {
6767 async fn put_opts(
6768 &self,
6769 location: &Path,
6770 bytes: PutPayload,
6771 opts: PutOptions,
6772 ) -> OSResult<PutResult> {
6773 self.target.put_opts(location, bytes, opts).await
6774 }
6775
6776 async fn put_multipart_opts(
6777 &self,
6778 location: &Path,
6779 opts: PutMultipartOptions,
6780 ) -> OSResult<Box<dyn MultipartUpload>> {
6781 self.target.put_multipart_opts(location, opts).await
6782 }
6783
6784 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
6785 self.target.get_opts(location, options).await
6786 }
6787
6788 async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> OSResult<Vec<Bytes>> {
6789 self.target.get_ranges(location, ranges).await
6790 }
6791
6792 fn delete_stream(
6793 &self,
6794 locations: BoxStream<'static, OSResult<Path>>,
6795 ) -> BoxStream<'static, OSResult<Path>> {
6796 self.target.delete_stream(locations)
6797 }
6798
6799 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
6800 self.target.list(prefix)
6801 }
6802
6803 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
6804 let behavior = *self.behavior.lock().unwrap();
6805 match behavior {
6806 None => self.target.list_with_delimiter(prefix).await,
6807 Some(ListBehavior::EmptyListing) => Ok(ListResult {
6808 common_prefixes: Vec::new(),
6809 objects: Vec::new(),
6810 extensions: Default::default(),
6811 }),
6812 Some(ListBehavior::Throttle) => Err(ObjectStoreError::Generic {
6815 store: "test",
6816 source: "Error performing list request: response error, after 3 retries, \
6817 max_retries: 3, retry_timeout: 180s - HTTP status server error \
6818 (503 Service Unavailable): ServerBusy: The server is busy"
6819 .into(),
6820 }),
6821 Some(ListBehavior::ServiceUnavailable) => Err(ObjectStoreError::Generic {
6822 store: "test",
6823 source: "Error performing list request: 503 Service Unavailable".into(),
6824 }),
6825 Some(ListBehavior::Internal) => Err(ObjectStoreError::Generic {
6826 store: "test",
6827 source: "Error performing list request: catastrophic unclassified failure"
6828 .into(),
6829 }),
6830 Some(ListBehavior::NotFound) => Err(ObjectStoreError::NotFound {
6831 path: "test_table.lance".to_string(),
6832 source: "entity not found".into(),
6833 }),
6834 }
6835 }
6836
6837 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
6838 self.target.copy_opts(from, to, opts).await
6839 }
6840 }
6841
6842 #[derive(Debug)]
6843 struct FailingListStoreProvider {
6844 behavior: Arc<Mutex<Option<ListBehavior>>>,
6845 }
6846
6847 #[async_trait]
6848 impl lance_io::object_store::ObjectStoreProvider for FailingListStoreProvider {
6849 async fn new_store(
6850 &self,
6851 base_path: Url,
6852 params: &ObjectStoreParams,
6853 ) -> Result<ObjectStore> {
6854 let mut store = FileStoreProvider.new_store(base_path, params).await?;
6855 store.inner = Arc::new(FailingListStore {
6856 target: store.inner.clone(),
6857 behavior: self.behavior.clone(),
6858 });
6859 Ok(store)
6860 }
6861
6862 fn extract_path(&self, url: &Url) -> Result<Path> {
6863 FileStoreProvider.extract_path(url)
6864 }
6865
6866 fn calculate_object_store_prefix(
6867 &self,
6868 url: &Url,
6869 storage_options: Option<&HashMap<String, String>>,
6870 ) -> Result<String> {
6871 FileStoreProvider.calculate_object_store_prefix(url, storage_options)
6872 }
6873 }
6874
6875 fn build_failing_list_session(behavior: Arc<Mutex<Option<ListBehavior>>>) -> Arc<Session> {
6876 let registry = Arc::new(ObjectStoreRegistry::default());
6877 registry.insert(
6878 "file-object-store",
6879 Arc::new(FailingListStoreProvider { behavior }),
6880 );
6881 Arc::new(Session::new(0, 0, registry))
6882 }
6883
6884 async fn failing_list_namespace() -> (
6888 DirectoryNamespace,
6889 TempStdDir,
6890 Arc<Mutex<Option<ListBehavior>>>,
6891 ) {
6892 let temp_dir = TempStdDir::default();
6893 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
6894 let behavior = Arc::new(Mutex::new(None));
6895 let session = build_failing_list_session(behavior.clone());
6896 let namespace = DirectoryNamespaceBuilder::new(root_uri)
6897 .session(session)
6898 .manifest_enabled(false)
6899 .dir_listing_enabled(true)
6900 .build()
6901 .await
6902 .unwrap();
6903 (namespace, temp_dir, behavior)
6904 }
6905
6906 async fn create_named_dir_table(namespace: &DirectoryNamespace, name: &str) {
6907 let schema = create_test_schema();
6908 let ipc_data = create_test_ipc_data(&schema);
6909 let mut create_req = CreateTableRequest::new();
6910 create_req.id = Some(vec![name.to_string()]);
6911 namespace
6912 .create_table(create_req, Bytes::from(ipc_data))
6913 .await
6914 .unwrap();
6915 }
6916
6917 #[tokio::test]
6922 async fn test_table_resolution_propagates_storage_errors_not_table_not_found() {
6923 for (behavior, expected_code, evidence) in [
6924 (ListBehavior::Throttle, ErrorCode::Throttling, "serverbusy"),
6925 (
6926 ListBehavior::ServiceUnavailable,
6927 ErrorCode::ServiceUnavailable,
6928 "503 service unavailable",
6929 ),
6930 (ListBehavior::Internal, ErrorCode::Internal, "catastrophic"),
6931 ] {
6932 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6933 create_named_dir_table(&namespace, "checkpoint").await;
6934 *toggle.lock().unwrap() = Some(behavior);
6935
6936 let mut describe_req = DescribeTableRequest::new();
6937 describe_req.id = Some(vec!["checkpoint".to_string()]);
6938 let err = namespace.describe_table(describe_req).await.unwrap_err();
6939 let msg = err.to_string();
6940 assert_eq!(
6941 mutation_error_code(err),
6942 expected_code,
6943 "describe_table under {behavior:?}; msg: {msg}"
6944 );
6945 assert!(
6946 msg.to_ascii_lowercase().contains(evidence),
6947 "describe_table message must carry storage evidence '{evidence}', got: {msg}"
6948 );
6949
6950 let mut exists_req = TableExistsRequest::new();
6951 exists_req.id = Some(vec!["checkpoint".to_string()]);
6952 let err = namespace.table_exists(exists_req).await.unwrap_err();
6953 let msg = err.to_string();
6954 assert_eq!(
6955 mutation_error_code(err),
6956 expected_code,
6957 "table_exists under {behavior:?}; msg: {msg}"
6958 );
6959 assert!(
6960 msg.to_ascii_lowercase().contains(evidence),
6961 "table_exists message must carry storage evidence '{evidence}', got: {msg}"
6962 );
6963 }
6964 }
6965
6966 #[tokio::test]
6969 async fn test_table_resolution_missing_table_yields_table_not_found() {
6970 for behavior in [ListBehavior::NotFound, ListBehavior::EmptyListing] {
6971 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6972 *toggle.lock().unwrap() = Some(behavior);
6973
6974 let mut describe_req = DescribeTableRequest::new();
6975 describe_req.id = Some(vec!["missing".to_string()]);
6976 let err = namespace.describe_table(describe_req).await.unwrap_err();
6977 assert_eq!(
6978 mutation_error_code(err),
6979 ErrorCode::TableNotFound,
6980 "describe_table under {behavior:?} should be TableNotFound"
6981 );
6982
6983 let mut exists_req = TableExistsRequest::new();
6984 exists_req.id = Some(vec!["missing".to_string()]);
6985 let err = namespace.table_exists(exists_req).await.unwrap_err();
6986 assert_eq!(
6987 mutation_error_code(err),
6988 ErrorCode::TableNotFound,
6989 "table_exists under {behavior:?} should be TableNotFound"
6990 );
6991 }
6992 }
6993
6994 #[tokio::test]
7004 async fn test_hybrid_resolution_falls_through_and_does_not_mask_throttle() {
7005 let temp_dir = TempStdDir::default();
7006 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
7007 let behavior = Arc::new(Mutex::new(None));
7008 let session = build_failing_list_session(behavior.clone());
7009
7010 let manifest_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
7012 .session(session.clone())
7013 .manifest_enabled(true)
7014 .dir_listing_enabled(true)
7015 .build()
7016 .await
7017 .unwrap();
7018 create_named_dir_table(&manifest_ns, "seed").await;
7019
7020 let dir_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
7022 .session(session.clone())
7023 .manifest_enabled(false)
7024 .dir_listing_enabled(true)
7025 .build()
7026 .await
7027 .unwrap();
7028 create_named_dir_table(&dir_ns, "checkpoint").await;
7029
7030 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
7033 .session(session)
7034 .manifest_enabled(true)
7035 .dir_listing_enabled(true)
7036 .dir_listing_to_manifest_migration_enabled(true)
7037 .build()
7038 .await
7039 .unwrap();
7040
7041 let mut describe_req = DescribeTableRequest::new();
7044 describe_req.id = Some(vec!["checkpoint".to_string()]);
7045 hybrid_ns
7046 .describe_table(describe_req)
7047 .await
7048 .expect("unregistered on-disk table should resolve via the manifest fall-through");
7049
7050 *behavior.lock().unwrap() = Some(ListBehavior::Throttle);
7054 let mut describe_req = DescribeTableRequest::new();
7055 describe_req.id = Some(vec!["checkpoint".to_string()]);
7056 let err = hybrid_ns.describe_table(describe_req).await.unwrap_err();
7057 let code = mutation_error_code(err);
7058 assert_ne!(
7059 code,
7060 ErrorCode::TableNotFound,
7061 "hybrid resolution masked a throttle as TableNotFound"
7062 );
7063 assert!(
7064 matches!(
7065 code,
7066 ErrorCode::Throttling | ErrorCode::ServiceUnavailable | ErrorCode::Internal
7067 ),
7068 "hybrid resolution should surface a storage error, got {code:?}"
7069 );
7070 }
7071
7072 #[test]
7073 fn test_classify_storage_error_maps_variants_and_preserves_evidence() {
7074 let throttle: Error = ObjectStoreError::Generic {
7075 store: "test",
7076 source: "list request failed, after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
7077 }
7078 .into();
7079 assert!(matches!(throttle, Error::IO { .. }));
7080 let classified = DirectoryNamespace::classify_storage_error(throttle);
7081 let msg = classified.to_string();
7082 assert_eq!(mutation_error_code(classified), ErrorCode::Throttling);
7083 assert!(
7084 msg.to_ascii_lowercase().contains("serverbusy"),
7085 "throttle evidence lost: {msg}"
7086 );
7087
7088 let service: Error = ObjectStoreError::Generic {
7089 store: "test",
7090 source: "504 Gateway Timeout".into(),
7091 }
7092 .into();
7093 assert_eq!(
7094 mutation_error_code(DirectoryNamespace::classify_storage_error(service)),
7095 ErrorCode::ServiceUnavailable
7096 );
7097
7098 let internal: Error = ObjectStoreError::Generic {
7099 store: "test",
7100 source: "disk caught fire".into(),
7101 }
7102 .into();
7103 assert_eq!(
7104 mutation_error_code(DirectoryNamespace::classify_storage_error(internal)),
7105 ErrorCode::Internal
7106 );
7107
7108 let preexisting: Error = NamespaceError::TableAlreadyExists {
7110 message: "t".to_string(),
7111 }
7112 .into();
7113 assert_eq!(
7114 mutation_error_code(DirectoryNamespace::classify_storage_error(preexisting)),
7115 ErrorCode::TableAlreadyExists
7116 );
7117 }
7118
7119 #[test]
7120 fn test_is_manifest_table_absent_error() {
7121 let table_not_found: Error = NamespaceError::TableNotFound {
7122 message: "t".to_string(),
7123 }
7124 .into();
7125 assert!(DirectoryNamespace::is_manifest_table_absent_error(
7126 &table_not_found
7127 ));
7128 let raw_not_found: Error = ObjectStoreError::NotFound {
7129 path: "t".to_string(),
7130 source: "x".into(),
7131 }
7132 .into();
7133 assert!(DirectoryNamespace::is_manifest_table_absent_error(
7134 &raw_not_found
7135 ));
7136
7137 let throttle: Error = ObjectStoreError::Generic {
7138 store: "test",
7139 source: "after 3 retries, max_retries: 3 ServerBusy".into(),
7140 }
7141 .into();
7142 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
7143 &throttle
7144 ));
7145 let internal: Error = NamespaceError::Internal {
7146 message: "boom".to_string(),
7147 }
7148 .into();
7149 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
7150 &internal
7151 ));
7152 }
7153
7154 #[test]
7155 fn test_map_open_error() {
7156 let not_found = || NamespaceError::TableNotFound {
7157 message: "table at 'x' not found: ...".to_string(),
7158 };
7159
7160 let throttle: Error = ObjectStoreError::Generic {
7161 store: "test",
7162 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
7163 }
7164 .into();
7165 assert_eq!(
7166 mutation_error_code(DirectoryNamespace::map_open_error(throttle, not_found())),
7167 ErrorCode::Throttling
7168 );
7169
7170 let generic_io: Error = ObjectStoreError::Generic {
7171 store: "test",
7172 source: "connection reset".into(),
7173 }
7174 .into();
7175 assert_eq!(
7176 mutation_error_code(DirectoryNamespace::map_open_error(generic_io, not_found())),
7177 ErrorCode::Internal
7178 );
7179
7180 let io_not_found: Error = ObjectStoreError::NotFound {
7181 path: "x".to_string(),
7182 source: "missing".into(),
7183 }
7184 .into();
7185 assert_eq!(
7186 mutation_error_code(DirectoryNamespace::map_open_error(
7187 io_not_found,
7188 not_found()
7189 )),
7190 ErrorCode::TableNotFound
7191 );
7192
7193 let dataset_not_found = Error::dataset_not_found("x".to_string(), "missing".into());
7194 assert_eq!(
7195 mutation_error_code(DirectoryNamespace::map_open_error(
7196 dataset_not_found,
7197 not_found()
7198 )),
7199 ErrorCode::TableNotFound
7200 );
7201
7202 let ref_not_found = Error::RefNotFound {
7204 message: "branch 'b' does not exist".to_string(),
7205 };
7206 assert_eq!(
7207 mutation_error_code(DirectoryNamespace::map_open_error(
7208 ref_not_found,
7209 not_found()
7210 )),
7211 ErrorCode::TableNotFound
7212 );
7213
7214 let version_miss = Error::RefNotFound {
7216 message: "version 5 does not exist".to_string(),
7217 };
7218 assert_eq!(
7219 mutation_error_code(DirectoryNamespace::map_open_error(
7220 version_miss,
7221 NamespaceError::TableVersionNotFound {
7222 message: "version 5 not found".to_string(),
7223 },
7224 )),
7225 ErrorCode::TableVersionNotFound
7226 );
7227 let version_throttle: Error = ObjectStoreError::Generic {
7228 store: "test",
7229 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
7230 }
7231 .into();
7232 assert_eq!(
7233 mutation_error_code(DirectoryNamespace::map_open_error(
7234 version_throttle,
7235 NamespaceError::TableVersionNotFound {
7236 message: "version 5 not found".to_string(),
7237 },
7238 )),
7239 ErrorCode::Throttling
7240 );
7241 }
7242
7243 fn create_test_ipc_data(schema: &JsonArrowSchema) -> Vec<u8> {
7245 use arrow::ipc::writer::StreamWriter;
7246
7247 let arrow_schema = convert_json_arrow_schema(schema).unwrap();
7248 let arrow_schema = Arc::new(arrow_schema);
7249 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
7250 let mut buffer = Vec::new();
7251 {
7252 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
7253 writer.write(&batch).unwrap();
7254 writer.finish().unwrap();
7255 }
7256 buffer
7257 }
7258
7259 fn create_ipc_data_from_batches(
7260 schema: Arc<arrow_schema::Schema>,
7261 batches: Vec<arrow::record_batch::RecordBatch>,
7262 ) -> Vec<u8> {
7263 use arrow::ipc::writer::StreamWriter;
7264
7265 let mut buffer = Vec::new();
7266 {
7267 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
7268 for batch in &batches {
7269 writer.write(batch).unwrap();
7270 }
7271 writer.finish().unwrap();
7272 }
7273 buffer
7274 }
7275
7276 fn create_non_empty_test_ipc_data() -> Vec<u8> {
7277 use arrow::array::{Int32Array, StringArray};
7278 use arrow::record_batch::RecordBatch;
7279
7280 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
7281 let batch = RecordBatch::try_new(
7282 schema.clone(),
7283 vec![
7284 Arc::new(Int32Array::from(vec![1, 2])),
7285 Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])),
7286 ],
7287 )
7288 .unwrap();
7289 create_ipc_data_from_batches(schema, vec![batch])
7290 }
7291
7292 fn create_single_row_test_ipc_data() -> Vec<u8> {
7293 use arrow::array::{Int32Array, StringArray};
7294 use arrow::record_batch::RecordBatch;
7295
7296 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
7297 let batch = RecordBatch::try_new(
7298 schema.clone(),
7299 vec![
7300 Arc::new(Int32Array::from(vec![10])),
7301 Arc::new(StringArray::from(vec![Some("carol")])),
7302 ],
7303 )
7304 .unwrap();
7305 create_ipc_data_from_batches(schema, vec![batch])
7306 }
7307
7308 fn create_test_schema() -> JsonArrowSchema {
7310 let int_type = JsonArrowDataType::new("int32".to_string());
7311 let string_type = JsonArrowDataType::new("utf8".to_string());
7312
7313 let id_field = JsonArrowField {
7314 name: "id".to_string(),
7315 r#type: Box::new(int_type),
7316 nullable: false,
7317 metadata: None,
7318 };
7319
7320 let name_field = JsonArrowField {
7321 name: "name".to_string(),
7322 r#type: Box::new(string_type),
7323 nullable: true,
7324 metadata: None,
7325 };
7326
7327 JsonArrowSchema {
7328 fields: vec![id_field, name_field],
7329 metadata: None,
7330 }
7331 }
7332
7333 fn create_scalar_table_ipc_data() -> Vec<u8> {
7334 use arrow::array::{Int32Array, StringArray};
7335 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7336
7337 let schema = Arc::new(ArrowSchema::new(vec![
7338 Field::new("id", DataType::Int32, false),
7339 Field::new("name", DataType::Utf8, true),
7340 ]));
7341 let batch = arrow::record_batch::RecordBatch::try_new(
7342 schema.clone(),
7343 vec![
7344 Arc::new(Int32Array::from(vec![1, 2, 3])),
7345 Arc::new(StringArray::from(vec!["alice", "bob", "cory"])),
7346 ],
7347 )
7348 .unwrap();
7349 create_ipc_data_from_batches(schema, vec![batch])
7350 }
7351
7352 async fn create_legacy_manifest_without_primary_key_metadata(root: &str) {
7353 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7354 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
7355
7356 let schema = Arc::new(ArrowSchema::new(vec![
7357 Field::new("object_id", DataType::Utf8, false),
7358 Field::new("object_type", DataType::Utf8, false),
7359 Field::new("location", DataType::Utf8, true),
7360 Field::new("metadata", DataType::Utf8, true),
7361 Field::new(
7362 "base_objects",
7363 DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))),
7364 true,
7365 ),
7366 ]));
7367 let batch = RecordBatch::new_empty(schema.clone());
7368 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
7369 Dataset::write(Box::new(reader), &format!("{}/__manifest", root), None)
7370 .await
7371 .unwrap();
7372 }
7373
7374 async fn manifest_has_primary_key_metadata(root: &str) -> bool {
7375 let dataset = Dataset::open(&format!("{}/__manifest", root))
7376 .await
7377 .unwrap();
7378 dataset
7379 .schema()
7380 .field("object_id")
7381 .map(|field| {
7382 field
7383 .metadata
7384 .contains_key(lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
7385 })
7386 .unwrap_or(false)
7387 }
7388
7389 fn create_vector_table_ipc_data() -> Vec<u8> {
7390 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
7391 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7392
7393 let schema = Arc::new(ArrowSchema::new(vec![
7394 Field::new("id", DataType::Int32, false),
7395 Field::new(
7396 "vector",
7397 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
7398 true,
7399 ),
7400 ]));
7401 let vector_field = Arc::new(Field::new("item", DataType::Float32, true));
7402 let vectors = FixedSizeListArray::try_new(
7403 vector_field,
7404 2,
7405 Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6])),
7406 None,
7407 )
7408 .unwrap();
7409 let batch = arrow::record_batch::RecordBatch::try_new(
7410 schema.clone(),
7411 vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(vectors)],
7412 )
7413 .unwrap();
7414 create_ipc_data_from_batches(schema, vec![batch])
7415 }
7416
7417 async fn create_scalar_table(namespace: &DirectoryNamespace, table_name: &str) {
7418 let mut create_table_request = CreateTableRequest::new();
7419 create_table_request.id = Some(vec![table_name.to_string()]);
7420 namespace
7421 .create_table(
7422 create_table_request,
7423 Bytes::from(create_scalar_table_ipc_data()),
7424 )
7425 .await
7426 .unwrap();
7427 }
7428
7429 async fn create_vector_table(namespace: &DirectoryNamespace, table_name: &str) {
7430 let mut create_table_request = CreateTableRequest::new();
7431 create_table_request.id = Some(vec![table_name.to_string()]);
7432 namespace
7433 .create_table(
7434 create_table_request,
7435 Bytes::from(create_vector_table_ipc_data()),
7436 )
7437 .await
7438 .unwrap();
7439 }
7440
7441 async fn open_dataset(namespace: &DirectoryNamespace, table_name: &str) -> Dataset {
7442 let mut describe_request = DescribeTableRequest::new();
7443 describe_request.id = Some(vec![table_name.to_string()]);
7444 let table_uri = namespace
7445 .describe_table(describe_request)
7446 .await
7447 .unwrap()
7448 .location
7449 .expect("table location should exist");
7450 Dataset::open(&table_uri).await.unwrap()
7451 }
7452
7453 async fn create_scalar_index(
7454 namespace: &DirectoryNamespace,
7455 table_name: &str,
7456 index_name: &str,
7457 ) -> Option<String> {
7458 use lance_namespace::models::CreateTableIndexRequest;
7459
7460 let mut create_index_request =
7461 CreateTableIndexRequest::new("id".to_string(), "BTREE".to_string());
7462 create_index_request.id = Some(vec![table_name.to_string()]);
7463 create_index_request.name = Some(index_name.to_string());
7464 namespace
7465 .create_table_scalar_index(create_index_request)
7466 .await
7467 .unwrap()
7468 .transaction_id
7469 }
7470
7471 async fn create_branch_with_commits(
7476 namespace: &DirectoryNamespace,
7477 table_name: &str,
7478 branch_name: &str,
7479 extra_versions: usize,
7480 ) -> String {
7481 let mut main = open_dataset(namespace, table_name).await;
7482 let fork_version = main.version().version;
7483 let branch = main
7484 .create_branch(branch_name, fork_version, None)
7485 .await
7486 .unwrap();
7487 let branch_uri = branch.uri().to_string();
7488 for i in 0..extra_versions {
7489 append_scalar_version(&branch_uri, (i as i32 + 1) * 100).await;
7490 }
7491 branch_uri
7492 }
7493
7494 async fn append_scalar_version(uri: &str, seed: i32) {
7497 use arrow::array::{Int32Array, StringArray};
7498 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7499 let schema = Arc::new(ArrowSchema::new(vec![
7500 Field::new("id", DataType::Int32, false),
7501 Field::new("name", DataType::Utf8, true),
7502 ]));
7503 let batch = arrow::record_batch::RecordBatch::try_new(
7504 schema.clone(),
7505 vec![
7506 Arc::new(Int32Array::from(vec![seed, seed + 1])),
7507 Arc::new(StringArray::from(vec![Some("x"), Some("y")])),
7508 ],
7509 )
7510 .unwrap();
7511 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
7512 Dataset::write(
7513 reader,
7514 uri,
7515 Some(WriteParams {
7516 mode: WriteMode::Append,
7517 ..Default::default()
7518 }),
7519 )
7520 .await
7521 .unwrap();
7522 }
7523
7524 async fn list_versions(
7526 namespace: &DirectoryNamespace,
7527 table_name: &str,
7528 branch: Option<&str>,
7529 ) -> Result<Vec<TableVersion>> {
7530 let req = ListTableVersionsRequest {
7531 id: Some(vec![table_name.to_string()]),
7532 branch: branch.map(|b| b.to_string()),
7533 ..Default::default()
7534 };
7535 namespace.list_table_versions(req).await.map(|r| r.versions)
7536 }
7537
7538 #[tokio::test]
7539 async fn test_list_table_versions_on_branch() {
7540 let (namespace, _temp_dir) = create_test_namespace().await;
7541 create_scalar_table(&namespace, "users").await;
7542 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7543
7544 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7547 .await
7548 .unwrap();
7549 assert!(branch_versions.len() >= 2);
7550 assert!(
7551 branch_versions
7552 .iter()
7553 .all(|v| v.manifest_path.contains("tree/exp")),
7554 "branch versions must resolve to branch manifests: {:?}",
7555 branch_versions
7556 );
7557
7558 let main_versions = list_versions(&namespace, "users", None).await.unwrap();
7560 let main_explicit = list_versions(&namespace, "users", Some("main"))
7561 .await
7562 .unwrap();
7563 assert_eq!(main_versions.len(), main_explicit.len());
7564 assert!(
7565 main_versions
7566 .iter()
7567 .all(|v| !v.manifest_path.contains("tree/"))
7568 );
7569
7570 let missing = list_versions(&namespace, "users", Some("does-not-exist")).await;
7572 assert!(missing.is_err());
7573 assert!(missing.unwrap_err().to_string().contains("not found"));
7574 }
7575
7576 #[tokio::test]
7577 async fn test_describe_table_version_on_branch() {
7578 let (namespace, _temp_dir) = create_test_namespace().await;
7579 create_scalar_table(&namespace, "users").await;
7580 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7581
7582 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7583 .await
7584 .unwrap();
7585 let latest = branch_versions.iter().map(|v| v.version).max().unwrap();
7586
7587 let req = DescribeTableVersionRequest {
7589 id: Some(vec!["users".to_string()]),
7590 branch: Some("exp".to_string()),
7591 ..Default::default()
7592 };
7593 let resp = namespace.describe_table_version(req).await.unwrap();
7594 assert_eq!(resp.version.version, latest);
7595 assert!(resp.version.manifest_path.contains("tree/exp"));
7596
7597 let req = DescribeTableVersionRequest {
7599 id: Some(vec!["users".to_string()]),
7600 version: Some(latest),
7601 branch: Some("exp".to_string()),
7602 ..Default::default()
7603 };
7604 assert!(namespace.describe_table_version(req).await.is_ok());
7605
7606 let req = DescribeTableVersionRequest {
7608 id: Some(vec!["users".to_string()]),
7609 version: Some(999_999),
7610 branch: Some("exp".to_string()),
7611 ..Default::default()
7612 };
7613 assert!(namespace.describe_table_version(req).await.is_err());
7614
7615 let req = DescribeTableVersionRequest {
7617 id: Some(vec!["users".to_string()]),
7618 branch: Some("nope".to_string()),
7619 ..Default::default()
7620 };
7621 let err = namespace.describe_table_version(req).await;
7622 assert!(err.is_err() && err.unwrap_err().to_string().contains("not found"));
7623 }
7624
7625 #[tokio::test]
7626 async fn test_restore_table_on_branch() {
7627 use lance_namespace::models::RestoreTableRequest;
7628
7629 let (namespace, _temp_dir) = create_test_namespace().await;
7630 create_scalar_table(&namespace, "users").await;
7631 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7632
7633 let before = list_versions(&namespace, "users", Some("exp"))
7634 .await
7635 .unwrap();
7636 let branch_latest = before.iter().map(|v| v.version).max().unwrap();
7637 let earliest = before.iter().map(|v| v.version).min().unwrap();
7638 let main_before = list_versions(&namespace, "users", None)
7639 .await
7640 .unwrap()
7641 .len();
7642
7643 let req = RestoreTableRequest {
7646 id: Some(vec!["users".to_string()]),
7647 version: earliest,
7648 branch: Some("exp".to_string()),
7649 ..Default::default()
7650 };
7651 let resp = namespace.restore_table(req).await.unwrap();
7652 assert!(resp.transaction_id.is_some());
7653
7654 let after = list_versions(&namespace, "users", Some("exp"))
7655 .await
7656 .unwrap();
7657 let new_latest = after.iter().map(|v| v.version).max().unwrap();
7658 assert!(
7659 new_latest > branch_latest,
7660 "restore should add a branch version"
7661 );
7662
7663 let main_after = list_versions(&namespace, "users", None)
7664 .await
7665 .unwrap()
7666 .len();
7667 assert_eq!(main_after, main_before, "main must be unaffected");
7668 }
7669
7670 #[tokio::test]
7671 async fn test_batch_delete_table_versions_on_branch() {
7672 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7673
7674 let (namespace, _temp_dir) = create_test_namespace().await;
7675 create_scalar_table(&namespace, "users").await;
7676 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7677
7678 let before = list_versions(&namespace, "users", Some("exp"))
7679 .await
7680 .unwrap();
7681 let main_before = list_versions(&namespace, "users", None).await.unwrap();
7682
7683 let req = BatchDeleteTableVersionsRequest {
7688 id: Some(vec!["users".to_string()]),
7689 branch: Some("exp".to_string()),
7690 ranges: vec![VersionRange::new(0, -1)],
7691 ..Default::default()
7692 };
7693 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7694 assert_eq!(
7695 resp.deleted_count,
7696 Some(before.len() as i64),
7697 "every branch manifest should be physically deleted"
7698 );
7699
7700 assert!(
7702 list_versions(&namespace, "users", Some("exp"))
7703 .await
7704 .is_err()
7705 );
7706 let main_after = list_versions(&namespace, "users", None).await.unwrap();
7707 assert_eq!(
7708 main_after.len(),
7709 main_before.len(),
7710 "main must be untouched"
7711 );
7712 }
7713
7714 #[tokio::test]
7715 async fn test_create_table_version_on_branch() {
7716 use futures::TryStreamExt;
7717 use lance_namespace::models::CreateTableVersionRequest;
7718
7719 let (namespace, _temp_dir) = create_test_namespace().await;
7720 create_scalar_table(&namespace, "users").await;
7721 let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 1).await;
7722
7723 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7725 let versions_dir = branch_ds.versions_dir();
7726 let store = branch_ds.object_store(None).await.unwrap();
7727 let existing = store
7728 .inner
7729 .list(Some(&versions_dir))
7730 .try_collect::<Vec<_>>()
7731 .await
7732 .unwrap()
7733 .into_iter()
7734 .find(|m| {
7735 m.location
7736 .filename()
7737 .map(|f| f.ends_with(".manifest"))
7738 .unwrap_or(false)
7739 })
7740 .expect("a branch manifest");
7741 let bytes = store
7742 .inner
7743 .get(&existing.location)
7744 .await
7745 .unwrap()
7746 .bytes()
7747 .await
7748 .unwrap();
7749 let staging = versions_dir.join("staging_manifest");
7750 store.inner.put(&staging, bytes.into()).await.unwrap();
7751
7752 let main_before = list_versions(&namespace, "users", None)
7753 .await
7754 .unwrap()
7755 .len();
7756 let new_version = list_versions(&namespace, "users", Some("exp"))
7757 .await
7758 .unwrap()
7759 .iter()
7760 .map(|v| v.version)
7761 .max()
7762 .unwrap()
7763 + 1;
7764
7765 let req = CreateTableVersionRequest {
7766 id: Some(vec!["users".to_string()]),
7767 version: new_version,
7768 manifest_path: staging.to_string(),
7769 naming_scheme: Some("V2".to_string()),
7770 branch: Some("exp".to_string()),
7771 ..Default::default()
7772 };
7773 let resp = namespace.create_table_version(req).await.unwrap();
7774 let info = resp.version.expect("version info");
7775 assert!(
7777 info.manifest_path.contains("tree/exp"),
7778 "got {}",
7779 info.manifest_path
7780 );
7781
7782 let after = list_versions(&namespace, "users", Some("exp"))
7784 .await
7785 .unwrap();
7786 assert!(after.iter().any(|v| v.version == new_version));
7787 let main_after = list_versions(&namespace, "users", None)
7788 .await
7789 .unwrap()
7790 .len();
7791 assert_eq!(main_after, main_before, "main must be unaffected");
7792 }
7793
7794 #[tokio::test]
7799 async fn test_external_manifest_store_resolves_branch_from_base_path() {
7800 use futures::TryStreamExt;
7801 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
7802 use lance_table::io::commit::external_manifest::ExternalManifestStore;
7803
7804 let (namespace, _temp_dir) = create_test_namespace().await;
7805 create_scalar_table(&namespace, "users").await; let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 2).await;
7807
7808 let namespace = Arc::new(namespace);
7809 let table_id = vec!["users".to_string()];
7810 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7811 let branch_base = branch_ds.branch_location().path;
7812 let root_base = branch_ds.branch_location().find_main().unwrap().path;
7813 let store = LanceNamespaceExternalManifestStore::new(
7814 namespace.clone(),
7815 table_id.clone(),
7816 root_base.clone(),
7817 );
7818
7819 let (branch_latest, branch_path) = store
7822 .get_latest_version(branch_base.as_ref())
7823 .await
7824 .unwrap()
7825 .expect("branch has versions");
7826 let (_main_latest, main_path) = store
7827 .get_latest_version(root_base.as_ref())
7828 .await
7829 .unwrap()
7830 .expect("main has versions");
7831 assert!(
7832 branch_path.contains("tree/exp"),
7833 "branch latest must resolve to the branch tree: {}",
7834 branch_path
7835 );
7836 assert!(
7837 !main_path.contains("tree/exp"),
7838 "main latest must not resolve to a branch tree: {}",
7839 main_path
7840 );
7841
7842 let described = store
7844 .get(branch_base.as_ref(), branch_latest)
7845 .await
7846 .unwrap();
7847 assert!(
7848 described.contains("tree/exp"),
7849 "describe on the branch must resolve to the branch tree: {}",
7850 described
7851 );
7852
7853 assert!(store.get_latest_version("somewhere/else").await.is_err());
7855
7856 let versions_dir = branch_ds.versions_dir();
7859 let obj = branch_ds.object_store(None).await.unwrap();
7860 let existing = obj
7861 .inner
7862 .list(Some(&versions_dir))
7863 .try_collect::<Vec<_>>()
7864 .await
7865 .unwrap()
7866 .into_iter()
7867 .find(|m| {
7868 m.location
7869 .filename()
7870 .map(|f| f.ends_with(".manifest"))
7871 .unwrap_or(false)
7872 })
7873 .expect("a branch manifest");
7874 let bytes = obj
7875 .inner
7876 .get(&existing.location)
7877 .await
7878 .unwrap()
7879 .bytes()
7880 .await
7881 .unwrap();
7882 let size = bytes.len() as u64;
7883 let staging = versions_dir.clone().join("staging_manifest");
7884 obj.inner.put(&staging, bytes.into()).await.unwrap();
7885
7886 let committed = store
7887 .put(
7888 &branch_base,
7889 branch_latest + 1,
7890 &staging,
7891 size,
7892 None,
7893 obj.inner.as_ref(),
7894 ManifestNamingScheme::V2,
7895 )
7896 .await
7897 .unwrap();
7898 assert!(
7899 committed.path.to_string().contains("tree/exp"),
7900 "a commit through a branch-qualified base must land on the branch tree: {}",
7901 committed.path
7902 );
7903 }
7904
7905 #[tokio::test]
7913 async fn test_write_into_namespace_on_branch_appends_to_branch() {
7914 use lance::dataset::builder::DatasetBuilder;
7915 use lance_namespace::models::CreateTableBranchRequest;
7916
7917 let temp = TempStdDir::default();
7918 let namespace = Arc::new(
7919 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
7920 .manifest_enabled(true)
7921 .table_version_tracking_enabled(true)
7922 .ops_metrics_enabled(true)
7923 .build()
7924 .await
7925 .unwrap(),
7926 );
7927 let ns: Arc<dyn LanceNamespace> = namespace.clone();
7928 let table_id = vec!["t".to_string()];
7929 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
7931 id: Some(table_id.clone()),
7932 name: "exp".to_string(),
7933 ..Default::default()
7934 })
7935 .await
7936 .unwrap();
7937
7938 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
7939 ns.list_table_versions(ListTableVersionsRequest {
7940 id: Some(table_id),
7941 ..Default::default()
7942 })
7943 .await
7944 .unwrap()
7945 .versions
7946 .len()
7947 };
7948 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
7949 let commits_before = namespace
7950 .retrieve_ops_metrics()
7951 .get("create_table_version")
7952 .copied()
7953 .unwrap_or(0);
7954
7955 let branch_ds = Dataset::write_into_namespace_on_branch(
7956 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
7957 ns.clone(),
7958 table_id.clone(),
7959 "exp",
7960 Some(WriteParams {
7961 mode: WriteMode::Append,
7962 ..Default::default()
7963 }),
7964 )
7965 .await
7966 .unwrap();
7967 assert_eq!(branch_ds.manifest.branch.as_deref(), Some("exp"));
7968 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
7969
7970 let commits_after = namespace
7973 .retrieve_ops_metrics()
7974 .get("create_table_version")
7975 .copied()
7976 .unwrap_or(0);
7977 assert_eq!(
7978 commits_after,
7979 commits_before + 1,
7980 "the branch append must register its version via create_table_version"
7981 );
7982 let exp_versions = ns
7983 .list_table_versions(ListTableVersionsRequest {
7984 id: Some(table_id.clone()),
7985 branch: Some("exp".to_string()),
7986 ..Default::default()
7987 })
7988 .await
7989 .unwrap()
7990 .versions;
7991 assert!(
7992 exp_versions
7993 .iter()
7994 .all(|v| v.manifest_path.contains("tree/exp")),
7995 "branch versions must resolve to the branch tree: {:?}",
7996 exp_versions
7997 );
7998 assert_eq!(
7999 main_chain_len(ns.clone(), table_id.clone()).await,
8000 main_before,
8001 "main's catalog must be untouched by the branch append"
8002 );
8003
8004 Dataset::write_into_namespace(
8007 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
8008 ns.clone(),
8009 table_id.clone(),
8010 Some(WriteParams {
8011 mode: WriteMode::Append,
8012 ..Default::default()
8013 }),
8014 )
8015 .await
8016 .unwrap();
8017 assert_eq!(
8018 main_chain_len(ns.clone(), table_id.clone()).await,
8019 main_before + 1,
8020 "a managed main append must register its version in the catalog"
8021 );
8022 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8023 .await
8024 .unwrap()
8025 .load()
8026 .await
8027 .unwrap();
8028 assert_eq!(
8029 scan_id_column(&fresh).await,
8030 vec![1, 2, 100],
8031 "a fresh managed open must resolve the appended version, not a stale latest"
8032 );
8033 }
8034
8035 #[tokio::test]
8037 async fn test_write_into_namespace_on_branch_rejects_create() {
8038 use arrow::array::{Int32Array, StringArray};
8039 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
8040
8041 let (namespace, _temp_dir) = create_test_namespace().await;
8042 let namespace = Arc::new(namespace);
8043
8044 let schema = Arc::new(ArrowSchema::new(vec![
8045 Field::new("id", DataType::Int32, false),
8046 Field::new("name", DataType::Utf8, true),
8047 ]));
8048 let batch = arrow::record_batch::RecordBatch::try_new(
8049 schema.clone(),
8050 vec![
8051 Arc::new(Int32Array::from(vec![1])),
8052 Arc::new(StringArray::from(vec![Some("a")])),
8053 ],
8054 )
8055 .unwrap();
8056 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
8057
8058 let result = Dataset::write_into_namespace_on_branch(
8059 reader,
8060 namespace.clone(),
8061 vec!["new_table".to_string()],
8062 "exp",
8063 Some(WriteParams {
8064 mode: WriteMode::Create,
8065 ..Default::default()
8066 }),
8067 )
8068 .await;
8069 assert!(result.is_err(), "create on a branch must be rejected");
8070 assert!(
8071 result.unwrap_err().to_string().contains("branch"),
8072 "error should mention the branch restriction"
8073 );
8074 }
8075
8076 #[tokio::test]
8077 async fn test_branch_name_validation_rejects_traversal() {
8078 let (namespace, _temp_dir) = create_test_namespace().await;
8079 create_scalar_table(&namespace, "users").await;
8080
8081 let err = list_versions(&namespace, "users", Some("../evil")).await;
8084 assert!(err.is_err());
8085 assert!(err.unwrap_err().to_string().contains("invalid branch name"));
8086 }
8087
8088 #[tokio::test]
8089 async fn test_branch_ops_reject_zombie_branch() {
8090 use futures::TryStreamExt;
8091 use lance_namespace::models::{
8092 BatchDeleteTableVersionsRequest, CreateTableVersionRequest, RestoreTableRequest,
8093 VersionRange,
8094 };
8095
8096 let (namespace, _temp_dir) = create_test_namespace().await;
8097 create_scalar_table(&namespace, "users").await;
8098
8099 let dataset = open_dataset(&namespace, "users").await;
8100 let store = dataset.object_store(None).await.unwrap();
8101 let manifest = store
8102 .inner
8103 .list(Some(&dataset.versions_dir()))
8104 .try_collect::<Vec<_>>()
8105 .await
8106 .unwrap()
8107 .into_iter()
8108 .find(|m| {
8109 m.location
8110 .filename()
8111 .map(|f| f.ends_with(".manifest"))
8112 .unwrap_or(false)
8113 })
8114 .expect("a manifest");
8115 let bytes = store
8116 .inner
8117 .get(&manifest.location)
8118 .await
8119 .unwrap()
8120 .bytes()
8121 .await
8122 .unwrap();
8123 let zombie = dataset
8124 .branch_location()
8125 .find_branch(Some("ghost"))
8126 .unwrap()
8127 .path
8128 .join(VERSIONS_DIR)
8129 .join(manifest.location.filename().unwrap());
8130 store.inner.put(&zombie, bytes.into()).await.unwrap();
8131
8132 assert!(dataset.branches().get("ghost").await.is_err());
8133
8134 fn rejected<T: std::fmt::Debug>(label: &str, r: Result<T>) {
8135 match r {
8136 Ok(v) => panic!("{label} must reject the zombie branch, got Ok({v:?})"),
8137 Err(e) => assert!(e.to_string().contains("not found"), "{label}: {e}"),
8138 }
8139 }
8140
8141 rejected(
8142 "list",
8143 list_versions(&namespace, "users", Some("ghost")).await,
8144 );
8145 rejected(
8146 "describe",
8147 namespace
8148 .describe_table_version(DescribeTableVersionRequest {
8149 id: Some(vec!["users".to_string()]),
8150 branch: Some("ghost".to_string()),
8151 ..Default::default()
8152 })
8153 .await,
8154 );
8155 rejected(
8156 "create",
8157 namespace
8158 .create_table_version(CreateTableVersionRequest {
8159 id: Some(vec!["users".to_string()]),
8160 version: 2,
8161 manifest_path: zombie.to_string(),
8162 branch: Some("ghost".to_string()),
8163 ..Default::default()
8164 })
8165 .await,
8166 );
8167 rejected(
8168 "restore",
8169 namespace
8170 .restore_table(RestoreTableRequest {
8171 id: Some(vec!["users".to_string()]),
8172 version: 1,
8173 branch: Some("ghost".to_string()),
8174 ..Default::default()
8175 })
8176 .await,
8177 );
8178 rejected(
8179 "batch_delete",
8180 namespace
8181 .batch_delete_table_versions(BatchDeleteTableVersionsRequest {
8182 id: Some(vec!["users".to_string()]),
8183 branch: Some("ghost".to_string()),
8184 ranges: vec![VersionRange::new(1, 1)],
8185 ..Default::default()
8186 })
8187 .await,
8188 );
8189 }
8190
8191 #[tokio::test]
8196 async fn test_batch_delete_table_versions_main_v2() {
8197 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
8198
8199 let (namespace, _temp_dir) = create_test_namespace().await;
8200 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
8202 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
8206 assert!(before.len() >= 3);
8207 assert!(
8210 before
8211 .iter()
8212 .all(|v| v.manifest_path.rsplit('/').next().unwrap().len() == 29),
8213 "expected V2-named manifests: {:?}",
8214 before
8215 );
8216 let min_v = before.iter().map(|v| v.version).min().unwrap();
8217 let max_v = before.iter().map(|v| v.version).max().unwrap();
8218
8219 let req = BatchDeleteTableVersionsRequest {
8222 id: Some(vec!["users".to_string()]),
8223 ranges: vec![VersionRange::new(min_v, max_v)],
8224 ..Default::default()
8225 };
8226 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
8227 assert_eq!(
8228 resp.deleted_count,
8229 Some((before.len() - 1) as i64),
8230 "V2 manifests must actually be deleted (was 0 before the fix)"
8231 );
8232
8233 let after = list_versions(&namespace, "users", None).await.unwrap();
8234 assert_eq!(after.len(), 1);
8235 assert_eq!(after[0].version, max_v);
8236 }
8237
8238 #[tokio::test]
8240 async fn test_batch_delete_end_is_exclusive() {
8241 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
8242
8243 let (namespace, _temp_dir) = create_test_namespace().await;
8244 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
8246 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
8250 let min_v = before.iter().map(|v| v.version).min().unwrap();
8251
8252 let req = BatchDeleteTableVersionsRequest {
8253 id: Some(vec!["users".to_string()]),
8254 ranges: vec![VersionRange::new(min_v, min_v + 1)],
8255 ..Default::default()
8256 };
8257 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
8258 assert_eq!(
8259 resp.deleted_count,
8260 Some(1),
8261 "only min_v is in [min_v, min_v+1)"
8262 );
8263
8264 let after = list_versions(&namespace, "users", None).await.unwrap();
8265 assert!(
8266 !after.iter().any(|v| v.version == min_v),
8267 "min_v must be deleted"
8268 );
8269 assert_eq!(after.len(), before.len() - 1, "exactly one version removed");
8270 }
8271
8272 #[tokio::test]
8273 async fn test_batch_delete_rejects_unbounded_range() {
8274 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
8275
8276 let (namespace, _temp_dir) = create_test_namespace().await;
8277 create_scalar_table(&namespace, "users").await;
8278
8279 let req = BatchDeleteTableVersionsRequest {
8282 id: Some(vec!["users".to_string()]),
8283 ranges: vec![VersionRange::new(0, i64::MAX)],
8284 ..Default::default()
8285 };
8286 let err = namespace.batch_delete_table_versions(req).await;
8287 assert!(err.is_err());
8288 assert!(
8289 err.unwrap_err().to_string().contains("limit"),
8290 "expected a range-too-large error"
8291 );
8292 }
8293
8294 async fn create_managed_namespace(path: &str) -> Arc<dyn LanceNamespace> {
8296 Arc::new(
8297 DirectoryNamespaceBuilder::new(path)
8298 .manifest_enabled(true)
8299 .table_version_tracking_enabled(true)
8300 .build()
8301 .await
8302 .unwrap(),
8303 )
8304 }
8305
8306 fn single_int_schema() -> Arc<arrow::datatypes::Schema> {
8307 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
8308 Arc::new(ArrowSchema::new(vec![Field::new(
8309 "id",
8310 DataType::Int32,
8311 false,
8312 )]))
8313 }
8314
8315 fn single_int_batch(seed: i32) -> arrow::record_batch::RecordBatch {
8316 use arrow::array::Int32Array;
8317 arrow::record_batch::RecordBatch::try_new(
8318 single_int_schema(),
8319 vec![Arc::new(Int32Array::from(vec![seed]))],
8320 )
8321 .unwrap()
8322 }
8323
8324 async fn create_managed_table(ns: &Arc<dyn LanceNamespace>, table_id: &[String]) -> Dataset {
8327 let mut ds = Dataset::write_into_namespace(
8328 RecordBatchIterator::new(vec![Ok(single_int_batch(1))], single_int_schema()),
8329 ns.clone(),
8330 table_id.to_vec(),
8331 Some(WriteParams {
8332 mode: WriteMode::Create,
8333 ..Default::default()
8334 }),
8335 )
8336 .await
8337 .unwrap();
8338 ds.append(
8339 RecordBatchIterator::new(vec![Ok(single_int_batch(2))], single_int_schema()),
8340 None,
8341 )
8342 .await
8343 .unwrap();
8344 ds
8345 }
8346
8347 async fn scan_id_column(ds: &Dataset) -> Vec<i32> {
8349 use arrow::array::Int32Array;
8350 use futures::TryStreamExt;
8351 let batches: Vec<arrow::record_batch::RecordBatch> = ds
8352 .scan()
8353 .try_into_stream()
8354 .await
8355 .unwrap()
8356 .try_collect()
8357 .await
8358 .unwrap();
8359 let mut ids: Vec<i32> = batches
8360 .iter()
8361 .flat_map(|b| {
8362 b.column(0)
8363 .as_any()
8364 .downcast_ref::<Int32Array>()
8365 .unwrap()
8366 .values()
8367 .to_vec()
8368 })
8369 .collect();
8370 ids.sort();
8371 ids
8372 }
8373
8374 #[tokio::test]
8379 async fn test_managed_branch_open_and_commit() {
8380 use futures::TryStreamExt;
8381 use lance::dataset::builder::DatasetBuilder;
8382 use lance_namespace::models::CreateTableBranchRequest;
8383
8384 let temp = TempStdDir::default();
8385 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8386 let table_id = vec!["t".to_string()];
8387 create_managed_table(&ns, &table_id).await;
8388 let main_before = ns
8389 .list_table_versions(ListTableVersionsRequest {
8390 id: Some(table_id.clone()),
8391 ..Default::default()
8392 })
8393 .await
8394 .unwrap()
8395 .versions
8396 .len();
8397
8398 ns.create_table_branch(CreateTableBranchRequest {
8401 id: Some(table_id.clone()),
8402 name: "exp".to_string(),
8403 ..Default::default()
8404 })
8405 .await
8406 .unwrap();
8407
8408 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8411 .await
8412 .unwrap()
8413 .with_branch("exp", None)
8414 .load()
8415 .await
8416 .unwrap();
8417 assert_eq!(
8418 branch_ds.manifest.branch.as_deref(),
8419 Some("exp"),
8420 "with_branch on a managed table must open the branch chain"
8421 );
8422 let branch_base = branch_ds.branch_location().path;
8423 assert!(
8424 branch_base.as_ref().ends_with("tree/exp"),
8425 "the branch dataset must be rooted at the branch chain: {}",
8426 branch_base
8427 );
8428 let branch_v_before = branch_ds.version().version;
8429
8430 branch_ds
8432 .append(
8433 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8434 None,
8435 )
8436 .await
8437 .unwrap();
8438 assert_eq!(
8439 branch_ds.manifest.branch.as_deref(),
8440 Some("exp"),
8441 "the commit must stay on the branch"
8442 );
8443 assert!(
8444 branch_ds.version().version > branch_v_before,
8445 "the branch version must advance after the commit"
8446 );
8447 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
8448
8449 let store = branch_ds.object_store(None).await.unwrap();
8453 let branch_data = branch_base.clone().join("data");
8454 let branch_files = store
8455 .inner
8456 .list(Some(&branch_data))
8457 .try_collect::<Vec<_>>()
8458 .await
8459 .unwrap();
8460 assert!(
8461 !branch_files.is_empty(),
8462 "the branch commit must place data files under the branch chain"
8463 );
8464
8465 let table_uri = ns
8467 .describe_table(DescribeTableRequest {
8468 id: Some(table_id.clone()),
8469 ..Default::default()
8470 })
8471 .await
8472 .unwrap()
8473 .location
8474 .unwrap();
8475 let fs_branch_ds = DatasetBuilder::from_uri(&table_uri)
8476 .with_branch("exp", None)
8477 .load()
8478 .await
8479 .unwrap();
8480 assert_eq!(fs_branch_ds.manifest.branch.as_deref(), Some("exp"));
8481 assert_eq!(scan_id_column(&fs_branch_ds).await, vec![1, 2, 3]);
8482
8483 let main_after = ns
8486 .list_table_versions(ListTableVersionsRequest {
8487 id: Some(table_id.clone()),
8488 ..Default::default()
8489 })
8490 .await
8491 .unwrap()
8492 .versions
8493 .len();
8494 assert_eq!(
8495 main_after, main_before,
8496 "committing on the branch must not change main's chain"
8497 );
8498 let main_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8499 .await
8500 .unwrap()
8501 .load()
8502 .await
8503 .unwrap();
8504 assert_eq!(main_ds.manifest.branch, None);
8505 assert_eq!(scan_id_column(&main_ds).await, vec![1, 2]);
8506 }
8507
8508 #[tokio::test]
8513 async fn test_managed_branch_tags() {
8514 use lance::dataset::builder::DatasetBuilder;
8515 use lance::dataset::refs::Ref;
8516 use lance_namespace::models::CreateTableBranchRequest;
8517
8518 let temp = TempStdDir::default();
8519 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8520 let table_id = vec!["t".to_string()];
8521 let main_ds = create_managed_table(&ns, &table_id).await;
8522 ns.create_table_branch(CreateTableBranchRequest {
8523 id: Some(table_id.clone()),
8524 name: "exp".to_string(),
8525 ..Default::default()
8526 })
8527 .await
8528 .unwrap();
8529 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8530 .await
8531 .unwrap()
8532 .with_branch("exp", None)
8533 .load()
8534 .await
8535 .unwrap();
8536 branch_ds
8537 .append(
8538 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8539 None,
8540 )
8541 .await
8542 .unwrap();
8543 let branch_version = branch_ds.version().version;
8544
8545 main_ds
8548 .tags()
8549 .create("exp-tag", ("exp", Some(branch_version)))
8550 .await
8551 .unwrap();
8552 let tag = main_ds.tags().get("exp-tag").await.unwrap();
8553 assert_eq!(tag.branch.as_deref(), Some("exp"));
8554 assert_eq!(tag.version, branch_version);
8555
8556 branch_ds
8558 .tags()
8559 .create("exp-tag2", branch_version)
8560 .await
8561 .unwrap();
8562 let tag2 = branch_ds.tags().get("exp-tag2").await.unwrap();
8563 assert_eq!(tag2.branch.as_deref(), Some("exp"));
8564
8565 let tag_open = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8568 .await
8569 .unwrap()
8570 .with_tag("exp-tag")
8571 .load()
8572 .await
8573 .unwrap();
8574 assert_eq!(tag_open.manifest.branch.as_deref(), Some("exp"));
8575 assert_eq!(tag_open.version().version, branch_version);
8576 assert_eq!(scan_id_column(&tag_open).await, vec![1, 2, 3]);
8577
8578 let tag_checkout = main_ds
8580 .checkout_version(Ref::Tag("exp-tag".to_string()))
8581 .await
8582 .unwrap();
8583 assert_eq!(tag_checkout.manifest.branch.as_deref(), Some("exp"));
8584 assert_eq!(scan_id_column(&tag_checkout).await, vec![1, 2, 3]);
8585
8586 let err = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8588 .await
8589 .unwrap()
8590 .with_tag("no-such-tag")
8591 .load()
8592 .await;
8593 assert!(err.is_err(), "a missing tag must error");
8594 }
8595
8596 #[tokio::test]
8601 async fn test_managed_cross_branch_checkout() {
8602 use lance::dataset::builder::DatasetBuilder;
8603 use lance::dataset::refs::Ref;
8604 use lance_namespace::models::CreateTableBranchRequest;
8605
8606 let temp = TempStdDir::default();
8607 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8608 let table_id = vec!["t".to_string()];
8609 let mut main_ds = create_managed_table(&ns, &table_id).await;
8610 ns.create_table_branch(CreateTableBranchRequest {
8611 id: Some(table_id.clone()),
8612 name: "exp".to_string(),
8613 ..Default::default()
8614 })
8615 .await
8616 .unwrap();
8617
8618 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8620 .await
8621 .unwrap()
8622 .with_branch("exp", None)
8623 .load()
8624 .await
8625 .unwrap();
8626 branch_ds
8627 .append(
8628 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8629 None,
8630 )
8631 .await
8632 .unwrap();
8633 let overlap_version = branch_ds.version().version;
8634 while main_ds.version().version < overlap_version {
8635 main_ds
8636 .append(
8637 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
8638 None,
8639 )
8640 .await
8641 .unwrap();
8642 }
8643
8644 let on_branch = main_ds
8647 .checkout_version(Ref::Version(Some("exp".to_string()), Some(overlap_version)))
8648 .await
8649 .unwrap();
8650 assert_eq!(on_branch.manifest.branch.as_deref(), Some("exp"));
8651 assert_eq!(scan_id_column(&on_branch).await, vec![1, 2, 3]);
8652
8653 let mut on_branch_latest = main_ds.checkout_branch("exp").await.unwrap();
8655 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8656 assert_eq!(on_branch_latest.version().version, overlap_version);
8657
8658 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
8661 ns.list_table_versions(ListTableVersionsRequest {
8662 id: Some(table_id),
8663 ..Default::default()
8664 })
8665 .await
8666 .unwrap()
8667 .versions
8668 .len()
8669 };
8670 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
8671 on_branch_latest
8672 .append(
8673 RecordBatchIterator::new(vec![Ok(single_int_batch(4))], single_int_schema()),
8674 None,
8675 )
8676 .await
8677 .unwrap();
8678 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8679 assert_eq!(scan_id_column(&on_branch_latest).await, vec![1, 2, 3, 4]);
8680 assert_eq!(
8681 main_chain_len(ns.clone(), table_id.clone()).await,
8682 main_before,
8683 "a commit on the checked-out branch must not advance main's chain"
8684 );
8685
8686 let on_main = branch_ds
8688 .checkout_version(Ref::Version(None, Some(1)))
8689 .await
8690 .unwrap();
8691 assert_eq!(on_main.manifest.branch, None);
8692 assert_eq!(scan_id_column(&on_main).await, vec![1]);
8693
8694 ns.create_table_branch(CreateTableBranchRequest {
8696 id: Some(table_id.clone()),
8697 name: "exp2".to_string(),
8698 ..Default::default()
8699 })
8700 .await
8701 .unwrap();
8702 let on_branch2 = branch_ds.checkout_branch("exp2").await.unwrap();
8703 assert_eq!(on_branch2.manifest.branch.as_deref(), Some("exp2"));
8704
8705 let err = main_ds
8707 .checkout_version(Ref::Version(Some("exp".to_string()), Some(999)))
8708 .await;
8709 assert!(err.is_err(), "a version missing from the branch must error");
8710 }
8711
8712 #[tokio::test]
8718 async fn test_commit_builder_honors_explicit_handler_for_dataset_dest() {
8719 use lance::dataset::write::{CommitBuilder, InsertBuilder};
8720 use lance::dataset::{WriteDestination, builder::DatasetBuilder};
8721 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
8722 use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
8723
8724 let temp = TempStdDir::default();
8725 let namespace = Arc::new(
8726 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
8727 .manifest_enabled(true)
8728 .table_version_tracking_enabled(true)
8729 .ops_metrics_enabled(true)
8730 .build()
8731 .await
8732 .unwrap(),
8733 );
8734 let ns: Arc<dyn LanceNamespace> = namespace.clone();
8735 let table_id = vec!["t".to_string()];
8736 create_managed_table(&ns, &table_id).await; let table_uri = ns
8740 .describe_table(DescribeTableRequest {
8741 id: Some(table_id.clone()),
8742 ..Default::default()
8743 })
8744 .await
8745 .unwrap()
8746 .location
8747 .unwrap();
8748 let plain_ds = Arc::new(Dataset::open(&table_uri).await.unwrap());
8749
8750 let transaction = InsertBuilder::new(WriteDestination::Dataset(plain_ds.clone()))
8751 .with_params(&WriteParams {
8752 mode: WriteMode::Append,
8753 ..Default::default()
8754 })
8755 .execute_uncommitted(vec![single_int_batch(3)])
8756 .await
8757 .unwrap();
8758
8759 let handler = Arc::new(ExternalManifestCommitHandler {
8760 external_manifest_store: Arc::new(
8761 LanceNamespaceExternalManifestStore::for_table_uri(
8762 ns.clone(),
8763 table_id.clone(),
8764 &table_uri,
8765 )
8766 .unwrap(),
8767 ),
8768 });
8769 let commits_before = namespace
8770 .retrieve_ops_metrics()
8771 .get("create_table_version")
8772 .copied()
8773 .unwrap_or(0);
8774 let committed = CommitBuilder::new(WriteDestination::Dataset(plain_ds))
8775 .with_commit_handler(handler)
8776 .execute(transaction)
8777 .await
8778 .unwrap();
8779 assert_eq!(scan_id_column(&committed).await, vec![1, 2, 3]);
8780
8781 let commits_after = namespace
8782 .retrieve_ops_metrics()
8783 .get("create_table_version")
8784 .copied()
8785 .unwrap_or(0);
8786 assert_eq!(
8787 commits_after,
8788 commits_before + 1,
8789 "the explicit handler must route the commit through create_table_version"
8790 );
8791 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8792 .await
8793 .unwrap()
8794 .load()
8795 .await
8796 .unwrap();
8797 assert_eq!(
8798 scan_id_column(&fresh).await,
8799 vec![1, 2, 3],
8800 "a fresh managed open must resolve the committed version"
8801 );
8802 }
8803
8804 #[tokio::test]
8806 async fn test_managed_branch_from_non_latest_fork() {
8807 use lance::dataset::builder::DatasetBuilder;
8808 use lance_namespace::models::CreateTableBranchRequest;
8809
8810 let temp = TempStdDir::default();
8811 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8812 let table_id = vec!["t".to_string()];
8813 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
8816 id: Some(table_id.clone()),
8817 name: "old".to_string(),
8818 from_version: Some(1),
8819 ..Default::default()
8820 })
8821 .await
8822 .unwrap();
8823
8824 let old_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8825 .await
8826 .unwrap()
8827 .with_branch("old", None)
8828 .load()
8829 .await
8830 .unwrap();
8831 assert_eq!(old_ds.manifest.branch.as_deref(), Some("old"));
8832 assert_eq!(
8833 scan_id_column(&old_ds).await,
8834 vec![1],
8835 "the fork must contain only the fork-point data"
8836 );
8837 }
8838
8839 #[test]
8842 fn test_manifest_version_from_filename() {
8843 assert_eq!(
8845 DirectoryNamespace::manifest_version_from_filename("5.manifest"),
8846 Some(5)
8847 );
8848 assert_eq!(
8849 DirectoryNamespace::manifest_version_from_filename("0.manifest"),
8850 Some(0)
8851 );
8852 let v2_five = format!("{:020}.manifest", u64::MAX - 5);
8854 assert_eq!(
8855 DirectoryNamespace::manifest_version_from_filename(&v2_five),
8856 Some(5)
8857 );
8858 let v2_zero = format!("{:020}.manifest", u64::MAX);
8859 assert_eq!(
8860 DirectoryNamespace::manifest_version_from_filename(&v2_zero),
8861 Some(0)
8862 );
8863 assert_eq!(
8865 DirectoryNamespace::manifest_version_from_filename("data.lance"),
8866 None
8867 );
8868 assert_eq!(
8869 DirectoryNamespace::manifest_version_from_filename("d5.manifest"),
8870 None
8871 );
8872 }
8873
8874 #[tokio::test]
8875 async fn test_create_table() {
8876 let (namespace, _temp_dir) = create_test_namespace().await;
8877
8878 let schema = create_test_schema();
8880 let ipc_data = create_test_ipc_data(&schema);
8881
8882 let mut request = CreateTableRequest::new();
8883 request.id = Some(vec!["test_table".to_string()]);
8884
8885 let response = namespace
8886 .create_table(request, bytes::Bytes::from(ipc_data))
8887 .await
8888 .unwrap();
8889
8890 assert!(response.location.is_some());
8891 assert!(response.location.unwrap().ends_with("test_table.lance"));
8892 assert_eq!(response.version, Some(1));
8893 }
8894
8895 #[tokio::test]
8896 async fn test_create_table_without_data() {
8897 let (namespace, _temp_dir) = create_test_namespace().await;
8898
8899 let mut request = CreateTableRequest::new();
8900 request.id = Some(vec!["test_table".to_string()]);
8901
8902 let result = namespace.create_table(request, bytes::Bytes::new()).await;
8903 assert!(result.is_err());
8904 assert!(
8905 result
8906 .unwrap_err()
8907 .to_string()
8908 .contains("Arrow IPC stream) is required")
8909 );
8910 }
8911
8912 #[tokio::test]
8913 async fn test_create_table_with_invalid_id() {
8914 let (namespace, _temp_dir) = create_test_namespace().await;
8915
8916 let schema = create_test_schema();
8918 let ipc_data = create_test_ipc_data(&schema);
8919
8920 let mut request = CreateTableRequest::new();
8922 request.id = Some(vec![]);
8923
8924 let result = namespace
8925 .create_table(request, bytes::Bytes::from(ipc_data.clone()))
8926 .await;
8927 assert!(result.is_err());
8928
8929 let mut create_ns_req = CreateNamespaceRequest::new();
8932 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
8933 namespace.create_namespace(create_ns_req).await.unwrap();
8934
8935 let mut request = CreateTableRequest::new();
8937 request.id = Some(vec!["test_namespace".to_string(), "table".to_string()]);
8938
8939 let result = namespace
8940 .create_table(request, bytes::Bytes::from(ipc_data))
8941 .await;
8942 assert!(
8944 result.is_ok(),
8945 "Multi-level table IDs should work with manifest enabled"
8946 );
8947 }
8948
8949 #[tokio::test]
8950 async fn test_list_tables() {
8951 let (namespace, _temp_dir) = create_test_namespace().await;
8952
8953 let mut request = ListTablesRequest::new();
8955 request.id = Some(vec![]);
8956 let response = namespace.list_tables(request).await.unwrap();
8957 assert_eq!(response.tables.len(), 0);
8958
8959 let schema = create_test_schema();
8961 let ipc_data = create_test_ipc_data(&schema);
8962
8963 let mut create_request = CreateTableRequest::new();
8965 create_request.id = Some(vec!["table1".to_string()]);
8966 namespace
8967 .create_table(create_request, bytes::Bytes::from(ipc_data.clone()))
8968 .await
8969 .unwrap();
8970
8971 let mut create_request = CreateTableRequest::new();
8973 create_request.id = Some(vec!["table2".to_string()]);
8974 namespace
8975 .create_table(create_request, bytes::Bytes::from(ipc_data))
8976 .await
8977 .unwrap();
8978
8979 let mut request = ListTablesRequest::new();
8981 request.id = Some(vec![]);
8982 let response = namespace.list_tables(request).await.unwrap();
8983 let tables = response.tables;
8984 assert_eq!(tables.len(), 2);
8985 assert!(tables.contains(&"table1".to_string()));
8986 assert!(tables.contains(&"table2".to_string()));
8987 }
8988
8989 #[tokio::test]
8990 async fn test_list_tables_pagination() {
8991 let (namespace, _temp_dir) = create_test_namespace().await;
8992
8993 let schema = create_test_schema();
8994 let ipc_data = create_test_ipc_data(&schema);
8995
8996 for name in ["alpha", "bravo", "charlie"] {
8997 let mut req = CreateTableRequest::new();
8998 req.id = Some(vec![name.to_string()]);
8999 namespace
9000 .create_table(req, bytes::Bytes::from(ipc_data.clone()))
9001 .await
9002 .unwrap();
9003 }
9004
9005 let first_page = namespace
9007 .list_tables(ListTablesRequest {
9008 id: Some(vec![]),
9009 limit: Some(2),
9010 ..Default::default()
9011 })
9012 .await
9013 .unwrap();
9014
9015 assert_eq!(first_page.tables, vec!["alpha", "bravo"]);
9016 assert_eq!(first_page.page_token.as_deref(), Some("bravo"));
9017
9018 let second_page = namespace
9020 .list_tables(ListTablesRequest {
9021 id: Some(vec![]),
9022 limit: Some(2),
9023 page_token: first_page.page_token.clone(),
9024 ..Default::default()
9025 })
9026 .await
9027 .unwrap();
9028
9029 assert_eq!(second_page.tables, vec!["charlie"]);
9030 assert!(second_page.page_token.is_none());
9031 }
9032
9033 #[tokio::test]
9034 async fn test_list_tables_pagination_limit_zero() {
9035 let (namespace, _temp_dir) = create_test_namespace().await;
9036
9037 let schema = create_test_schema();
9038 let ipc_data = create_test_ipc_data(&schema);
9039
9040 let mut req = CreateTableRequest::new();
9041 req.id = Some(vec!["alpha".to_string()]);
9042 namespace
9043 .create_table(req, bytes::Bytes::from(ipc_data))
9044 .await
9045 .unwrap();
9046
9047 let response = namespace
9048 .list_tables(ListTablesRequest {
9049 id: Some(vec![]),
9050 limit: Some(0),
9051 ..Default::default()
9052 })
9053 .await
9054 .unwrap();
9055
9056 assert!(response.tables.is_empty());
9057 assert!(response.page_token.is_none());
9058 }
9059
9060 #[tokio::test]
9061 async fn test_list_tables_with_namespace_id() {
9062 let (namespace, _temp_dir) = create_test_namespace().await;
9063
9064 let mut create_ns_req = CreateNamespaceRequest::new();
9066 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
9067 namespace.create_namespace(create_ns_req).await.unwrap();
9068
9069 let mut request = ListTablesRequest::new();
9071 request.id = Some(vec!["test_namespace".to_string()]);
9072
9073 let result = namespace.list_tables(request).await;
9074 assert!(
9076 result.is_ok(),
9077 "list_tables should work with child namespace when manifest is enabled"
9078 );
9079 let response = result.unwrap();
9080 assert_eq!(
9081 response.tables.len(),
9082 0,
9083 "Namespace should have no tables yet"
9084 );
9085 }
9086
9087 #[tokio::test]
9088 async fn test_create_scalar_index() {
9089 let (namespace, _temp_dir) = create_test_namespace().await;
9090 create_scalar_table(&namespace, "users").await;
9091
9092 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9093 let dataset = open_dataset(&namespace, "users").await;
9094 let expected_transaction_id = dataset
9095 .read_transaction()
9096 .await
9097 .unwrap()
9098 .map(|transaction| transaction.uuid);
9099 assert_eq!(transaction_id, expected_transaction_id);
9100 let indices = dataset.load_indices().await.unwrap();
9101 assert!(indices.iter().any(|index| index.name == "users_id_idx"));
9102 }
9103
9104 #[tokio::test]
9105 async fn test_create_vector_index() {
9106 use lance_namespace::models::CreateTableIndexRequest;
9107
9108 let (namespace, _temp_dir) = create_test_namespace().await;
9109 create_vector_table(&namespace, "vectors").await;
9110
9111 let mut create_index_request =
9112 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
9113 create_index_request.id = Some(vec!["vectors".to_string()]);
9114 create_index_request.name = Some("vector_idx".to_string());
9115 create_index_request.distance_type = Some("l2".to_string());
9116 let transaction_id = namespace
9117 .create_table_index(create_index_request)
9118 .await
9119 .unwrap()
9120 .transaction_id;
9121
9122 let dataset = open_dataset(&namespace, "vectors").await;
9123 let expected_transaction_id = dataset
9124 .read_transaction()
9125 .await
9126 .unwrap()
9127 .map(|transaction| transaction.uuid);
9128 assert_eq!(transaction_id, expected_transaction_id);
9129 let indices = dataset.load_indices().await.unwrap();
9130 assert!(indices.iter().any(|index| index.name == "vector_idx"));
9131 }
9132
9133 #[tokio::test]
9134 async fn test_list_table_indices() {
9135 use lance_namespace::models::{CreateTableIndexRequest, ListTableIndicesRequest};
9136
9137 let (namespace, _temp_dir) = create_test_namespace().await;
9138 create_scalar_table(&namespace, "users").await;
9139 create_scalar_index(&namespace, "users", "a_idx").await;
9140 create_scalar_index(&namespace, "users", "b_idx").await;
9141 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9142
9143 let response = namespace
9144 .list_table_indices(ListTableIndicesRequest {
9145 id: Some(vec!["users".to_string()]),
9146 ..Default::default()
9147 })
9148 .await
9149 .unwrap();
9150
9151 assert_eq!(response.indexes.len(), 3);
9152 assert_eq!(response.indexes[0].index_name, "a_idx");
9153 assert_eq!(response.indexes[1].index_name, "b_idx");
9154 assert_eq!(response.indexes[2].index_name, "users_id_idx");
9155 assert!(response.page_token.is_none());
9156 let users_id_idx = response
9157 .indexes
9158 .iter()
9159 .find(|index| index.index_name == "users_id_idx")
9160 .unwrap();
9161 assert_eq!(users_id_idx.columns, vec!["id"]);
9162 assert_eq!(users_id_idx.status, "SUCCEEDED");
9163
9164 assert_eq!(users_id_idx.index_type.as_deref(), Some("BTree"));
9166 assert!(
9167 users_id_idx
9168 .type_url
9169 .as_deref()
9170 .is_some_and(|s| !s.is_empty())
9171 );
9172 assert_eq!(users_id_idx.num_indexed_rows, Some(3));
9173 assert_eq!(users_id_idx.num_unindexed_rows, Some(0));
9174 assert_eq!(users_id_idx.num_segments, Some(1));
9175 assert!(users_id_idx.size_bytes.is_some_and(|size| size > 0));
9176 assert!(users_id_idx.created_at.is_some());
9177 assert!(users_id_idx.index_version.is_some());
9178 assert!(users_id_idx.index_details.is_some());
9179
9180 let dataset = open_dataset(&namespace, "users").await;
9181 let expected_transaction_id = dataset
9182 .read_transaction()
9183 .await
9184 .unwrap()
9185 .map(|transaction| transaction.uuid);
9186 assert_eq!(transaction_id, expected_transaction_id);
9187 let indices = dataset.load_indices().await.unwrap();
9188 assert_eq!(
9189 indices
9190 .iter()
9191 .filter(|index| index.name == "users_id_idx")
9192 .count(),
9193 1
9194 );
9195
9196 let first_page = namespace
9197 .list_table_indices(ListTableIndicesRequest {
9198 id: Some(vec!["users".to_string()]),
9199 limit: Some(2),
9200 ..Default::default()
9201 })
9202 .await
9203 .unwrap();
9204
9205 assert_eq!(first_page.indexes.len(), 2);
9206 assert_eq!(first_page.indexes[0].index_name, "a_idx");
9207 assert_eq!(first_page.indexes[1].index_name, "b_idx");
9208 assert_eq!(first_page.page_token.as_deref(), Some("b_idx"));
9209
9210 let second_page = namespace
9211 .list_table_indices(ListTableIndicesRequest {
9212 id: Some(vec!["users".to_string()]),
9213 page_token: first_page.page_token.clone(),
9214 limit: Some(2),
9215 ..Default::default()
9216 })
9217 .await
9218 .unwrap();
9219
9220 assert_eq!(second_page.indexes.len(), 1);
9221 assert_eq!(second_page.indexes[0].index_name, "users_id_idx");
9222 assert!(second_page.page_token.is_none());
9223
9224 create_vector_table(&namespace, "vectors").await;
9226 let mut create_index_request =
9227 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
9228 create_index_request.id = Some(vec!["vectors".to_string()]);
9229 create_index_request.name = Some("vector_idx".to_string());
9230 create_index_request.distance_type = Some("l2".to_string());
9231 namespace
9232 .create_table_index(create_index_request)
9233 .await
9234 .unwrap();
9235
9236 let vector_response = namespace
9237 .list_table_indices(ListTableIndicesRequest {
9238 id: Some(vec!["vectors".to_string()]),
9239 ..Default::default()
9240 })
9241 .await
9242 .unwrap();
9243
9244 assert_eq!(vector_response.indexes.len(), 1);
9245 let vector_idx = &vector_response.indexes[0];
9246 assert_eq!(vector_idx.index_name, "vector_idx");
9247 assert_eq!(vector_idx.columns, vec!["vector"]);
9248 assert_eq!(vector_idx.index_type.as_deref(), Some("IVF_FLAT"));
9249 assert!(
9250 vector_idx
9251 .type_url
9252 .as_deref()
9253 .is_some_and(|s| !s.is_empty())
9254 );
9255 assert!(vector_idx.num_indexed_rows.is_some());
9256 assert!(vector_idx.num_unindexed_rows.is_some());
9257 assert_eq!(vector_idx.num_segments, Some(1));
9258 assert!(vector_idx.created_at.is_some());
9259 assert!(vector_idx.index_version.is_some());
9260 assert!(vector_idx.index_details.is_some());
9261 }
9262
9263 #[tokio::test]
9264 async fn test_describe_table_index_stats() {
9265 use lance_namespace::models::DescribeTableIndexStatsRequest;
9266
9267 let (namespace, _temp_dir) = create_test_namespace().await;
9268 create_scalar_table(&namespace, "users").await;
9269 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9270
9271 let response = namespace
9272 .describe_table_index_stats(DescribeTableIndexStatsRequest {
9273 id: Some(vec!["users".to_string()]),
9274 index_name: Some("users_id_idx".to_string()),
9275 ..Default::default()
9276 })
9277 .await
9278 .unwrap();
9279 assert_eq!(response.index_type, Some("BTree".to_string()));
9280 assert_eq!(response.num_indices, Some(1));
9281 assert_eq!(response.num_indexed_rows, Some(3));
9282 assert_eq!(response.num_unindexed_rows, Some(0));
9283
9284 let dataset = open_dataset(&namespace, "users").await;
9285 let expected_transaction_id = dataset
9286 .read_transaction()
9287 .await
9288 .unwrap()
9289 .map(|transaction| transaction.uuid);
9290 assert_eq!(transaction_id, expected_transaction_id);
9291 let stats: serde_json::Value =
9292 serde_json::from_str(&dataset.index_statistics("users_id_idx").await.unwrap()).unwrap();
9293 assert_eq!(stats["index_type"], "BTree");
9294 assert_eq!(stats["num_indices"], 1);
9295 assert_eq!(stats["num_indexed_rows"], 3);
9296 assert_eq!(stats["num_unindexed_rows"], 0);
9297 }
9298
9299 #[tokio::test]
9300 async fn test_describe_transaction() {
9301 use lance_namespace::models::DescribeTransactionRequest;
9302
9303 let (namespace, _temp_dir) = create_test_namespace().await;
9304 create_scalar_table(&namespace, "users").await;
9305 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9306 let dataset = open_dataset(&namespace, "users").await;
9307 let latest_transaction = dataset.read_transaction().await.unwrap();
9308 assert_eq!(
9309 transaction_id,
9310 latest_transaction
9311 .as_ref()
9312 .map(|transaction| transaction.uuid.clone())
9313 );
9314
9315 if let Some(transaction_id) = transaction_id {
9316 let response = namespace
9317 .describe_transaction(DescribeTransactionRequest {
9318 id: Some(vec!["users".to_string(), transaction_id.clone()]),
9319 ..Default::default()
9320 })
9321 .await
9322 .unwrap();
9323 assert_eq!(response.status, "SUCCEEDED");
9324 assert_eq!(
9325 response
9326 .properties
9327 .as_ref()
9328 .and_then(|props| props.get("operation")),
9329 Some(&"CreateIndex".to_string())
9330 );
9331 assert_eq!(
9332 response
9333 .properties
9334 .as_ref()
9335 .and_then(|props| props.get("uuid")),
9336 Some(&transaction_id)
9337 );
9338 } else {
9339 assert!(latest_transaction.is_none());
9340 }
9341 }
9342
9343 #[tokio::test]
9344 async fn test_drop_table_index() {
9345 use lance_namespace::models::{DropTableIndexRequest, ListTableIndicesRequest};
9346
9347 let (namespace, _temp_dir) = create_test_namespace().await;
9348 create_scalar_table(&namespace, "users").await;
9349 let create_transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9350
9351 let drop_transaction_id = namespace
9352 .drop_table_index(DropTableIndexRequest {
9353 id: Some(vec!["users".to_string()]),
9354 index_name: Some("users_id_idx".to_string()),
9355 ..Default::default()
9356 })
9357 .await
9358 .unwrap()
9359 .transaction_id;
9360
9361 let dataset = open_dataset(&namespace, "users").await;
9362 let previous_dataset = dataset
9363 .checkout_version(dataset.version().version - 1)
9364 .await
9365 .unwrap();
9366 let previous_transaction_id = previous_dataset
9367 .read_transaction()
9368 .await
9369 .unwrap()
9370 .map(|transaction| transaction.uuid);
9371 assert_eq!(create_transaction_id, previous_transaction_id);
9372 let expected_drop_transaction_id = dataset
9373 .read_transaction()
9374 .await
9375 .unwrap()
9376 .map(|transaction| transaction.uuid);
9377 assert_eq!(drop_transaction_id, expected_drop_transaction_id);
9378 let indices = dataset.load_indices().await.unwrap();
9379 assert!(!indices.iter().any(|index| index.name == "users_id_idx"));
9380
9381 let list_response = namespace
9382 .list_table_indices(ListTableIndicesRequest {
9383 id: Some(vec!["users".to_string()]),
9384 ..Default::default()
9385 })
9386 .await
9387 .unwrap();
9388 assert!(list_response.indexes.is_empty());
9389 }
9390
9391 #[tokio::test]
9392 async fn test_describe_table() {
9393 let (namespace, _temp_dir) = create_test_namespace().await;
9394
9395 let schema = create_test_schema();
9397 let ipc_data = create_test_ipc_data(&schema);
9398
9399 let mut create_request = CreateTableRequest::new();
9400 create_request.id = Some(vec!["test_table".to_string()]);
9401 namespace
9402 .create_table(create_request, bytes::Bytes::from(ipc_data))
9403 .await
9404 .unwrap();
9405
9406 let mut request = DescribeTableRequest::new();
9408 request.id = Some(vec!["test_table".to_string()]);
9409 let response = namespace.describe_table(request).await.unwrap();
9410
9411 assert!(response.location.is_some());
9412 assert!(response.location.unwrap().ends_with("test_table.lance"));
9413 }
9414
9415 #[tokio::test]
9416 async fn test_describe_nonexistent_table() {
9417 let (namespace, _temp_dir) = create_test_namespace().await;
9418
9419 let mut request = DescribeTableRequest::new();
9420 request.id = Some(vec!["nonexistent".to_string()]);
9421
9422 let result = namespace.describe_table(request).await;
9423 assert!(result.is_err());
9424 assert!(result.unwrap_err().to_string().contains("Table not found"));
9425 }
9426
9427 #[tokio::test]
9428 async fn test_table_exists() {
9429 let (namespace, _temp_dir) = create_test_namespace().await;
9430
9431 let schema = create_test_schema();
9433 let ipc_data = create_test_ipc_data(&schema);
9434
9435 let mut create_request = CreateTableRequest::new();
9436 create_request.id = Some(vec!["existing_table".to_string()]);
9437 namespace
9438 .create_table(create_request, bytes::Bytes::from(ipc_data))
9439 .await
9440 .unwrap();
9441
9442 let mut request = TableExistsRequest::new();
9444 request.id = Some(vec!["existing_table".to_string()]);
9445 let result = namespace.table_exists(request).await;
9446 assert!(result.is_ok());
9447
9448 let mut request = TableExistsRequest::new();
9450 request.id = Some(vec!["nonexistent".to_string()]);
9451 let result = namespace.table_exists(request).await;
9452 assert!(result.is_err());
9453 assert!(result.unwrap_err().to_string().contains("Table not found"));
9454 }
9455
9456 #[tokio::test]
9457 async fn test_drop_table() {
9458 let (namespace, _temp_dir) = create_test_namespace().await;
9459
9460 let schema = create_test_schema();
9462 let ipc_data = create_test_ipc_data(&schema);
9463
9464 let mut create_request = CreateTableRequest::new();
9465 create_request.id = Some(vec!["table_to_drop".to_string()]);
9466 namespace
9467 .create_table(create_request, bytes::Bytes::from(ipc_data))
9468 .await
9469 .unwrap();
9470
9471 let mut exists_request = TableExistsRequest::new();
9473 exists_request.id = Some(vec!["table_to_drop".to_string()]);
9474 assert!(namespace.table_exists(exists_request.clone()).await.is_ok());
9475
9476 let mut drop_request = DropTableRequest::new();
9478 drop_request.id = Some(vec!["table_to_drop".to_string()]);
9479 let response = namespace.drop_table(drop_request).await.unwrap();
9480 assert!(response.location.is_some());
9481
9482 assert!(namespace.table_exists(exists_request).await.is_err());
9484 }
9485
9486 #[tokio::test]
9487 async fn test_drop_nonexistent_table() {
9488 let (namespace, _temp_dir) = create_test_namespace().await;
9489
9490 let mut request = DropTableRequest::new();
9491 request.id = Some(vec!["nonexistent".to_string()]);
9492
9493 let result = namespace.drop_table(request).await;
9495 let _ = result;
9498 }
9499
9500 #[tokio::test]
9501 async fn test_root_namespace_operations() {
9502 let (namespace, _temp_dir) = create_test_namespace().await;
9503
9504 let mut request = ListNamespacesRequest::new();
9506 request.id = Some(vec![]);
9507 let result = namespace.list_namespaces(request).await;
9508 assert!(result.is_ok());
9509 assert_eq!(result.unwrap().namespaces.len(), 0);
9510
9511 let mut request = DescribeNamespaceRequest::new();
9513 request.id = Some(vec![]);
9514 let result = namespace.describe_namespace(request).await;
9515 assert!(result.is_ok());
9516
9517 let mut request = NamespaceExistsRequest::new();
9519 request.id = Some(vec![]);
9520 let result = namespace.namespace_exists(request).await;
9521 assert!(result.is_ok());
9522
9523 let mut request = CreateNamespaceRequest::new();
9525 request.id = Some(vec![]);
9526 let result = namespace.create_namespace(request).await;
9527 assert!(result.is_err());
9528 assert!(result.unwrap_err().to_string().contains("already exists"));
9529
9530 let mut request = DropNamespaceRequest::new();
9532 request.id = Some(vec![]);
9533 let result = namespace.drop_namespace(request).await;
9534 assert!(result.is_err());
9535 assert!(
9536 result
9537 .unwrap_err()
9538 .to_string()
9539 .contains("cannot be dropped")
9540 );
9541 }
9542
9543 #[tokio::test]
9544 async fn test_non_root_namespace_operations() {
9545 let (namespace, _temp_dir) = create_test_namespace().await;
9546
9547 let mut request = CreateNamespaceRequest::new();
9550 request.id = Some(vec!["child".to_string()]);
9551 let result = namespace.create_namespace(request).await;
9552 assert!(
9553 result.is_ok(),
9554 "Child namespace creation should succeed with manifest enabled"
9555 );
9556
9557 let mut request = NamespaceExistsRequest::new();
9559 request.id = Some(vec!["child".to_string()]);
9560 let result = namespace.namespace_exists(request).await;
9561 assert!(
9562 result.is_ok(),
9563 "Child namespace should exist after creation"
9564 );
9565
9566 let mut request = DropNamespaceRequest::new();
9568 request.id = Some(vec!["child".to_string()]);
9569 let result = namespace.drop_namespace(request).await;
9570 assert!(
9571 result.is_ok(),
9572 "Child namespace drop should succeed with manifest enabled"
9573 );
9574
9575 let mut request = NamespaceExistsRequest::new();
9577 request.id = Some(vec!["child".to_string()]);
9578 let result = namespace.namespace_exists(request).await;
9579 assert!(
9580 result.is_err(),
9581 "Child namespace should not exist after drop"
9582 );
9583 }
9584
9585 #[tokio::test]
9586 async fn test_config_custom_root() {
9587 let temp_dir = TempStdDir::default();
9588 let custom_path = temp_dir.join("custom");
9589 std::fs::create_dir(&custom_path).unwrap();
9590
9591 let namespace = DirectoryNamespaceBuilder::new(custom_path.to_string_lossy().to_string())
9592 .build()
9593 .await
9594 .unwrap();
9595
9596 let schema = create_test_schema();
9598 let ipc_data = create_test_ipc_data(&schema);
9599
9600 let mut request = CreateTableRequest::new();
9602 request.id = Some(vec!["test_table".to_string()]);
9603
9604 let response = namespace
9605 .create_table(request, bytes::Bytes::from(ipc_data))
9606 .await
9607 .unwrap();
9608
9609 assert!(response.location.unwrap().contains("custom"));
9610 }
9611
9612 #[tokio::test]
9613 async fn test_config_storage_options() {
9614 let temp_dir = TempStdDir::default();
9615
9616 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9617 .storage_option("option1", "value1")
9618 .storage_option("option2", "value2")
9619 .build()
9620 .await
9621 .unwrap();
9622
9623 let schema = create_test_schema();
9625 let ipc_data = create_test_ipc_data(&schema);
9626
9627 let mut request = CreateTableRequest::new();
9629 request.id = Some(vec!["test_table".to_string()]);
9630
9631 let response = namespace
9632 .create_table(request, bytes::Bytes::from(ipc_data))
9633 .await
9634 .unwrap();
9635
9636 let storage_options = response.storage_options.unwrap();
9637 assert_eq!(storage_options.get("option1"), Some(&"value1".to_string()));
9638 assert_eq!(storage_options.get("option2"), Some(&"value2".to_string()));
9639 }
9640
9641 #[tokio::test]
9645 async fn test_no_storage_options_without_vendor() {
9646 use lance_namespace::models::DeclareTableRequest;
9647
9648 let temp_dir = TempStdDir::default();
9649
9650 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9652 .manifest_enabled(false)
9653 .storage_option("aws_access_key_id", "AKID")
9654 .storage_option("aws_secret_access_key", "SECRET")
9655 .storage_option("region", "us-east-1")
9656 .build()
9657 .await
9658 .unwrap();
9659
9660 let schema = create_test_schema();
9661 let ipc_data = create_test_ipc_data(&schema);
9662
9663 let mut create_req = CreateTableRequest::new();
9665 create_req.id = Some(vec!["t1".to_string()]);
9666 namespace
9667 .create_table(create_req, bytes::Bytes::from(ipc_data))
9668 .await
9669 .unwrap();
9670
9671 let mut desc_req = DescribeTableRequest::new();
9673 desc_req.id = Some(vec!["t1".to_string()]);
9674 let resp = namespace.describe_table(desc_req).await.unwrap();
9675 assert!(resp.storage_options.is_none());
9676
9677 let mut decl_req = DeclareTableRequest::new();
9679 decl_req.id = Some(vec!["t2".to_string()]);
9680 let resp = namespace.declare_table(decl_req).await.unwrap();
9681 assert!(resp.storage_options.is_none());
9682 }
9683
9684 #[tokio::test]
9686 async fn test_no_storage_options_without_vendor_manifest() {
9687 let temp_dir = TempStdDir::default();
9688
9689 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9690 .storage_option("aws_access_key_id", "AKID")
9691 .storage_option("aws_secret_access_key", "SECRET")
9692 .storage_option("region", "us-east-1")
9693 .build()
9694 .await
9695 .unwrap();
9696
9697 let schema = create_test_schema();
9698 let ipc_data = create_test_ipc_data(&schema);
9699
9700 let mut create_req = CreateTableRequest::new();
9701 create_req.id = Some(vec!["t1".to_string()]);
9702 namespace
9703 .create_table(create_req, bytes::Bytes::from(ipc_data))
9704 .await
9705 .unwrap();
9706
9707 let mut desc_req = DescribeTableRequest::new();
9709 desc_req.id = Some(vec!["t1".to_string()]);
9710 let resp = namespace.describe_table(desc_req).await.unwrap();
9711 assert!(resp.storage_options.is_none());
9712 }
9713
9714 #[tokio::test]
9715 async fn test_from_properties_manifest_enabled() {
9716 let temp_dir = TempStdDir::default();
9717
9718 let mut properties = HashMap::new();
9719 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9720 properties.insert("manifest_enabled".to_string(), "true".to_string());
9721 properties.insert("dir_listing_enabled".to_string(), "false".to_string());
9722
9723 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9724 assert!(builder.manifest_enabled);
9725 assert!(!builder.dir_listing_enabled);
9726
9727 let namespace = builder.build().await.unwrap();
9728
9729 let schema = create_test_schema();
9731 let ipc_data = create_test_ipc_data(&schema);
9732
9733 let mut request = CreateTableRequest::new();
9735 request.id = Some(vec!["test_table".to_string()]);
9736
9737 let response = namespace
9738 .create_table(request, bytes::Bytes::from(ipc_data))
9739 .await
9740 .unwrap();
9741
9742 assert!(response.location.is_some());
9743 }
9744
9745 #[tokio::test]
9746 async fn test_from_properties_dir_listing_enabled() {
9747 let temp_dir = TempStdDir::default();
9748
9749 let mut properties = HashMap::new();
9750 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9751 properties.insert("manifest_enabled".to_string(), "false".to_string());
9752 properties.insert("dir_listing_enabled".to_string(), "true".to_string());
9753
9754 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9755 assert!(!builder.manifest_enabled);
9756 assert!(builder.dir_listing_enabled);
9757
9758 let namespace = builder.build().await.unwrap();
9759
9760 let schema = create_test_schema();
9762 let ipc_data = create_test_ipc_data(&schema);
9763
9764 let mut request = CreateTableRequest::new();
9766 request.id = Some(vec!["test_table".to_string()]);
9767
9768 let response = namespace
9769 .create_table(request, bytes::Bytes::from(ipc_data))
9770 .await
9771 .unwrap();
9772
9773 assert!(response.location.is_some());
9774 }
9775
9776 #[tokio::test]
9777 async fn test_from_properties_defaults() {
9778 let temp_dir = TempStdDir::default();
9779
9780 let mut properties = HashMap::new();
9781 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9782
9783 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9784 assert!(builder.manifest_enabled);
9785 assert!(builder.dir_listing_enabled);
9786 assert!(!builder.inline_optimization_enabled);
9787 }
9788
9789 #[test]
9790 fn test_builder_disables_inline_optimization_by_default() {
9791 let builder = DirectoryNamespaceBuilder::new("memory://");
9792 assert!(!builder.inline_optimization_enabled);
9793 }
9794
9795 #[tokio::test]
9796 async fn test_from_properties_with_storage_options() {
9797 let temp_dir = TempStdDir::default();
9798
9799 let mut properties = HashMap::new();
9800 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9801 properties.insert("manifest_enabled".to_string(), "true".to_string());
9802 properties.insert("storage.region".to_string(), "us-west-2".to_string());
9803 properties.insert("storage.bucket".to_string(), "my-bucket".to_string());
9804
9805 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9806 assert!(builder.manifest_enabled);
9807 assert!(builder.storage_options.is_some());
9808
9809 let storage_options = builder.storage_options.unwrap();
9810 assert_eq!(
9811 storage_options.get("region"),
9812 Some(&"us-west-2".to_string())
9813 );
9814 assert_eq!(
9815 storage_options.get("bucket"),
9816 Some(&"my-bucket".to_string())
9817 );
9818 }
9819
9820 #[tokio::test]
9821 async fn test_various_arrow_types() {
9822 let (namespace, _temp_dir) = create_test_namespace().await;
9823
9824 let fields = vec![
9826 JsonArrowField {
9827 name: "bool_col".to_string(),
9828 r#type: Box::new(JsonArrowDataType::new("bool".to_string())),
9829 nullable: true,
9830 metadata: None,
9831 },
9832 JsonArrowField {
9833 name: "int8_col".to_string(),
9834 r#type: Box::new(JsonArrowDataType::new("int8".to_string())),
9835 nullable: true,
9836 metadata: None,
9837 },
9838 JsonArrowField {
9839 name: "float64_col".to_string(),
9840 r#type: Box::new(JsonArrowDataType::new("float64".to_string())),
9841 nullable: true,
9842 metadata: None,
9843 },
9844 JsonArrowField {
9845 name: "binary_col".to_string(),
9846 r#type: Box::new(JsonArrowDataType::new("binary".to_string())),
9847 nullable: true,
9848 metadata: None,
9849 },
9850 ];
9851
9852 let schema = JsonArrowSchema {
9853 fields,
9854 metadata: None,
9855 };
9856
9857 let ipc_data = create_test_ipc_data(&schema);
9859
9860 let mut request = CreateTableRequest::new();
9861 request.id = Some(vec!["complex_table".to_string()]);
9862
9863 let response = namespace
9864 .create_table(request, bytes::Bytes::from(ipc_data))
9865 .await
9866 .unwrap();
9867
9868 assert!(response.location.is_some());
9869 }
9870
9871 #[tokio::test]
9872 async fn test_connect_dir() {
9873 let temp_dir = TempStdDir::default();
9874
9875 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9876 .build()
9877 .await
9878 .unwrap();
9879
9880 let mut request = ListTablesRequest::new();
9882 request.id = Some(vec![]);
9883 let response = namespace.list_tables(request).await.unwrap();
9884 assert_eq!(response.tables.len(), 0);
9885 }
9886
9887 #[tokio::test]
9888 async fn test_create_table_with_ipc_data() {
9889 use arrow::array::{Int32Array, StringArray};
9890 use arrow::ipc::writer::StreamWriter;
9891
9892 let (namespace, _temp_dir) = create_test_namespace().await;
9893
9894 let schema = create_test_schema();
9896
9897 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
9899 let arrow_schema = Arc::new(arrow_schema);
9900
9901 let id_array = Int32Array::from(vec![1, 2, 3]);
9903 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
9904 let batch = arrow::record_batch::RecordBatch::try_new(
9905 arrow_schema.clone(),
9906 vec![Arc::new(id_array), Arc::new(name_array)],
9907 )
9908 .unwrap();
9909
9910 let mut buffer = Vec::new();
9912 {
9913 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
9914 writer.write(&batch).unwrap();
9915 writer.finish().unwrap();
9916 }
9917
9918 let mut request = CreateTableRequest::new();
9920 request.id = Some(vec!["test_table_with_data".to_string()]);
9921
9922 let response = namespace
9923 .create_table(request, Bytes::from(buffer))
9924 .await
9925 .unwrap();
9926
9927 assert_eq!(response.version, Some(1));
9928 assert!(
9929 response
9930 .location
9931 .unwrap()
9932 .contains("test_table_with_data.lance")
9933 );
9934
9935 let mut exists_request = TableExistsRequest::new();
9937 exists_request.id = Some(vec!["test_table_with_data".to_string()]);
9938 namespace.table_exists(exists_request).await.unwrap();
9939 }
9940
9941 #[tokio::test]
9942 async fn test_child_namespace_create_and_list() {
9943 let (namespace, _temp_dir) = create_test_namespace().await;
9944
9945 for i in 1..=3 {
9947 let mut create_req = CreateNamespaceRequest::new();
9948 create_req.id = Some(vec![format!("ns{}", i)]);
9949 let result = namespace.create_namespace(create_req).await;
9950 assert!(result.is_ok(), "Failed to create child namespace ns{}", i);
9951 }
9952
9953 let list_req = ListNamespacesRequest {
9955 id: Some(vec![]),
9956 ..Default::default()
9957 };
9958 let result = namespace.list_namespaces(list_req).await;
9959 assert!(result.is_ok());
9960 let namespaces = result.unwrap().namespaces;
9961 assert_eq!(namespaces.len(), 3);
9962 assert!(namespaces.contains(&"ns1".to_string()));
9963 assert!(namespaces.contains(&"ns2".to_string()));
9964 assert!(namespaces.contains(&"ns3".to_string()));
9965 }
9966
9967 #[tokio::test]
9968 async fn test_nested_namespace_hierarchy() {
9969 let (namespace, _temp_dir) = create_test_namespace().await;
9970
9971 let mut create_req = CreateNamespaceRequest::new();
9973 create_req.id = Some(vec!["parent".to_string()]);
9974 namespace.create_namespace(create_req).await.unwrap();
9975
9976 let mut create_req = CreateNamespaceRequest::new();
9978 create_req.id = Some(vec!["parent".to_string(), "child1".to_string()]);
9979 namespace.create_namespace(create_req).await.unwrap();
9980
9981 let mut create_req = CreateNamespaceRequest::new();
9982 create_req.id = Some(vec!["parent".to_string(), "child2".to_string()]);
9983 namespace.create_namespace(create_req).await.unwrap();
9984
9985 let list_req = ListNamespacesRequest {
9987 id: Some(vec!["parent".to_string()]),
9988 ..Default::default()
9989 };
9990 let result = namespace.list_namespaces(list_req).await;
9991 assert!(result.is_ok());
9992 let children = result.unwrap().namespaces;
9993 assert_eq!(children.len(), 2);
9994 assert!(children.contains(&"child1".to_string()));
9995 assert!(children.contains(&"child2".to_string()));
9996
9997 let list_req = ListNamespacesRequest {
9999 id: Some(vec![]),
10000 ..Default::default()
10001 };
10002 let result = namespace.list_namespaces(list_req).await;
10003 assert!(result.is_ok());
10004 let root_namespaces = result.unwrap().namespaces;
10005 assert_eq!(root_namespaces.len(), 1);
10006 assert_eq!(root_namespaces[0], "parent");
10007 }
10008
10009 #[tokio::test]
10010 async fn test_table_in_child_namespace() {
10011 let (namespace, _temp_dir) = create_test_namespace().await;
10012
10013 let mut create_ns_req = CreateNamespaceRequest::new();
10015 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10016 namespace.create_namespace(create_ns_req).await.unwrap();
10017
10018 let schema = create_test_schema();
10020 let ipc_data = create_test_ipc_data(&schema);
10021 let mut create_table_req = CreateTableRequest::new();
10022 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10023 let result = namespace
10024 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10025 .await;
10026 assert!(result.is_ok(), "Failed to create table in child namespace");
10027
10028 let list_req = ListTablesRequest {
10030 id: Some(vec!["test_ns".to_string()]),
10031 ..Default::default()
10032 };
10033 let result = namespace.list_tables(list_req).await;
10034 assert!(result.is_ok());
10035 let tables = result.unwrap().tables;
10036 assert_eq!(tables.len(), 1);
10037 assert_eq!(tables[0], "table1");
10038
10039 let mut exists_req = TableExistsRequest::new();
10041 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10042 let result = namespace.table_exists(exists_req).await;
10043 assert!(result.is_ok());
10044
10045 let mut describe_req = DescribeTableRequest::new();
10047 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10048 let result = namespace.describe_table(describe_req).await;
10049 assert!(result.is_ok());
10050 let response = result.unwrap();
10051 assert!(response.location.is_some());
10052 }
10053
10054 #[tokio::test]
10063 async fn test_read_self_heals_after_manifest_created_by_other_connection() {
10064 let temp_dir = TempStdDir::default();
10065 let root = temp_dir.to_str().unwrap();
10066
10067 let reader = DirectoryNamespaceBuilder::new(root).build().await.unwrap();
10070
10071 let writer = DirectoryNamespaceBuilder::new(root).build().await.unwrap();
10074 let mut create_ns_req = CreateNamespaceRequest::new();
10075 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10076 writer.create_namespace(create_ns_req).await.unwrap();
10077 let ipc_data = create_test_ipc_data(&create_test_schema());
10078 let mut create_table_req = CreateTableRequest::new();
10079 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10080 writer
10081 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10082 .await
10083 .unwrap();
10084
10085 let mut describe_req = DescribeTableRequest::new();
10088 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10089 let resp = reader
10090 .describe_table(describe_req)
10091 .await
10092 .expect("describe_table must resolve a table registered after build");
10093 assert!(resp.location.is_some());
10094
10095 let mut exists_req = TableExistsRequest::new();
10096 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10097 reader
10098 .table_exists(exists_req)
10099 .await
10100 .expect("table_exists must resolve a table registered after build");
10101
10102 let list_req = ListTablesRequest {
10103 id: Some(vec!["test_ns".to_string()]),
10104 ..Default::default()
10105 };
10106 let tables = reader.list_tables(list_req).await.unwrap().tables;
10107 assert_eq!(tables, vec!["table1".to_string()]);
10108 }
10109
10110 #[tokio::test]
10118 async fn test_migration_root_read_self_heals_registered_alias() {
10119 use lance_namespace::models::RegisterTableRequest;
10120
10121 let temp_dir = TempStdDir::default();
10122 let temp_path = temp_dir.to_str().unwrap();
10123
10124 let reader = DirectoryNamespaceBuilder::new(temp_path)
10127 .dir_listing_enabled(true)
10128 .dir_listing_to_manifest_migration_enabled(true)
10129 .build()
10130 .await
10131 .unwrap();
10132
10133 let writer = DirectoryNamespaceBuilder::new(temp_path)
10137 .dir_listing_enabled(true)
10138 .dir_listing_to_manifest_migration_enabled(true)
10139 .build()
10140 .await
10141 .unwrap();
10142 let ipc_data = create_test_ipc_data(&create_test_schema());
10143 let table_uri = format!("{}/external_table.lance", temp_path);
10144 let cursor = Cursor::new(ipc_data);
10145 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
10146 let batches: Vec<_> = stream_reader
10147 .collect::<std::result::Result<Vec<_>, _>>()
10148 .unwrap();
10149 let schema = batches[0].schema();
10150 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
10151 let batch_reader = RecordBatchIterator::new(batch_results, schema);
10152 Dataset::write(Box::new(batch_reader), &table_uri, None)
10153 .await
10154 .unwrap();
10155 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
10156 register_req.id = Some(vec!["registered_table".to_string()]);
10157 writer.register_table(register_req).await.unwrap();
10158
10159 let mut describe_req = DescribeTableRequest::new();
10162 describe_req.id = Some(vec!["registered_table".to_string()]);
10163 reader
10164 .describe_table(describe_req)
10165 .await
10166 .expect("describe_table must resolve a manifest alias registered after build");
10167
10168 let mut exists_req = TableExistsRequest::new();
10169 exists_req.id = Some(vec!["registered_table".to_string()]);
10170 reader
10171 .table_exists(exists_req)
10172 .await
10173 .expect("table_exists must resolve a manifest alias registered after build");
10174
10175 let list_req = ListTablesRequest {
10176 id: Some(vec![]),
10177 ..Default::default()
10178 };
10179 let tables = reader.list_tables(list_req).await.unwrap().tables;
10180 assert!(
10181 tables.contains(&"registered_table".to_string()),
10182 "list_tables must include the manifest alias registered after build, got {:?}",
10183 tables
10184 );
10185 }
10186
10187 #[tokio::test]
10188 async fn test_multiple_tables_in_child_namespace() {
10189 let (namespace, _temp_dir) = create_test_namespace().await;
10190
10191 let mut create_ns_req = CreateNamespaceRequest::new();
10193 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10194 namespace.create_namespace(create_ns_req).await.unwrap();
10195
10196 let schema = create_test_schema();
10198 let ipc_data = create_test_ipc_data(&schema);
10199 for i in 1..=3 {
10200 let mut create_table_req = CreateTableRequest::new();
10201 create_table_req.id = Some(vec!["test_ns".to_string(), format!("table{}", i)]);
10202 namespace
10203 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
10204 .await
10205 .unwrap();
10206 }
10207
10208 let list_req = ListTablesRequest {
10210 id: Some(vec!["test_ns".to_string()]),
10211 ..Default::default()
10212 };
10213 let result = namespace.list_tables(list_req).await;
10214 assert!(result.is_ok());
10215 let tables = result.unwrap().tables;
10216 assert_eq!(tables.len(), 3);
10217 assert!(tables.contains(&"table1".to_string()));
10218 assert!(tables.contains(&"table2".to_string()));
10219 assert!(tables.contains(&"table3".to_string()));
10220 }
10221
10222 #[tokio::test]
10223 async fn test_drop_table_in_child_namespace() {
10224 let (namespace, _temp_dir) = create_test_namespace().await;
10225
10226 let mut create_ns_req = CreateNamespaceRequest::new();
10228 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10229 namespace.create_namespace(create_ns_req).await.unwrap();
10230
10231 let schema = create_test_schema();
10233 let ipc_data = create_test_ipc_data(&schema);
10234 let mut create_table_req = CreateTableRequest::new();
10235 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10236 namespace
10237 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10238 .await
10239 .unwrap();
10240
10241 let mut drop_req = DropTableRequest::new();
10243 drop_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10244 let result = namespace.drop_table(drop_req).await;
10245 assert!(result.is_ok(), "Failed to drop table in child namespace");
10246
10247 let mut exists_req = TableExistsRequest::new();
10249 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10250 let result = namespace.table_exists(exists_req).await;
10251 assert!(result.is_err());
10252 }
10253
10254 #[tokio::test]
10255 async fn test_deeply_nested_namespace() {
10256 let (namespace, _temp_dir) = create_test_namespace().await;
10257
10258 let mut create_req = CreateNamespaceRequest::new();
10260 create_req.id = Some(vec!["level1".to_string()]);
10261 namespace.create_namespace(create_req).await.unwrap();
10262
10263 let mut create_req = CreateNamespaceRequest::new();
10264 create_req.id = Some(vec!["level1".to_string(), "level2".to_string()]);
10265 namespace.create_namespace(create_req).await.unwrap();
10266
10267 let mut create_req = CreateNamespaceRequest::new();
10268 create_req.id = Some(vec![
10269 "level1".to_string(),
10270 "level2".to_string(),
10271 "level3".to_string(),
10272 ]);
10273 namespace.create_namespace(create_req).await.unwrap();
10274
10275 let schema = create_test_schema();
10277 let ipc_data = create_test_ipc_data(&schema);
10278 let mut create_table_req = CreateTableRequest::new();
10279 create_table_req.id = Some(vec![
10280 "level1".to_string(),
10281 "level2".to_string(),
10282 "level3".to_string(),
10283 "table1".to_string(),
10284 ]);
10285 let result = namespace
10286 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10287 .await;
10288 assert!(
10289 result.is_ok(),
10290 "Failed to create table in deeply nested namespace"
10291 );
10292
10293 let mut exists_req = TableExistsRequest::new();
10295 exists_req.id = Some(vec![
10296 "level1".to_string(),
10297 "level2".to_string(),
10298 "level3".to_string(),
10299 "table1".to_string(),
10300 ]);
10301 let result = namespace.table_exists(exists_req).await;
10302 assert!(result.is_ok());
10303 }
10304
10305 #[tokio::test]
10306 async fn test_namespace_with_properties() {
10307 let (namespace, _temp_dir) = create_test_namespace().await;
10308
10309 let mut properties = HashMap::new();
10311 properties.insert("owner".to_string(), "test_user".to_string());
10312 properties.insert("description".to_string(), "Test namespace".to_string());
10313
10314 let mut create_req = CreateNamespaceRequest::new();
10315 create_req.id = Some(vec!["test_ns".to_string()]);
10316 create_req.properties = Some(properties.clone());
10317 namespace.create_namespace(create_req).await.unwrap();
10318
10319 let describe_req = DescribeNamespaceRequest {
10321 id: Some(vec!["test_ns".to_string()]),
10322 ..Default::default()
10323 };
10324 let result = namespace.describe_namespace(describe_req).await;
10325 assert!(result.is_ok());
10326 let response = result.unwrap();
10327 assert!(response.properties.is_some());
10328 let props = response.properties.unwrap();
10329 assert_eq!(props.get("owner"), Some(&"test_user".to_string()));
10330 assert_eq!(
10331 props.get("description"),
10332 Some(&"Test namespace".to_string())
10333 );
10334 }
10335
10336 #[tokio::test]
10337 async fn test_cannot_drop_namespace_with_tables() {
10338 let (namespace, _temp_dir) = create_test_namespace().await;
10339
10340 let mut create_ns_req = CreateNamespaceRequest::new();
10342 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10343 namespace.create_namespace(create_ns_req).await.unwrap();
10344
10345 let schema = create_test_schema();
10347 let ipc_data = create_test_ipc_data(&schema);
10348 let mut create_table_req = CreateTableRequest::new();
10349 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10350 namespace
10351 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10352 .await
10353 .unwrap();
10354
10355 let mut drop_req = DropNamespaceRequest::new();
10357 drop_req.id = Some(vec!["test_ns".to_string()]);
10358 let result = namespace.drop_namespace(drop_req).await;
10359 assert!(
10360 result.is_err(),
10361 "Should not be able to drop namespace with tables"
10362 );
10363 }
10364
10365 #[tokio::test]
10366 async fn test_isolation_between_namespaces() {
10367 let (namespace, _temp_dir) = create_test_namespace().await;
10368
10369 let mut create_req = CreateNamespaceRequest::new();
10371 create_req.id = Some(vec!["ns1".to_string()]);
10372 namespace.create_namespace(create_req).await.unwrap();
10373
10374 let mut create_req = CreateNamespaceRequest::new();
10375 create_req.id = Some(vec!["ns2".to_string()]);
10376 namespace.create_namespace(create_req).await.unwrap();
10377
10378 let schema = create_test_schema();
10380 let ipc_data = create_test_ipc_data(&schema);
10381
10382 let mut create_table_req = CreateTableRequest::new();
10383 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10384 namespace
10385 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
10386 .await
10387 .unwrap();
10388
10389 let mut create_table_req = CreateTableRequest::new();
10390 create_table_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
10391 namespace
10392 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10393 .await
10394 .unwrap();
10395
10396 let list_req = ListTablesRequest {
10398 id: Some(vec!["ns1".to_string()]),
10399 page_token: None,
10400 limit: None,
10401 ..Default::default()
10402 };
10403 let result = namespace.list_tables(list_req).await.unwrap();
10404 assert_eq!(result.tables.len(), 1);
10405 assert_eq!(result.tables[0], "table1");
10406
10407 let list_req = ListTablesRequest {
10408 id: Some(vec!["ns2".to_string()]),
10409 page_token: None,
10410 limit: None,
10411 ..Default::default()
10412 };
10413 let result = namespace.list_tables(list_req).await.unwrap();
10414 assert_eq!(result.tables.len(), 1);
10415 assert_eq!(result.tables[0], "table1");
10416
10417 let mut drop_req = DropTableRequest::new();
10419 drop_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10420 namespace.drop_table(drop_req).await.unwrap();
10421
10422 let mut exists_req = TableExistsRequest::new();
10424 exists_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10425 assert!(namespace.table_exists(exists_req).await.is_err());
10426
10427 let mut exists_req = TableExistsRequest::new();
10428 exists_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
10429 assert!(namespace.table_exists(exists_req).await.is_ok());
10430 }
10431
10432 #[tokio::test]
10433 async fn test_migrate_directory_tables() {
10434 let temp_dir = TempStdDir::default();
10435 let temp_path = temp_dir.to_str().unwrap();
10436
10437 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
10439 .manifest_enabled(false)
10440 .dir_listing_enabled(true)
10441 .build()
10442 .await
10443 .unwrap();
10444
10445 let schema = create_test_schema();
10447 let ipc_data = create_test_ipc_data(&schema);
10448
10449 for i in 1..=3 {
10450 let mut create_req = CreateTableRequest::new();
10451 create_req.id = Some(vec![format!("table{}", i)]);
10452 dir_only_ns
10453 .create_table(create_req, bytes::Bytes::from(ipc_data.clone()))
10454 .await
10455 .unwrap();
10456 }
10457
10458 drop(dir_only_ns);
10459
10460 let dual_mode_ns = DirectoryNamespaceBuilder::new(temp_path)
10462 .manifest_enabled(true)
10463 .dir_listing_enabled(true)
10464 .build()
10465 .await
10466 .unwrap();
10467
10468 let mut list_req = ListTablesRequest::new();
10470 list_req.id = Some(vec![]);
10471 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
10472 assert_eq!(tables.len(), 3);
10473
10474 let migrated_count = dual_mode_ns.migrate().await.unwrap();
10476 assert_eq!(migrated_count, 3, "Should migrate all 3 tables");
10477
10478 let mut list_req = ListTablesRequest::new();
10480 list_req.id = Some(vec![]);
10481 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
10482 assert_eq!(tables.len(), 3);
10483
10484 let migrated_count = dual_mode_ns.migrate().await.unwrap();
10486 assert_eq!(
10487 migrated_count, 0,
10488 "Should not migrate already-migrated tables"
10489 );
10490
10491 drop(dual_mode_ns);
10492
10493 let manifest_only_ns = DirectoryNamespaceBuilder::new(temp_path)
10495 .manifest_enabled(true)
10496 .dir_listing_enabled(false)
10497 .build()
10498 .await
10499 .unwrap();
10500
10501 let mut list_req = ListTablesRequest::new();
10503 list_req.id = Some(vec![]);
10504 let tables = manifest_only_ns.list_tables(list_req).await.unwrap().tables;
10505 assert_eq!(tables.len(), 3);
10506 assert!(tables.contains(&"table1".to_string()));
10507 assert!(tables.contains(&"table2".to_string()));
10508 assert!(tables.contains(&"table3".to_string()));
10509 }
10510
10511 #[tokio::test]
10512 async fn test_migrate_without_manifest() {
10513 let temp_dir = TempStdDir::default();
10514 let temp_path = temp_dir.to_str().unwrap();
10515
10516 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10518 .manifest_enabled(false)
10519 .dir_listing_enabled(true)
10520 .build()
10521 .await
10522 .unwrap();
10523
10524 let migrated_count = namespace.migrate().await.unwrap();
10526 assert_eq!(migrated_count, 0);
10527 }
10528
10529 #[tokio::test]
10530 async fn test_register_table() {
10531 use lance_namespace::models::{RegisterTableRequest, TableExistsRequest};
10532
10533 let temp_dir = TempStdDir::default();
10534 let temp_path = temp_dir.to_str().unwrap();
10535
10536 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10537 .dir_listing_to_manifest_migration_enabled(true)
10538 .build()
10539 .await
10540 .unwrap();
10541
10542 let schema = create_test_schema();
10544 let ipc_data = create_test_ipc_data(&schema);
10545
10546 let table_uri = format!("{}/external_table.lance", temp_path);
10547 let cursor = Cursor::new(ipc_data);
10548 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
10549 let batches: Vec<_> = stream_reader
10550 .collect::<std::result::Result<Vec<_>, _>>()
10551 .unwrap();
10552 let schema = batches[0].schema();
10553 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
10554 let reader = RecordBatchIterator::new(batch_results, schema);
10555 Dataset::write(Box::new(reader), &table_uri, None)
10556 .await
10557 .unwrap();
10558
10559 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
10561 register_req.id = Some(vec!["registered_table".to_string()]);
10562
10563 let response = namespace.register_table(register_req).await.unwrap();
10564 assert_eq!(response.location, Some("external_table.lance".to_string()));
10565
10566 let mut exists_req = TableExistsRequest::new();
10568 exists_req.id = Some(vec!["registered_table".to_string()]);
10569 assert!(namespace.table_exists(exists_req).await.is_ok());
10570
10571 let mut list_req = ListTablesRequest::new();
10573 list_req.id = Some(vec![]);
10574 let tables = namespace.list_tables(list_req).await.unwrap();
10575 assert!(tables.tables.contains(&"registered_table".to_string()));
10576 }
10577
10578 #[tokio::test]
10579 async fn test_register_table_duplicate_fails() {
10580 use lance_namespace::models::RegisterTableRequest;
10581
10582 let temp_dir = TempStdDir::default();
10583 let temp_path = temp_dir.to_str().unwrap();
10584
10585 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10586 .build()
10587 .await
10588 .unwrap();
10589
10590 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10592 register_req.id = Some(vec!["test_table".to_string()]);
10593
10594 namespace
10595 .register_table(register_req.clone())
10596 .await
10597 .unwrap();
10598
10599 let result = namespace.register_table(register_req).await;
10601 assert!(result.is_err());
10602 assert!(result.unwrap_err().to_string().contains("already exists"));
10603 }
10604
10605 #[tokio::test]
10606 async fn test_deregister_table() {
10607 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
10608
10609 let temp_dir = TempStdDir::default();
10610 let temp_path = temp_dir.to_str().unwrap();
10611
10612 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10615 .manifest_enabled(true)
10616 .dir_listing_enabled(false)
10617 .build()
10618 .await
10619 .unwrap();
10620
10621 let schema = create_test_schema();
10623 let ipc_data = create_test_ipc_data(&schema);
10624
10625 let mut create_req = CreateTableRequest::new();
10626 create_req.id = Some(vec!["test_table".to_string()]);
10627 namespace
10628 .create_table(create_req, bytes::Bytes::from(ipc_data))
10629 .await
10630 .unwrap();
10631
10632 let mut exists_req = TableExistsRequest::new();
10634 exists_req.id = Some(vec!["test_table".to_string()]);
10635 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
10636
10637 let mut deregister_req = DeregisterTableRequest::new();
10639 deregister_req.id = Some(vec!["test_table".to_string()]);
10640 let response = namespace.deregister_table(deregister_req).await.unwrap();
10641
10642 assert!(
10644 response.location.is_some(),
10645 "Deregister should return location"
10646 );
10647 let location = response.location.as_ref().unwrap();
10648 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10651 .expect("Failed to convert temp path to URL");
10652 let expected_prefix = expected_url.to_string();
10653 assert!(
10654 location.starts_with(&expected_prefix),
10655 "Location should start with '{}', got: {}",
10656 expected_prefix,
10657 location
10658 );
10659 assert!(
10660 location.contains("test_table"),
10661 "Location should contain table name: {}",
10662 location
10663 );
10664 assert_eq!(response.id, Some(vec!["test_table".to_string()]));
10665
10666 assert!(namespace.table_exists(exists_req).await.is_err());
10668
10669 let dataset = Dataset::open(location).await;
10671 assert!(
10672 dataset.is_ok(),
10673 "Physical table data should still exist at {}",
10674 location
10675 );
10676 }
10677
10678 #[tokio::test]
10679 async fn test_deregister_table_in_child_namespace() {
10680 use lance_namespace::models::{
10681 CreateNamespaceRequest, DeregisterTableRequest, TableExistsRequest,
10682 };
10683
10684 let temp_dir = TempStdDir::default();
10685 let temp_path = temp_dir.to_str().unwrap();
10686
10687 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10688 .build()
10689 .await
10690 .unwrap();
10691
10692 let mut create_ns_req = CreateNamespaceRequest::new();
10694 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10695 namespace.create_namespace(create_ns_req).await.unwrap();
10696
10697 let schema = create_test_schema();
10699 let ipc_data = create_test_ipc_data(&schema);
10700
10701 let mut create_req = CreateTableRequest::new();
10702 create_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10703 namespace
10704 .create_table(create_req, bytes::Bytes::from(ipc_data))
10705 .await
10706 .unwrap();
10707
10708 let mut deregister_req = DeregisterTableRequest::new();
10710 deregister_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10711 let response = namespace.deregister_table(deregister_req).await.unwrap();
10712
10713 assert!(
10715 response.location.is_some(),
10716 "Deregister should return location"
10717 );
10718 let location = response.location.as_ref().unwrap();
10719 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10722 .expect("Failed to convert temp path to URL");
10723 let expected_prefix = expected_url.to_string();
10724 assert!(
10725 location.starts_with(&expected_prefix),
10726 "Location should start with '{}', got: {}",
10727 expected_prefix,
10728 location
10729 );
10730 assert!(
10731 location.contains("test_ns") && location.contains("test_table"),
10732 "Location should contain namespace and table name: {}",
10733 location
10734 );
10735 assert_eq!(
10736 response.id,
10737 Some(vec!["test_ns".to_string(), "test_table".to_string()])
10738 );
10739
10740 let mut exists_req = TableExistsRequest::new();
10742 exists_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10743 assert!(namespace.table_exists(exists_req).await.is_err());
10744 }
10745
10746 #[tokio::test]
10747 async fn test_register_without_manifest_fails() {
10748 use lance_namespace::models::RegisterTableRequest;
10749
10750 let temp_dir = TempStdDir::default();
10751 let temp_path = temp_dir.to_str().unwrap();
10752
10753 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10755 .manifest_enabled(false)
10756 .build()
10757 .await
10758 .unwrap();
10759
10760 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10762 register_req.id = Some(vec!["test_table".to_string()]);
10763 let result = namespace.register_table(register_req).await;
10764 assert!(result.is_err());
10765 assert!(
10766 result
10767 .unwrap_err()
10768 .to_string()
10769 .contains("manifest mode is enabled")
10770 );
10771
10772 }
10775
10776 #[tokio::test]
10777 async fn test_register_table_rejects_absolute_uri() {
10778 use lance_namespace::models::RegisterTableRequest;
10779
10780 let temp_dir = TempStdDir::default();
10781 let temp_path = temp_dir.to_str().unwrap();
10782
10783 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10784 .build()
10785 .await
10786 .unwrap();
10787
10788 let mut register_req = RegisterTableRequest::new("s3://bucket/table.lance".to_string());
10790 register_req.id = Some(vec!["test_table".to_string()]);
10791 let result = namespace.register_table(register_req).await;
10792 assert!(result.is_err());
10793 let err_msg = result.unwrap_err().to_string();
10794 assert!(err_msg.contains("Absolute URIs are not allowed"));
10795 }
10796
10797 #[tokio::test]
10798 async fn test_register_table_rejects_absolute_path() {
10799 use lance_namespace::models::RegisterTableRequest;
10800
10801 let temp_dir = TempStdDir::default();
10802 let temp_path = temp_dir.to_str().unwrap();
10803
10804 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10805 .build()
10806 .await
10807 .unwrap();
10808
10809 let mut register_req = RegisterTableRequest::new("/tmp/table.lance".to_string());
10811 register_req.id = Some(vec!["test_table".to_string()]);
10812 let result = namespace.register_table(register_req).await;
10813 assert!(result.is_err());
10814 let err_msg = result.unwrap_err().to_string();
10815 assert!(err_msg.contains("Absolute paths are not allowed"));
10816 }
10817
10818 #[tokio::test]
10819 async fn test_register_table_rejects_path_traversal() {
10820 use lance_namespace::models::RegisterTableRequest;
10821
10822 let temp_dir = TempStdDir::default();
10823 let temp_path = temp_dir.to_str().unwrap();
10824
10825 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10826 .build()
10827 .await
10828 .unwrap();
10829
10830 let mut register_req = RegisterTableRequest::new("../outside/table.lance".to_string());
10832 register_req.id = Some(vec!["test_table".to_string()]);
10833 let result = namespace.register_table(register_req).await;
10834 assert!(result.is_err());
10835 let err_msg = result.unwrap_err().to_string();
10836 assert!(err_msg.contains("Path traversal is not allowed"));
10837 }
10838
10839 #[tokio::test]
10840 async fn test_namespace_write() {
10841 use arrow::array::Int32Array;
10842 use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
10843 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
10844 use lance::dataset::{Dataset, WriteMode, WriteParams};
10845 use lance_namespace::LanceNamespace;
10846
10847 let (namespace, _temp_dir) = create_test_namespace().await;
10848 let namespace = Arc::new(namespace) as Arc<dyn LanceNamespace>;
10849
10850 let table_id = vec!["test_ns".to_string(), "test_table".to_string()];
10852 let schema = Arc::new(ArrowSchema::new(vec![
10853 ArrowField::new("a", DataType::Int32, false),
10854 ArrowField::new("b", DataType::Int32, false),
10855 ]));
10856
10857 let data1 = RecordBatch::try_new(
10859 schema.clone(),
10860 vec![
10861 Arc::new(Int32Array::from(vec![1, 2, 3])),
10862 Arc::new(Int32Array::from(vec![10, 20, 30])),
10863 ],
10864 )
10865 .unwrap();
10866
10867 let reader1 = RecordBatchIterator::new(vec![data1].into_iter().map(Ok), schema.clone());
10868 let dataset =
10869 Dataset::write_into_namespace(reader1, namespace.clone(), table_id.clone(), None)
10870 .await
10871 .unwrap();
10872
10873 assert_eq!(dataset.count_rows(None).await.unwrap(), 3);
10874 assert_eq!(dataset.version().version, 1);
10875
10876 let data2 = RecordBatch::try_new(
10878 schema.clone(),
10879 vec![
10880 Arc::new(Int32Array::from(vec![4, 5])),
10881 Arc::new(Int32Array::from(vec![40, 50])),
10882 ],
10883 )
10884 .unwrap();
10885
10886 let params_append = WriteParams {
10887 mode: WriteMode::Append,
10888 ..Default::default()
10889 };
10890
10891 let reader2 = RecordBatchIterator::new(vec![data2].into_iter().map(Ok), schema.clone());
10892 let dataset = Dataset::write_into_namespace(
10893 reader2,
10894 namespace.clone(),
10895 table_id.clone(),
10896 Some(params_append),
10897 )
10898 .await
10899 .unwrap();
10900
10901 assert_eq!(dataset.count_rows(None).await.unwrap(), 5);
10902 assert_eq!(dataset.version().version, 2);
10903
10904 let data3 = RecordBatch::try_new(
10906 schema.clone(),
10907 vec![
10908 Arc::new(Int32Array::from(vec![100, 200])),
10909 Arc::new(Int32Array::from(vec![1000, 2000])),
10910 ],
10911 )
10912 .unwrap();
10913
10914 let params_overwrite = WriteParams {
10915 mode: WriteMode::Overwrite,
10916 ..Default::default()
10917 };
10918
10919 let reader3 = RecordBatchIterator::new(vec![data3].into_iter().map(Ok), schema.clone());
10920 let dataset = Dataset::write_into_namespace(
10921 reader3,
10922 namespace.clone(),
10923 table_id.clone(),
10924 Some(params_overwrite),
10925 )
10926 .await
10927 .unwrap();
10928
10929 assert_eq!(dataset.count_rows(None).await.unwrap(), 2);
10930 assert_eq!(dataset.version().version, 3);
10931
10932 let result = dataset.scan().try_into_batch().await.unwrap();
10934 let a_col = result
10935 .column_by_name("a")
10936 .unwrap()
10937 .as_any()
10938 .downcast_ref::<Int32Array>()
10939 .unwrap();
10940 assert_eq!(a_col.values(), &[100, 200]);
10941 }
10942
10943 #[tokio::test]
10948 async fn test_declare_table_v1_mode() {
10949 use lance_namespace::models::{
10950 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
10951 };
10952
10953 let temp_dir = TempStdDir::default();
10954 let temp_path = temp_dir.to_str().unwrap();
10955
10956 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10958 .manifest_enabled(false)
10959 .build()
10960 .await
10961 .unwrap();
10962
10963 let mut declare_req = DeclareTableRequest::new();
10965 declare_req.id = Some(vec!["test_table".to_string()]);
10966 let response = namespace.declare_table(declare_req).await.unwrap();
10967
10968 assert!(response.location.is_some());
10970 let location = response.location.as_ref().unwrap();
10971 assert!(location.ends_with("test_table.lance"));
10972
10973 let mut exists_req = TableExistsRequest::new();
10975 exists_req.id = Some(vec!["test_table".to_string()]);
10976 assert!(namespace.table_exists(exists_req).await.is_ok());
10977
10978 let mut describe_req = DescribeTableRequest::new();
10980 describe_req.id = Some(vec!["test_table".to_string()]);
10981 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10982 assert!(describe_response.location.is_some());
10983 assert!(describe_response.version.is_none()); assert!(describe_response.schema.is_none()); assert_eq!(describe_response.is_only_declared, None);
10986
10987 let mut describe_req = DescribeTableRequest::new();
10988 describe_req.id = Some(vec!["test_table".to_string()]);
10989 describe_req.check_declared = Some(true);
10990 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10991 assert_eq!(describe_response.is_only_declared, Some(true));
10992
10993 let mut list_req = ListTablesRequest::new();
10994 list_req.id = Some(vec![]);
10995 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
10996 assert_eq!(list_response.tables, vec!["test_table".to_string()]);
10997
10998 list_req.include_declared = Some(false);
10999 let list_response = namespace.list_tables(list_req).await.unwrap();
11000 assert!(list_response.tables.is_empty());
11001 }
11002
11003 #[tokio::test]
11004 async fn test_insert_into_declared_table_promotes_it_from_declared_state() {
11005 use lance_namespace::models::{
11006 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest,
11007 };
11008
11009 let temp_dir = TempStdDir::default();
11010 let temp_path = temp_dir.to_str().unwrap();
11011
11012 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11013 .manifest_enabled(false)
11014 .build()
11015 .await
11016 .unwrap();
11017
11018 let mut declare_req = DeclareTableRequest::new();
11019 declare_req.id = Some(vec!["test_table".to_string()]);
11020 namespace.declare_table(declare_req).await.unwrap();
11021
11022 let schema = create_test_schema();
11023 let ipc_data = create_test_ipc_data(&schema);
11024 let mut insert_req = InsertIntoTableRequest::new();
11025 insert_req.id = Some(vec!["test_table".to_string()]);
11026 namespace
11027 .insert_into_table(insert_req, bytes::Bytes::from(ipc_data))
11028 .await
11029 .unwrap();
11030
11031 let mut describe_req = DescribeTableRequest::new();
11032 describe_req.id = Some(vec!["test_table".to_string()]);
11033 describe_req.load_detailed_metadata = Some(true);
11034 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11035
11036 assert_eq!(describe_response.is_only_declared, Some(false));
11037 assert_eq!(describe_response.version, Some(1));
11038 assert!(describe_response.schema.is_some());
11039
11040 let mut list_req = ListTablesRequest::new();
11041 list_req.id = Some(vec![]);
11042 list_req.include_declared = Some(false);
11043 assert_eq!(
11044 namespace.list_tables(list_req).await.unwrap().tables,
11045 vec!["test_table".to_string()]
11046 );
11047 }
11048
11049 #[tokio::test]
11050 async fn test_create_table_after_declare_table_v1_mode_creates_table() {
11051 use lance_namespace::models::{
11052 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11053 };
11054
11055 let temp_dir = TempStdDir::default();
11056 let temp_path = temp_dir.to_str().unwrap();
11057
11058 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11059 .manifest_enabled(false)
11060 .build()
11061 .await
11062 .unwrap();
11063
11064 let mut declare_req = DeclareTableRequest::new();
11065 declare_req.id = Some(vec!["test_table".to_string()]);
11066 namespace.declare_table(declare_req).await.unwrap();
11067
11068 let mut create_req = CreateTableRequest::new();
11069 create_req.id = Some(vec!["test_table".to_string()]);
11070 let response = namespace
11071 .create_table(
11072 create_req,
11073 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11074 )
11075 .await
11076 .unwrap();
11077
11078 assert_eq!(response.version, Some(1));
11079
11080 let mut describe_req = DescribeTableRequest::new();
11081 describe_req.id = Some(vec!["test_table".to_string()]);
11082 describe_req.load_detailed_metadata = Some(true);
11083 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11084 assert_eq!(describe_response.is_only_declared, Some(false));
11085 assert_eq!(describe_response.version, Some(1));
11086
11087 let mut list_req = ListTablesRequest::new();
11088 list_req.id = Some(vec![]);
11089 list_req.include_declared = Some(false);
11090 assert_eq!(
11091 namespace.list_tables(list_req).await.unwrap().tables,
11092 vec!["test_table".to_string()]
11093 );
11094 }
11095
11096 #[tokio::test]
11097 async fn test_insert_into_declared_table_with_manifest_promotes_it() {
11098 use lance_namespace::models::{
11099 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest, ListTablesRequest,
11100 };
11101
11102 let temp_dir = TempStdDir::default();
11103 let temp_path = temp_dir.to_str().unwrap();
11104
11105 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11106 .manifest_enabled(true)
11107 .dir_listing_enabled(false)
11108 .build()
11109 .await
11110 .unwrap();
11111
11112 let mut declare_req = DeclareTableRequest::new();
11113 declare_req.id = Some(vec!["test_table".to_string()]);
11114 namespace.declare_table(declare_req).await.unwrap();
11115
11116 let mut insert_req = InsertIntoTableRequest::new();
11117 insert_req.id = Some(vec!["test_table".to_string()]);
11118 namespace
11119 .insert_into_table(
11120 insert_req,
11121 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11122 )
11123 .await
11124 .unwrap();
11125
11126 let mut describe_req = DescribeTableRequest::new();
11127 describe_req.id = Some(vec!["test_table".to_string()]);
11128 describe_req.load_detailed_metadata = Some(true);
11129 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11130 assert_eq!(describe_response.is_only_declared, Some(false));
11131 assert_eq!(describe_response.version, Some(1));
11132
11133 let mut list_req = ListTablesRequest::new();
11134 list_req.id = Some(vec![]);
11135 list_req.include_declared = Some(false);
11136 assert_eq!(
11137 namespace.list_tables(list_req).await.unwrap().tables,
11138 vec!["test_table".to_string()]
11139 );
11140 }
11141
11142 #[tokio::test]
11143 async fn test_create_table_after_declare_table_with_manifest_creates_table() {
11144 use lance_namespace::models::{
11145 CreateTableRequest, DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11146 };
11147
11148 let temp_dir = TempStdDir::default();
11149 let temp_path = temp_dir.to_str().unwrap();
11150
11151 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11152 .manifest_enabled(true)
11153 .dir_listing_enabled(false)
11154 .build()
11155 .await
11156 .unwrap();
11157
11158 let mut declare_req = DeclareTableRequest::new();
11159 declare_req.id = Some(vec!["test_table".to_string()]);
11160 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11161 namespace.declare_table(declare_req).await.unwrap();
11162
11163 let mut create_req = CreateTableRequest::new();
11164 create_req.id = Some(vec!["test_table".to_string()]);
11165 create_req.mode = Some("Overwrite".to_string());
11166 let response = namespace
11167 .create_table(
11168 create_req,
11169 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11170 )
11171 .await
11172 .unwrap();
11173
11174 assert_eq!(response.version, Some(1));
11175 assert_eq!(
11176 response
11177 .properties
11178 .as_ref()
11179 .and_then(|properties| properties.get("owner")),
11180 Some(&"alice".to_string())
11181 );
11182
11183 let mut describe_req = DescribeTableRequest::new();
11184 describe_req.id = Some(vec!["test_table".to_string()]);
11185 describe_req.load_detailed_metadata = Some(true);
11186 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11187 assert_eq!(describe_response.is_only_declared, Some(false));
11188 assert_eq!(describe_response.version, Some(1));
11189 assert_eq!(
11190 describe_response
11191 .properties
11192 .as_ref()
11193 .and_then(|properties| properties.get("owner")),
11194 Some(&"alice".to_string())
11195 );
11196
11197 let mut list_req = ListTablesRequest::new();
11198 list_req.id = Some(vec![]);
11199 list_req.include_declared = Some(false);
11200 assert_eq!(
11201 namespace.list_tables(list_req).await.unwrap().tables,
11202 vec!["test_table".to_string()]
11203 );
11204 }
11205
11206 #[tokio::test]
11207 async fn test_create_table_after_declare_table_with_manifest_rejects_new_properties() {
11208 use lance_namespace::models::{CreateTableRequest, DeclareTableRequest};
11209
11210 let temp_dir = TempStdDir::default();
11211 let temp_path = temp_dir.to_str().unwrap();
11212
11213 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11214 .manifest_enabled(true)
11215 .dir_listing_enabled(false)
11216 .build()
11217 .await
11218 .unwrap();
11219
11220 let mut declare_req = DeclareTableRequest::new();
11221 declare_req.id = Some(vec!["test_table".to_string()]);
11222 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11223 namespace.declare_table(declare_req).await.unwrap();
11224
11225 let mut create_req = CreateTableRequest::new();
11226 create_req.id = Some(vec!["test_table".to_string()]);
11227 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
11228
11229 let result = namespace
11230 .create_table(
11231 create_req,
11232 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11233 )
11234 .await;
11235
11236 assert!(result.is_err());
11237 assert!(
11238 result
11239 .unwrap_err()
11240 .to_string()
11241 .contains("cannot set properties for already declared table")
11242 );
11243 }
11244
11245 #[tokio::test]
11246 async fn test_create_table_with_manifest_exist_ok_keeps_existing_table() {
11247 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
11248
11249 let temp_dir = TempStdDir::default();
11250 let temp_path = temp_dir.to_str().unwrap();
11251
11252 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11253 .manifest_enabled(true)
11254 .dir_listing_enabled(false)
11255 .build()
11256 .await
11257 .unwrap();
11258
11259 let mut create_req = CreateTableRequest::new();
11260 create_req.id = Some(vec!["test_table".to_string()]);
11261 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11262 namespace
11263 .create_table(
11264 create_req,
11265 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11266 )
11267 .await
11268 .unwrap();
11269
11270 let mut create_req = CreateTableRequest::new();
11271 create_req.id = Some(vec!["test_table".to_string()]);
11272 create_req.mode = Some("ExistOk".to_string());
11273 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
11274 let response = namespace
11275 .create_table(
11276 create_req,
11277 bytes::Bytes::from(create_single_row_test_ipc_data()),
11278 )
11279 .await
11280 .unwrap();
11281
11282 assert_eq!(
11283 response
11284 .properties
11285 .as_ref()
11286 .and_then(|properties| properties.get("owner")),
11287 Some(&"alice".to_string())
11288 );
11289 assert_eq!(
11290 open_dataset(&namespace, "test_table")
11291 .await
11292 .count_rows(None)
11293 .await
11294 .unwrap(),
11295 2
11296 );
11297
11298 let mut describe_req = DescribeTableRequest::new();
11299 describe_req.id = Some(vec!["test_table".to_string()]);
11300 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11301 assert_eq!(
11302 describe_response
11303 .properties
11304 .as_ref()
11305 .and_then(|properties| properties.get("owner")),
11306 Some(&"alice".to_string())
11307 );
11308 }
11309
11310 #[tokio::test]
11311 async fn test_create_table_with_manifest_overwrite_replaces_existing_table() {
11312 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
11313
11314 let temp_dir = TempStdDir::default();
11315 let temp_path = temp_dir.to_str().unwrap();
11316
11317 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11318 .manifest_enabled(true)
11319 .dir_listing_enabled(false)
11320 .build()
11321 .await
11322 .unwrap();
11323
11324 let mut create_req = CreateTableRequest::new();
11325 create_req.id = Some(vec!["test_table".to_string()]);
11326 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11327 namespace
11328 .create_table(
11329 create_req,
11330 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11331 )
11332 .await
11333 .unwrap();
11334
11335 let mut create_req = CreateTableRequest::new();
11336 create_req.id = Some(vec!["test_table".to_string()]);
11337 create_req.mode = Some("overwrite".to_string());
11338 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
11339 let response = namespace
11340 .create_table(
11341 create_req,
11342 bytes::Bytes::from(create_single_row_test_ipc_data()),
11343 )
11344 .await
11345 .unwrap();
11346
11347 assert_eq!(response.version, Some(2));
11348 assert_eq!(
11349 response
11350 .properties
11351 .as_ref()
11352 .and_then(|properties| properties.get("owner")),
11353 Some(&"bob".to_string())
11354 );
11355 assert_eq!(
11356 open_dataset(&namespace, "test_table")
11357 .await
11358 .count_rows(None)
11359 .await
11360 .unwrap(),
11361 1
11362 );
11363
11364 let mut describe_req = DescribeTableRequest::new();
11365 describe_req.id = Some(vec!["test_table".to_string()]);
11366 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11367 assert_eq!(
11368 describe_response
11369 .properties
11370 .as_ref()
11371 .and_then(|properties| properties.get("owner")),
11372 Some(&"bob".to_string())
11373 );
11374 }
11375
11376 #[tokio::test]
11377 async fn test_create_table_with_manifest_invalid_mode_rejected() {
11378 use lance_namespace::models::CreateTableRequest;
11379
11380 let temp_dir = TempStdDir::default();
11381 let temp_path = temp_dir.to_str().unwrap();
11382
11383 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11384 .manifest_enabled(true)
11385 .dir_listing_enabled(false)
11386 .build()
11387 .await
11388 .unwrap();
11389
11390 let mut create_req = CreateTableRequest::new();
11391 create_req.id = Some(vec!["test_table".to_string()]);
11392 create_req.mode = Some("append".to_string());
11393 let result = namespace
11394 .create_table(
11395 create_req,
11396 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11397 )
11398 .await;
11399
11400 assert!(result.is_err());
11401 assert!(
11402 result
11403 .unwrap_err()
11404 .to_string()
11405 .contains("Unsupported create_table mode")
11406 );
11407 }
11408
11409 #[tokio::test]
11410 async fn test_merge_insert_into_declared_table_v1_mode_creates_table() {
11411 use lance_namespace::models::{
11412 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11413 MergeInsertIntoTableRequest,
11414 };
11415
11416 let temp_dir = TempStdDir::default();
11417 let temp_path = temp_dir.to_str().unwrap();
11418
11419 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11420 .manifest_enabled(false)
11421 .build()
11422 .await
11423 .unwrap();
11424
11425 let mut declare_req = DeclareTableRequest::new();
11426 declare_req.id = Some(vec!["test_table".to_string()]);
11427 namespace.declare_table(declare_req).await.unwrap();
11428
11429 let mut merge_req = MergeInsertIntoTableRequest::new();
11430 merge_req.id = Some(vec!["test_table".to_string()]);
11431 merge_req.on = Some(vec!["id".to_string()]);
11432 let response = namespace
11433 .merge_insert_into_table(
11434 merge_req,
11435 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11436 )
11437 .await
11438 .unwrap();
11439
11440 assert_eq!(response.num_inserted_rows, Some(2));
11441 assert_eq!(response.num_updated_rows, Some(0));
11442
11443 let mut describe_req = DescribeTableRequest::new();
11444 describe_req.id = Some(vec!["test_table".to_string()]);
11445 describe_req.load_detailed_metadata = Some(true);
11446 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11447 assert_eq!(describe_response.is_only_declared, Some(false));
11448 assert_eq!(describe_response.version, Some(1));
11449
11450 let mut list_req = ListTablesRequest::new();
11451 list_req.id = Some(vec![]);
11452 list_req.include_declared = Some(false);
11453 assert_eq!(
11454 namespace.list_tables(list_req).await.unwrap().tables,
11455 vec!["test_table".to_string()]
11456 );
11457 }
11458
11459 #[tokio::test]
11460 async fn test_merge_insert_into_declared_table_with_manifest_creates_table() {
11461 use lance_namespace::models::{
11462 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11463 MergeInsertIntoTableRequest,
11464 };
11465
11466 let temp_dir = TempStdDir::default();
11467 let temp_path = temp_dir.to_str().unwrap();
11468
11469 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11470 .manifest_enabled(true)
11471 .dir_listing_enabled(false)
11472 .build()
11473 .await
11474 .unwrap();
11475
11476 let mut declare_req = DeclareTableRequest::new();
11477 declare_req.id = Some(vec!["test_table".to_string()]);
11478 namespace.declare_table(declare_req).await.unwrap();
11479
11480 let mut merge_req = MergeInsertIntoTableRequest::new();
11481 merge_req.id = Some(vec!["test_table".to_string()]);
11482 merge_req.on = Some(vec!["id".to_string()]);
11483 let response = namespace
11484 .merge_insert_into_table(
11485 merge_req,
11486 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11487 )
11488 .await
11489 .unwrap();
11490
11491 assert_eq!(response.num_inserted_rows, Some(2));
11492 assert_eq!(response.num_updated_rows, Some(0));
11493
11494 let mut describe_req = DescribeTableRequest::new();
11495 describe_req.id = Some(vec!["test_table".to_string()]);
11496 describe_req.load_detailed_metadata = Some(true);
11497 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11498 assert_eq!(describe_response.is_only_declared, Some(false));
11499 assert_eq!(describe_response.version, Some(1));
11500
11501 let mut list_req = ListTablesRequest::new();
11502 list_req.id = Some(vec![]);
11503 list_req.include_declared = Some(false);
11504 assert_eq!(
11505 namespace.list_tables(list_req).await.unwrap().tables,
11506 vec!["test_table".to_string()]
11507 );
11508 }
11509
11510 fn create_composite_key_ipc_data(rows: &[(Option<&str>, i32, &str)]) -> Vec<u8> {
11514 use arrow::array::{Int32Array, StringArray};
11515 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11516 use arrow::record_batch::RecordBatch;
11517
11518 let schema = Arc::new(ArrowSchema::new(vec![
11519 Field::new("region", DataType::Utf8, true),
11520 Field::new("id", DataType::Int32, false),
11521 Field::new("value", DataType::Utf8, false),
11522 ]));
11523 let batch = RecordBatch::try_new(
11524 schema.clone(),
11525 vec![
11526 Arc::new(StringArray::from_iter(
11527 rows.iter().map(|(region, _, _)| *region),
11528 )),
11529 Arc::new(Int32Array::from_iter_values(
11530 rows.iter().map(|(_, id, _)| *id),
11531 )),
11532 Arc::new(StringArray::from_iter_values(
11533 rows.iter().map(|(_, _, value)| *value),
11534 )),
11535 ],
11536 )
11537 .unwrap();
11538 create_ipc_data_from_batches(schema, vec![batch])
11539 }
11540
11541 async fn read_composite_key_rows(root: &str) -> Vec<(Option<String>, i32, String)> {
11543 use arrow::array::Array;
11544
11545 let dataset = Dataset::open(&format!("{}/test_table.lance", root))
11546 .await
11547 .unwrap();
11548 let batch = dataset.scan().try_into_batch().await.unwrap();
11549 let regions = batch["region"]
11550 .as_any()
11551 .downcast_ref::<arrow::array::StringArray>()
11552 .unwrap();
11553 let ids = batch["id"]
11554 .as_any()
11555 .downcast_ref::<arrow::array::Int32Array>()
11556 .unwrap();
11557 let values = batch["value"]
11558 .as_any()
11559 .downcast_ref::<arrow::array::StringArray>()
11560 .unwrap();
11561
11562 let mut rows: Vec<_> = (0..batch.num_rows())
11563 .map(|row| {
11564 (
11565 regions
11566 .is_valid(row)
11567 .then(|| regions.value(row).to_string()),
11568 ids.value(row),
11569 values.value(row).to_string(),
11570 )
11571 })
11572 .collect();
11573 rows.sort_unstable();
11574 rows
11575 }
11576
11577 #[tokio::test]
11578 async fn test_merge_insert_matches_on_every_column_of_a_composite_key() {
11579 use lance_namespace::models::{DeclareTableRequest, MergeInsertIntoTableRequest};
11580
11581 let temp_dir = TempStdDir::default();
11582 let temp_path = temp_dir.to_str().unwrap();
11583 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11584 .manifest_enabled(false)
11585 .build()
11586 .await
11587 .unwrap();
11588
11589 let mut declare_req = DeclareTableRequest::new();
11590 declare_req.id = Some(vec!["test_table".to_string()]);
11591 namespace.declare_table(declare_req).await.unwrap();
11592
11593 let seed = create_composite_key_ipc_data(&[
11594 (Some("us"), 1, "a"),
11595 (Some("us"), 2, "b"),
11596 (Some("eu"), 1, "c"),
11597 ]);
11598 let mut merge_req = MergeInsertIntoTableRequest::new();
11599 merge_req.id = Some(vec!["test_table".to_string()]);
11600 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11601 namespace
11602 .merge_insert_into_table(merge_req, bytes::Bytes::from(seed))
11603 .await
11604 .unwrap();
11605
11606 let mut merge_req = MergeInsertIntoTableRequest::new();
11609 merge_req.id = Some(vec!["test_table".to_string()]);
11610 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11611 merge_req.when_matched_update_all = Some(true);
11612 let response = namespace
11613 .merge_insert_into_table(
11614 merge_req,
11615 bytes::Bytes::from(create_composite_key_ipc_data(&[
11616 (Some("us"), 1, "updated"),
11617 (Some("eu"), 2, "inserted"),
11618 ])),
11619 )
11620 .await
11621 .unwrap();
11622
11623 assert_eq!(response.num_updated_rows, Some(1));
11624 assert_eq!(response.num_inserted_rows, Some(1));
11625
11626 assert_eq!(
11627 read_composite_key_rows(temp_path).await,
11628 vec![
11629 (Some("eu".into()), 1, "c".into()),
11631 (Some("eu".into()), 2, "inserted".into()),
11632 (Some("us".into()), 1, "updated".into()),
11633 (Some("us".into()), 2, "b".into()),
11634 ]
11635 );
11636 }
11637
11638 #[tokio::test]
11643 async fn test_merge_insert_composite_key_never_matches_a_null_key_column() {
11644 use lance_namespace::models::{DeclareTableRequest, MergeInsertIntoTableRequest};
11645
11646 let temp_dir = TempStdDir::default();
11647 let temp_path = temp_dir.to_str().unwrap();
11648 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11649 .manifest_enabled(false)
11650 .build()
11651 .await
11652 .unwrap();
11653
11654 let mut declare_req = DeclareTableRequest::new();
11655 declare_req.id = Some(vec!["test_table".to_string()]);
11656 namespace.declare_table(declare_req).await.unwrap();
11657
11658 let mut merge_req = MergeInsertIntoTableRequest::new();
11659 merge_req.id = Some(vec!["test_table".to_string()]);
11660 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11661 namespace
11662 .merge_insert_into_table(
11663 merge_req,
11664 bytes::Bytes::from(create_composite_key_ipc_data(&[
11665 (None, 1, "seeded"),
11666 (Some("us"), 1, "us-seeded"),
11667 ])),
11668 )
11669 .await
11670 .unwrap();
11671
11672 let mut merge_req = MergeInsertIntoTableRequest::new();
11673 merge_req.id = Some(vec!["test_table".to_string()]);
11674 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11675 merge_req.when_matched_update_all = Some(true);
11676 let response = namespace
11677 .merge_insert_into_table(
11678 merge_req,
11679 bytes::Bytes::from(create_composite_key_ipc_data(&[(None, 1, "not-a-match")])),
11680 )
11681 .await
11682 .unwrap();
11683
11684 assert_eq!(response.num_updated_rows, Some(0));
11686 assert_eq!(response.num_inserted_rows, Some(1));
11687
11688 assert_eq!(
11689 read_composite_key_rows(temp_path).await,
11690 vec![
11691 (None, 1, "not-a-match".into()),
11692 (None, 1, "seeded".into()),
11693 (Some("us".into()), 1, "us-seeded".into()),
11694 ]
11695 );
11696 }
11697
11698 #[rstest::rstest]
11699 #[case::missing(None, "'on' field is required")]
11700 #[case::empty(Some(vec![]), "must name at least one column")]
11701 #[case::duplicate(
11702 Some(vec!["region".to_string(), "region".to_string()]),
11703 "names column 'region' more than once"
11704 )]
11705 #[tokio::test]
11706 async fn test_merge_insert_rejects_an_invalid_on_key(
11707 #[case] on: Option<Vec<String>>,
11708 #[case] expected_message: &str,
11709 ) {
11710 use lance_namespace::models::{DeclareTableRequest, MergeInsertIntoTableRequest};
11711
11712 let temp_dir = TempStdDir::default();
11713 let temp_path = temp_dir.to_str().unwrap();
11714 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11715 .manifest_enabled(false)
11716 .build()
11717 .await
11718 .unwrap();
11719
11720 let mut declare_req = DeclareTableRequest::new();
11721 declare_req.id = Some(vec!["test_table".to_string()]);
11722 namespace.declare_table(declare_req).await.unwrap();
11723
11724 let mut merge_req = MergeInsertIntoTableRequest::new();
11725 merge_req.id = Some(vec!["test_table".to_string()]);
11726 merge_req.on = on;
11727 let error = namespace
11728 .merge_insert_into_table(
11729 merge_req,
11730 bytes::Bytes::from(create_composite_key_ipc_data(&[(Some("us"), 1, "a")])),
11731 )
11732 .await
11733 .unwrap_err();
11734
11735 let lance_core::Error::Namespace { source, .. } = &error else {
11736 panic!("expected a Namespace error, got: {}", error);
11737 };
11738 let ns_err = source
11739 .downcast_ref::<NamespaceError>()
11740 .expect("expected a NamespaceError source");
11741 assert_eq!(ns_err.code(), lance_namespace::ErrorCode::InvalidInput);
11742 assert!(
11743 error.to_string().contains(expected_message),
11744 "unexpected error message: {error}"
11745 );
11746 }
11747
11748 #[tokio::test]
11749 async fn test_declare_table_with_manifest() {
11750 use lance_namespace::models::{
11751 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
11752 };
11753
11754 let temp_dir = TempStdDir::default();
11755 let temp_path = temp_dir.to_str().unwrap();
11756
11757 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11759 .manifest_enabled(true)
11760 .dir_listing_enabled(false)
11761 .build()
11762 .await
11763 .unwrap();
11764
11765 let mut declare_req = DeclareTableRequest::new();
11767 declare_req.id = Some(vec!["test_table".to_string()]);
11768 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11769 let response = namespace.declare_table(declare_req).await.unwrap();
11770
11771 assert!(response.location.is_some());
11773 assert_eq!(
11774 response
11775 .properties
11776 .as_ref()
11777 .and_then(|properties| properties.get("owner")),
11778 Some(&"alice".to_string())
11779 );
11780
11781 let mut exists_req = TableExistsRequest::new();
11783 exists_req.id = Some(vec!["test_table".to_string()]);
11784 assert!(namespace.table_exists(exists_req).await.is_ok());
11785
11786 let mut describe_req = DescribeTableRequest::new();
11787 describe_req.id = Some(vec!["test_table".to_string()]);
11788 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11789 assert_eq!(describe_response.is_only_declared, None);
11790
11791 let mut describe_req = DescribeTableRequest::new();
11792 describe_req.id = Some(vec!["test_table".to_string()]);
11793 describe_req.check_declared = Some(true);
11794 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11795 assert_eq!(describe_response.is_only_declared, Some(true));
11796 assert_eq!(
11797 describe_response
11798 .properties
11799 .as_ref()
11800 .and_then(|properties| properties.get("owner")),
11801 Some(&"alice".to_string())
11802 );
11803
11804 let mut list_req = ListTablesRequest::new();
11805 list_req.id = Some(vec![]);
11806 assert_eq!(
11807 namespace
11808 .list_tables(list_req.clone())
11809 .await
11810 .unwrap()
11811 .tables,
11812 vec!["test_table".to_string()]
11813 );
11814 list_req.include_declared = Some(false);
11815 assert!(
11816 namespace
11817 .list_tables(list_req)
11818 .await
11819 .unwrap()
11820 .tables
11821 .is_empty()
11822 );
11823 }
11824
11825 #[tokio::test]
11826 async fn test_declare_table_with_manifest_marker_already_exists() {
11827 use lance_namespace::error::ErrorCode;
11830 use lance_namespace::models::DeclareTableRequest;
11831
11832 let temp_dir = TempStdDir::default();
11833 let temp_path = temp_dir.to_str().unwrap();
11834
11835 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11836 .manifest_enabled(true)
11837 .dir_listing_enabled(true)
11838 .build()
11839 .await
11840 .unwrap();
11841
11842 let table_dir = temp_dir.join("test_table.lance");
11843 std::fs::create_dir_all(&table_dir).unwrap();
11844 std::fs::write(table_dir.join(".lance-reserved"), b"reserved").unwrap();
11845
11846 let mut declare_req = DeclareTableRequest::new();
11847 declare_req.id = Some(vec!["test_table".to_string()]);
11848 let err = namespace
11849 .declare_table(declare_req)
11850 .await
11851 .expect_err("declare with existing marker must fail");
11852 let msg = err.to_string();
11853 assert!(
11854 msg.contains("already exists") || msg.contains("TableAlreadyExists"),
11855 "expected TableAlreadyExists, got: {msg}"
11856 );
11857 assert_eq!(
11858 mutation_error_code(err),
11859 ErrorCode::TableAlreadyExists,
11860 "expected TableAlreadyExists error code"
11861 );
11862 }
11863
11864 #[tokio::test]
11865 async fn test_declare_table_when_table_exists() {
11866 use lance_namespace::models::DeclareTableRequest;
11867
11868 let temp_dir = TempStdDir::default();
11869 let temp_path = temp_dir.to_str().unwrap();
11870
11871 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11872 .manifest_enabled(false)
11873 .build()
11874 .await
11875 .unwrap();
11876
11877 let schema = create_test_schema();
11879 let ipc_data = create_test_ipc_data(&schema);
11880 let mut create_req = CreateTableRequest::new();
11881 create_req.id = Some(vec!["test_table".to_string()]);
11882 namespace
11883 .create_table(create_req, bytes::Bytes::from(ipc_data))
11884 .await
11885 .unwrap();
11886
11887 let mut declare_req = DeclareTableRequest::new();
11889 declare_req.id = Some(vec!["test_table".to_string()]);
11890 let result = namespace.declare_table(declare_req).await;
11891 assert!(result.is_err());
11892 }
11893
11894 #[tokio::test]
11899 async fn test_deregister_table_v1_mode() {
11900 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
11901
11902 let temp_dir = TempStdDir::default();
11903 let temp_path = temp_dir.to_str().unwrap();
11904
11905 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11907 .manifest_enabled(false)
11908 .dir_listing_enabled(true)
11909 .build()
11910 .await
11911 .unwrap();
11912
11913 let schema = create_test_schema();
11915 let ipc_data = create_test_ipc_data(&schema);
11916 let mut create_req = CreateTableRequest::new();
11917 create_req.id = Some(vec!["test_table".to_string()]);
11918 namespace
11919 .create_table(create_req, bytes::Bytes::from(ipc_data))
11920 .await
11921 .unwrap();
11922
11923 let mut exists_req = TableExistsRequest::new();
11925 exists_req.id = Some(vec!["test_table".to_string()]);
11926 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
11927
11928 let mut deregister_req = DeregisterTableRequest::new();
11930 deregister_req.id = Some(vec!["test_table".to_string()]);
11931 let response = namespace.deregister_table(deregister_req).await.unwrap();
11932
11933 assert!(response.location.is_some());
11935 let location = response.location.as_ref().unwrap();
11936 assert!(location.contains("test_table"));
11937
11938 let result = namespace.table_exists(exists_req).await;
11940 assert!(result.is_err());
11941 assert!(result.unwrap_err().to_string().contains("deregistered"));
11942
11943 let dataset = Dataset::open(location).await;
11945 assert!(dataset.is_ok(), "Physical table data should still exist");
11946 }
11947
11948 #[tokio::test]
11949 async fn test_deregister_table_v1_already_deregistered() {
11950 use lance_namespace::models::DeregisterTableRequest;
11951
11952 let temp_dir = TempStdDir::default();
11953 let temp_path = temp_dir.to_str().unwrap();
11954
11955 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11956 .manifest_enabled(false)
11957 .dir_listing_enabled(true)
11958 .build()
11959 .await
11960 .unwrap();
11961
11962 let schema = create_test_schema();
11964 let ipc_data = create_test_ipc_data(&schema);
11965 let mut create_req = CreateTableRequest::new();
11966 create_req.id = Some(vec!["test_table".to_string()]);
11967 namespace
11968 .create_table(create_req, bytes::Bytes::from(ipc_data))
11969 .await
11970 .unwrap();
11971
11972 let mut deregister_req = DeregisterTableRequest::new();
11974 deregister_req.id = Some(vec!["test_table".to_string()]);
11975 namespace
11976 .deregister_table(deregister_req.clone())
11977 .await
11978 .unwrap();
11979
11980 let result = namespace.deregister_table(deregister_req).await;
11982 assert!(result.is_err());
11983 assert!(
11984 result
11985 .unwrap_err()
11986 .to_string()
11987 .contains("already deregistered")
11988 );
11989 }
11990
11991 #[tokio::test]
11996 async fn test_list_tables_skips_deregistered_v1() {
11997 use lance_namespace::models::DeregisterTableRequest;
11998
11999 let temp_dir = TempStdDir::default();
12000 let temp_path = temp_dir.to_str().unwrap();
12001
12002 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12003 .manifest_enabled(false)
12004 .dir_listing_enabled(true)
12005 .build()
12006 .await
12007 .unwrap();
12008
12009 let schema = create_test_schema();
12011 let ipc_data = create_test_ipc_data(&schema);
12012
12013 let mut create_req1 = CreateTableRequest::new();
12014 create_req1.id = Some(vec!["table1".to_string()]);
12015 namespace
12016 .create_table(create_req1, bytes::Bytes::from(ipc_data.clone()))
12017 .await
12018 .unwrap();
12019
12020 let mut create_req2 = CreateTableRequest::new();
12021 create_req2.id = Some(vec!["table2".to_string()]);
12022 namespace
12023 .create_table(create_req2, bytes::Bytes::from(ipc_data))
12024 .await
12025 .unwrap();
12026
12027 let mut list_req = ListTablesRequest::new();
12029 list_req.id = Some(vec![]);
12030 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
12031 assert_eq!(list_response.tables.len(), 2);
12032
12033 let mut deregister_req = DeregisterTableRequest::new();
12035 deregister_req.id = Some(vec!["table1".to_string()]);
12036 namespace.deregister_table(deregister_req).await.unwrap();
12037
12038 let list_response = namespace.list_tables(list_req).await.unwrap();
12040 assert_eq!(list_response.tables.len(), 1);
12041 assert!(list_response.tables.contains(&"table2".to_string()));
12042 assert!(!list_response.tables.contains(&"table1".to_string()));
12043 }
12044
12045 #[tokio::test]
12050 async fn test_describe_table_fails_for_deregistered_v1() {
12051 use lance_namespace::models::{DeregisterTableRequest, DescribeTableRequest};
12052
12053 let temp_dir = TempStdDir::default();
12054 let temp_path = temp_dir.to_str().unwrap();
12055
12056 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12057 .manifest_enabled(false)
12058 .dir_listing_enabled(true)
12059 .build()
12060 .await
12061 .unwrap();
12062
12063 let schema = create_test_schema();
12065 let ipc_data = create_test_ipc_data(&schema);
12066 let mut create_req = CreateTableRequest::new();
12067 create_req.id = Some(vec!["test_table".to_string()]);
12068 namespace
12069 .create_table(create_req, bytes::Bytes::from(ipc_data))
12070 .await
12071 .unwrap();
12072
12073 let mut describe_req = DescribeTableRequest::new();
12075 describe_req.id = Some(vec!["test_table".to_string()]);
12076 assert!(namespace.describe_table(describe_req.clone()).await.is_ok());
12077
12078 let mut deregister_req = DeregisterTableRequest::new();
12080 deregister_req.id = Some(vec!["test_table".to_string()]);
12081 namespace.deregister_table(deregister_req).await.unwrap();
12082
12083 let result = namespace.describe_table(describe_req).await;
12085 assert!(result.is_err());
12086 let err = result.unwrap_err();
12087 assert!(matches!(err, Error::Namespace { .. }));
12088 let err_msg = err.to_string();
12089 assert!(err_msg.contains("deregistered"));
12090 assert!(err_msg.contains("table id 'test_table'"));
12091 }
12092
12093 #[tokio::test]
12094 async fn test_table_exists_fails_for_deregistered_v1() {
12095 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
12096
12097 let temp_dir = TempStdDir::default();
12098 let temp_path = temp_dir.to_str().unwrap();
12099
12100 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12101 .manifest_enabled(false)
12102 .dir_listing_enabled(true)
12103 .build()
12104 .await
12105 .unwrap();
12106
12107 let schema = create_test_schema();
12109 let ipc_data = create_test_ipc_data(&schema);
12110 let mut create_req = CreateTableRequest::new();
12111 create_req.id = Some(vec!["test_table".to_string()]);
12112 namespace
12113 .create_table(create_req, bytes::Bytes::from(ipc_data))
12114 .await
12115 .unwrap();
12116
12117 let mut exists_req = TableExistsRequest::new();
12119 exists_req.id = Some(vec!["test_table".to_string()]);
12120 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
12121
12122 let mut deregister_req = DeregisterTableRequest::new();
12124 deregister_req.id = Some(vec!["test_table".to_string()]);
12125 namespace.deregister_table(deregister_req).await.unwrap();
12126
12127 let result = namespace.table_exists(exists_req).await;
12129 assert!(result.is_err());
12130 let err = result.unwrap_err();
12131 assert!(matches!(err, Error::Namespace { .. }));
12132 let err_msg = err.to_string();
12133 assert!(err_msg.contains("deregistered"));
12134 assert!(err_msg.contains("table id 'test_table'"));
12135 }
12136
12137 #[tokio::test]
12138 async fn test_atomic_table_status_check() {
12139 let temp_dir = TempStdDir::default();
12143 let temp_path = temp_dir.to_str().unwrap();
12144
12145 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12146 .manifest_enabled(false)
12147 .dir_listing_enabled(true)
12148 .build()
12149 .await
12150 .unwrap();
12151
12152 let schema = create_test_schema();
12154 let ipc_data = create_test_ipc_data(&schema);
12155 let mut create_req = CreateTableRequest::new();
12156 create_req.id = Some(vec!["test_table".to_string()]);
12157 namespace
12158 .create_table(create_req, bytes::Bytes::from(ipc_data))
12159 .await
12160 .unwrap();
12161
12162 let status = namespace.check_table_status("test_table").await.unwrap();
12164 assert!(status.exists);
12165 assert!(!status.is_deregistered);
12166 assert!(!status.has_reserved_file);
12167 }
12168
12169 #[tokio::test]
12170 async fn test_table_version_tracking_enabled_managed_versioning() {
12171 use lance_namespace::models::DescribeTableRequest;
12172
12173 let temp_dir = TempStdDir::default();
12174 let temp_path = temp_dir.to_str().unwrap();
12175
12176 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12178 .table_version_tracking_enabled(true)
12179 .build()
12180 .await
12181 .unwrap();
12182
12183 let schema = create_test_schema();
12185 let ipc_data = create_test_ipc_data(&schema);
12186 let mut create_req = CreateTableRequest::new();
12187 create_req.id = Some(vec!["test_table".to_string()]);
12188 namespace
12189 .create_table(create_req, bytes::Bytes::from(ipc_data))
12190 .await
12191 .unwrap();
12192
12193 let mut describe_req = DescribeTableRequest::new();
12195 describe_req.id = Some(vec!["test_table".to_string()]);
12196 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
12197
12198 assert_eq!(
12200 describe_resp.managed_versioning,
12201 Some(true),
12202 "managed_versioning should be true when table_version_tracking_enabled=true"
12203 );
12204 }
12205
12206 #[tokio::test]
12207 async fn test_table_version_tracking_disabled_no_managed_versioning() {
12208 use lance_namespace::models::DescribeTableRequest;
12209
12210 let temp_dir = TempStdDir::default();
12211 let temp_path = temp_dir.to_str().unwrap();
12212
12213 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12215 .table_version_tracking_enabled(false)
12216 .build()
12217 .await
12218 .unwrap();
12219
12220 let schema = create_test_schema();
12222 let ipc_data = create_test_ipc_data(&schema);
12223 let mut create_req = CreateTableRequest::new();
12224 create_req.id = Some(vec!["test_table".to_string()]);
12225 namespace
12226 .create_table(create_req, bytes::Bytes::from(ipc_data))
12227 .await
12228 .unwrap();
12229
12230 let mut describe_req = DescribeTableRequest::new();
12232 describe_req.id = Some(vec!["test_table".to_string()]);
12233 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
12234
12235 assert!(
12237 describe_resp.managed_versioning.is_none(),
12238 "managed_versioning should be None when table_version_tracking_enabled=false, got: {:?}",
12239 describe_resp.managed_versioning
12240 );
12241 }
12242
12243 #[tokio::test]
12244 async fn test_list_table_versions() {
12245 use arrow::array::{Int32Array, RecordBatchIterator};
12246 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12247 use arrow::record_batch::RecordBatch;
12248 use lance::dataset::{Dataset, WriteMode, WriteParams};
12249 use lance_namespace::models::{CreateNamespaceRequest, ListTableVersionsRequest};
12250
12251 let temp_dir = TempStrDir::default();
12252 let temp_path: &str = &temp_dir;
12253
12254 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12255 DirectoryNamespaceBuilder::new(temp_path)
12256 .table_version_tracking_enabled(true)
12257 .build()
12258 .await
12259 .unwrap(),
12260 );
12261
12262 let mut create_ns_req = CreateNamespaceRequest::new();
12264 create_ns_req.id = Some(vec!["workspace".to_string()]);
12265 namespace.create_namespace(create_ns_req).await.unwrap();
12266
12267 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12269 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12270 "id",
12271 DataType::Int32,
12272 false,
12273 )]));
12274 let batch = RecordBatch::try_new(
12275 arrow_schema.clone(),
12276 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12277 )
12278 .unwrap();
12279 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
12280 let write_params = WriteParams {
12281 mode: WriteMode::Create,
12282 ..Default::default()
12283 };
12284 let mut dataset = Dataset::write_into_namespace(
12285 batches,
12286 namespace.clone(),
12287 table_id.clone(),
12288 Some(write_params),
12289 )
12290 .await
12291 .unwrap();
12292
12293 let batch2 = RecordBatch::try_new(
12295 arrow_schema.clone(),
12296 vec![Arc::new(Int32Array::from(vec![100, 200]))],
12297 )
12298 .unwrap();
12299 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
12300 dataset.append(batches, None).await.unwrap();
12301
12302 let batch3 = RecordBatch::try_new(
12304 arrow_schema.clone(),
12305 vec![Arc::new(Int32Array::from(vec![300, 400]))],
12306 )
12307 .unwrap();
12308 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
12309 dataset.append(batches, None).await.unwrap();
12310
12311 let mut list_req = ListTableVersionsRequest::new();
12313 list_req.id = Some(table_id.clone());
12314 let list_resp = namespace.list_table_versions(list_req).await.unwrap();
12315
12316 assert_eq!(
12317 list_resp.versions.len(),
12318 3,
12319 "Should have 3 versions, got: {:?}",
12320 list_resp.versions
12321 );
12322
12323 for expected_version in 1..=3 {
12325 let version = list_resp
12326 .versions
12327 .iter()
12328 .find(|v| v.version == expected_version)
12329 .unwrap_or_else(|| panic!("Expected version {}", expected_version));
12330
12331 assert!(
12332 !version.manifest_path.is_empty(),
12333 "manifest_path should be set for version {}",
12334 expected_version
12335 );
12336 assert!(
12337 version.manifest_path.contains(".manifest"),
12338 "manifest_path should contain .manifest for version {}",
12339 expected_version
12340 );
12341 assert!(
12342 version.manifest_size.is_some(),
12343 "manifest_size should be set for version {}",
12344 expected_version
12345 );
12346 assert!(
12347 version.manifest_size.unwrap() > 0,
12348 "manifest_size should be > 0 for version {}",
12349 expected_version
12350 );
12351 assert!(
12352 version.timestamp_millis.is_some(),
12353 "timestamp_millis should be set for version {}",
12354 expected_version
12355 );
12356 }
12357 }
12358
12359 #[tokio::test]
12360 async fn test_describe_table_version() {
12361 use arrow::array::{Int32Array, RecordBatchIterator};
12362 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12363 use arrow::record_batch::RecordBatch;
12364 use lance::dataset::{Dataset, WriteMode, WriteParams};
12365 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
12366
12367 let temp_dir = TempStrDir::default();
12368 let temp_path: &str = &temp_dir;
12369
12370 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12371 DirectoryNamespaceBuilder::new(temp_path)
12372 .table_version_tracking_enabled(true)
12373 .build()
12374 .await
12375 .unwrap(),
12376 );
12377
12378 let mut create_ns_req = CreateNamespaceRequest::new();
12380 create_ns_req.id = Some(vec!["workspace".to_string()]);
12381 namespace.create_namespace(create_ns_req).await.unwrap();
12382
12383 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12385 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12386 "id",
12387 DataType::Int32,
12388 false,
12389 )]));
12390 let batch = RecordBatch::try_new(
12391 arrow_schema.clone(),
12392 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12393 )
12394 .unwrap();
12395 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12396 let write_params = WriteParams {
12397 mode: WriteMode::Create,
12398 ..Default::default()
12399 };
12400 let mut dataset = Dataset::write_into_namespace(
12401 batches,
12402 namespace.clone(),
12403 table_id.clone(),
12404 Some(write_params),
12405 )
12406 .await
12407 .unwrap();
12408
12409 let batch2 = RecordBatch::try_new(
12411 arrow_schema.clone(),
12412 vec![Arc::new(Int32Array::from(vec![100, 200]))],
12413 )
12414 .unwrap();
12415 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
12416 dataset.append(batches, None).await.unwrap();
12417
12418 let mut describe_req = DescribeTableVersionRequest::new();
12420 describe_req.id = Some(table_id.clone());
12421 describe_req.version = Some(1);
12422 let describe_resp = namespace
12423 .describe_table_version(describe_req)
12424 .await
12425 .unwrap();
12426
12427 let version = &describe_resp.version;
12428 assert_eq!(version.version, 1);
12429 assert!(version.timestamp_millis.is_some());
12430 assert!(
12431 !version.manifest_path.is_empty(),
12432 "manifest_path should be set"
12433 );
12434 assert!(
12435 version.manifest_path.contains(".manifest"),
12436 "manifest_path should contain .manifest"
12437 );
12438 assert!(
12439 version.manifest_size.is_some(),
12440 "manifest_size should be set"
12441 );
12442 assert!(
12443 version.manifest_size.unwrap() > 0,
12444 "manifest_size should be > 0"
12445 );
12446
12447 let mut describe_req = DescribeTableVersionRequest::new();
12449 describe_req.id = Some(table_id.clone());
12450 describe_req.version = Some(2);
12451 let describe_resp = namespace
12452 .describe_table_version(describe_req)
12453 .await
12454 .unwrap();
12455
12456 let version = &describe_resp.version;
12457 assert_eq!(version.version, 2);
12458 assert!(version.timestamp_millis.is_some());
12459 assert!(
12460 !version.manifest_path.is_empty(),
12461 "manifest_path should be set"
12462 );
12463 assert!(
12464 version.manifest_size.is_some(),
12465 "manifest_size should be set"
12466 );
12467 assert!(
12468 version.manifest_size.unwrap() > 0,
12469 "manifest_size should be > 0"
12470 );
12471 }
12472
12473 #[tokio::test]
12474 async fn test_describe_table_version_latest() {
12475 use arrow::array::{Int32Array, RecordBatchIterator};
12476 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12477 use arrow::record_batch::RecordBatch;
12478 use lance::dataset::{Dataset, WriteMode, WriteParams};
12479 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
12480
12481 let temp_dir = TempStrDir::default();
12482 let temp_path: &str = &temp_dir;
12483
12484 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12485 DirectoryNamespaceBuilder::new(temp_path)
12486 .table_version_tracking_enabled(true)
12487 .build()
12488 .await
12489 .unwrap(),
12490 );
12491
12492 let mut create_ns_req = CreateNamespaceRequest::new();
12494 create_ns_req.id = Some(vec!["workspace".to_string()]);
12495 namespace.create_namespace(create_ns_req).await.unwrap();
12496
12497 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12499 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12500 "id",
12501 DataType::Int32,
12502 false,
12503 )]));
12504 let batch = RecordBatch::try_new(
12505 arrow_schema.clone(),
12506 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12507 )
12508 .unwrap();
12509 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12510 let write_params = WriteParams {
12511 mode: WriteMode::Create,
12512 ..Default::default()
12513 };
12514 let mut dataset = Dataset::write_into_namespace(
12515 batches,
12516 namespace.clone(),
12517 table_id.clone(),
12518 Some(write_params),
12519 )
12520 .await
12521 .unwrap();
12522
12523 let batch2 = RecordBatch::try_new(
12525 arrow_schema.clone(),
12526 vec![Arc::new(Int32Array::from(vec![100, 200]))],
12527 )
12528 .unwrap();
12529 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
12530 dataset.append(batches, None).await.unwrap();
12531
12532 let batch3 = RecordBatch::try_new(
12534 arrow_schema.clone(),
12535 vec![Arc::new(Int32Array::from(vec![300, 400]))],
12536 )
12537 .unwrap();
12538 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
12539 dataset.append(batches, None).await.unwrap();
12540
12541 let mut describe_req = DescribeTableVersionRequest::new();
12543 describe_req.id = Some(table_id.clone());
12544 describe_req.version = None;
12545 let describe_resp = namespace
12546 .describe_table_version(describe_req)
12547 .await
12548 .unwrap();
12549
12550 assert_eq!(describe_resp.version.version, 3);
12552 }
12553
12554 #[tokio::test]
12555 async fn test_create_table_version() {
12556 use futures::TryStreamExt;
12557 use lance::dataset::builder::DatasetBuilder;
12558 use lance_namespace::models::CreateTableVersionRequest;
12559
12560 let temp_dir = TempStrDir::default();
12561 let temp_path: &str = &temp_dir;
12562
12563 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12564 DirectoryNamespaceBuilder::new(temp_path)
12565 .table_version_tracking_enabled(true)
12566 .build()
12567 .await
12568 .unwrap(),
12569 );
12570
12571 let schema = create_test_schema();
12573 let ipc_data = create_test_ipc_data(&schema);
12574 let mut create_req = CreateTableRequest::new();
12575 create_req.id = Some(vec!["test_table".to_string()]);
12576 namespace
12577 .create_table(create_req, bytes::Bytes::from(ipc_data))
12578 .await
12579 .unwrap();
12580
12581 let table_id = vec!["test_table".to_string()];
12583 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12584 .await
12585 .unwrap()
12586 .load()
12587 .await
12588 .unwrap();
12589
12590 let versions_path = dataset.versions_dir();
12592 let manifest_metas: Vec<_> = dataset
12593 .object_store(None)
12594 .await
12595 .unwrap()
12596 .inner
12597 .list(Some(&versions_path))
12598 .try_collect()
12599 .await
12600 .unwrap();
12601
12602 let manifest_meta = manifest_metas
12603 .iter()
12604 .find(|m| {
12605 m.location
12606 .filename()
12607 .map(|f| f.ends_with(".manifest"))
12608 .unwrap_or(false)
12609 })
12610 .expect("No manifest file found");
12611
12612 let manifest_data = dataset
12614 .object_store(None)
12615 .await
12616 .unwrap()
12617 .inner
12618 .get(&manifest_meta.location)
12619 .await
12620 .unwrap()
12621 .bytes()
12622 .await
12623 .unwrap();
12624
12625 let staging_path = dataset.versions_dir().join("staging_manifest");
12627 dataset
12628 .object_store(None)
12629 .await
12630 .unwrap()
12631 .inner
12632 .put(&staging_path, manifest_data.into())
12633 .await
12634 .unwrap();
12635
12636 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12639 create_version_req.id = Some(table_id.clone());
12640 create_version_req.naming_scheme = Some("V2".to_string());
12641
12642 let result = namespace.create_table_version(create_version_req).await;
12643 assert!(
12644 result.is_ok(),
12645 "create_table_version should succeed: {:?}",
12646 result
12647 );
12648
12649 let response = result.unwrap();
12651 let version_info = response
12652 .version
12653 .expect("response should contain version info");
12654 let version_2_path = Path::parse(&version_info.manifest_path).unwrap();
12655 let head_result = dataset
12656 .object_store(None)
12657 .await
12658 .unwrap()
12659 .inner
12660 .head(&version_2_path)
12661 .await;
12662 assert!(
12663 head_result.is_ok(),
12664 "Version 2 manifest should exist at {}",
12665 version_2_path
12666 );
12667
12668 let staging_head_result = dataset
12670 .object_store(None)
12671 .await
12672 .unwrap()
12673 .inner
12674 .head(&staging_path)
12675 .await;
12676 assert!(
12677 staging_head_result.is_err(),
12678 "Staging manifest should have been deleted after create_table_version"
12679 );
12680 }
12681
12682 #[tokio::test]
12683 async fn test_create_table_version_idempotent() {
12684 use futures::TryStreamExt;
12687 use lance::dataset::builder::DatasetBuilder;
12688 use lance_namespace::models::CreateTableVersionRequest;
12689
12690 let temp_dir = TempStrDir::default();
12691 let temp_path: &str = &temp_dir;
12692
12693 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12694 DirectoryNamespaceBuilder::new(temp_path)
12695 .table_version_tracking_enabled(true)
12696 .build()
12697 .await
12698 .unwrap(),
12699 );
12700
12701 let schema = create_test_schema();
12702 let ipc_data = create_test_ipc_data(&schema);
12703 let mut create_req = CreateTableRequest::new();
12704 create_req.id = Some(vec!["test_table".to_string()]);
12705 namespace
12706 .create_table(create_req, bytes::Bytes::from(ipc_data))
12707 .await
12708 .unwrap();
12709
12710 let table_id = vec!["test_table".to_string()];
12711 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12712 .await
12713 .unwrap()
12714 .load()
12715 .await
12716 .unwrap();
12717
12718 let versions_path = dataset.versions_dir();
12719 let manifest_metas: Vec<_> = dataset
12720 .object_store(None)
12721 .await
12722 .unwrap()
12723 .inner
12724 .list(Some(&versions_path))
12725 .try_collect()
12726 .await
12727 .unwrap();
12728
12729 let manifest_meta = manifest_metas
12730 .iter()
12731 .find(|m| {
12732 m.location
12733 .filename()
12734 .map(|f| f.ends_with(".manifest"))
12735 .unwrap_or(false)
12736 })
12737 .expect("No manifest file found");
12738
12739 let manifest_data = dataset
12740 .object_store(None)
12741 .await
12742 .unwrap()
12743 .inner
12744 .get(&manifest_meta.location)
12745 .await
12746 .unwrap()
12747 .bytes()
12748 .await
12749 .unwrap();
12750
12751 let staging_path = dataset.versions_dir().join("staging_manifest");
12752 dataset
12753 .object_store(None)
12754 .await
12755 .unwrap()
12756 .inner
12757 .put(&staging_path, manifest_data.clone().into())
12758 .await
12759 .unwrap();
12760
12761 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12762 create_version_req.id = Some(table_id.clone());
12763 create_version_req.naming_scheme = Some("V2".to_string());
12764 let first = namespace
12765 .create_table_version(create_version_req)
12766 .await
12767 .expect("first create_table_version should succeed");
12768
12769 let retry_staging = dataset.versions_dir().join("staging_manifest_retry");
12771 dataset
12772 .object_store(None)
12773 .await
12774 .unwrap()
12775 .inner
12776 .put(&retry_staging, manifest_data.into())
12777 .await
12778 .unwrap();
12779
12780 let mut retry_req = CreateTableVersionRequest::new(2, retry_staging.to_string());
12781 retry_req.id = Some(table_id.clone());
12782 retry_req.naming_scheme = Some("V2".to_string());
12783 let second = namespace
12784 .create_table_version(retry_req)
12785 .await
12786 .expect("idempotent retry must succeed");
12787
12788 assert_eq!(
12789 first.version.as_ref().map(|v| v.version),
12790 second.version.as_ref().map(|v| v.version)
12791 );
12792 assert_eq!(
12793 first.version.as_ref().map(|v| &v.manifest_path),
12794 second.version.as_ref().map(|v| &v.manifest_path)
12795 );
12796 }
12797
12798 #[tokio::test]
12799 async fn test_create_table_version_conflict() {
12800 use futures::TryStreamExt;
12802 use lance::dataset::builder::DatasetBuilder;
12803 use lance_namespace::models::CreateTableVersionRequest;
12804
12805 let temp_dir = TempStrDir::default();
12806 let temp_path: &str = &temp_dir;
12807
12808 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12809 DirectoryNamespaceBuilder::new(temp_path)
12810 .table_version_tracking_enabled(true)
12811 .build()
12812 .await
12813 .unwrap(),
12814 );
12815
12816 let schema = create_test_schema();
12818 let ipc_data = create_test_ipc_data(&schema);
12819 let mut create_req = CreateTableRequest::new();
12820 create_req.id = Some(vec!["test_table".to_string()]);
12821 namespace
12822 .create_table(create_req, bytes::Bytes::from(ipc_data))
12823 .await
12824 .unwrap();
12825
12826 let table_id = vec!["test_table".to_string()];
12828 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12829 .await
12830 .unwrap()
12831 .load()
12832 .await
12833 .unwrap();
12834
12835 let versions_path = dataset.versions_dir();
12837 let manifest_metas: Vec<_> = dataset
12838 .object_store(None)
12839 .await
12840 .unwrap()
12841 .inner
12842 .list(Some(&versions_path))
12843 .try_collect()
12844 .await
12845 .unwrap();
12846
12847 let manifest_meta = manifest_metas
12848 .iter()
12849 .find(|m| {
12850 m.location
12851 .filename()
12852 .map(|f| f.ends_with(".manifest"))
12853 .unwrap_or(false)
12854 })
12855 .expect("No manifest file found");
12856
12857 let manifest_data = dataset
12859 .object_store(None)
12860 .await
12861 .unwrap()
12862 .inner
12863 .get(&manifest_meta.location)
12864 .await
12865 .unwrap()
12866 .bytes()
12867 .await
12868 .unwrap();
12869
12870 let staging_path = dataset.versions_dir().join("staging_manifest");
12872 dataset
12873 .object_store(None)
12874 .await
12875 .unwrap()
12876 .inner
12877 .put(&staging_path, manifest_data.into())
12878 .await
12879 .unwrap();
12880
12881 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12883 create_version_req.id = Some(table_id.clone());
12884 create_version_req.naming_scheme = Some("V2".to_string());
12885 let first_result = namespace.create_table_version(create_version_req).await;
12886 assert!(
12887 first_result.is_ok(),
12888 "First create_table_version for version 2 should succeed: {:?}",
12889 first_result
12890 );
12891
12892 let version_2_path = Path::parse(
12894 &first_result
12895 .unwrap()
12896 .version
12897 .expect("response should contain version info")
12898 .manifest_path,
12899 )
12900 .unwrap();
12901
12902 let conflict_staging = dataset.versions_dir().join("staging_manifest_conflict");
12904 dataset
12905 .object_store(None)
12906 .await
12907 .unwrap()
12908 .inner
12909 .put(
12910 &conflict_staging,
12911 bytes::Bytes::from_static(b"not-a-real-manifest").into(),
12912 )
12913 .await
12914 .unwrap();
12915
12916 let mut create_version_req =
12917 CreateTableVersionRequest::new(2, conflict_staging.to_string());
12918 create_version_req.id = Some(table_id.clone());
12919 create_version_req.naming_scheme = Some("V2".to_string());
12920
12921 let result = namespace.create_table_version(create_version_req).await;
12922 assert!(
12923 result.is_err(),
12924 "create_table_version should fail for existing version with different content"
12925 );
12926 let err = result.unwrap_err().to_string();
12927 assert!(
12928 err.contains("already exists") || err.contains("ConcurrentModification"),
12929 "expected ConcurrentModification, got: {err}"
12930 );
12931
12932 let head_result = dataset
12934 .object_store(None)
12935 .await
12936 .unwrap()
12937 .inner
12938 .head(&version_2_path)
12939 .await;
12940 assert!(
12941 head_result.is_ok(),
12942 "Version 2 manifest should still exist at {}",
12943 version_2_path
12944 );
12945 }
12946
12947 #[tokio::test]
12948 async fn test_create_table_version_cas_rejects_gap() {
12949 use futures::TryStreamExt;
12951 use lance::dataset::builder::DatasetBuilder;
12952 use lance_namespace::models::CreateTableVersionRequest;
12953
12954 let temp_dir = TempStrDir::default();
12955 let temp_path: &str = &temp_dir;
12956
12957 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12958 DirectoryNamespaceBuilder::new(temp_path)
12959 .table_version_tracking_enabled(true)
12960 .build()
12961 .await
12962 .unwrap(),
12963 );
12964
12965 let schema = create_test_schema();
12966 let ipc_data = create_test_ipc_data(&schema);
12967 let mut create_req = CreateTableRequest::new();
12968 create_req.id = Some(vec!["test_table".to_string()]);
12969 namespace
12970 .create_table(create_req, bytes::Bytes::from(ipc_data))
12971 .await
12972 .unwrap();
12973
12974 let table_id = vec!["test_table".to_string()];
12975 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12976 .await
12977 .unwrap()
12978 .load()
12979 .await
12980 .unwrap();
12981
12982 let versions_path = dataset.versions_dir();
12983 let manifest_metas: Vec<_> = dataset
12984 .object_store(None)
12985 .await
12986 .unwrap()
12987 .inner
12988 .list(Some(&versions_path))
12989 .try_collect()
12990 .await
12991 .unwrap();
12992 let manifest_meta = manifest_metas
12993 .iter()
12994 .find(|m| {
12995 m.location
12996 .filename()
12997 .map(|f| f.ends_with(".manifest"))
12998 .unwrap_or(false)
12999 })
13000 .expect("No manifest file found");
13001 let manifest_data = dataset
13002 .object_store(None)
13003 .await
13004 .unwrap()
13005 .inner
13006 .get(&manifest_meta.location)
13007 .await
13008 .unwrap()
13009 .bytes()
13010 .await
13011 .unwrap();
13012
13013 let staging_path = dataset.versions_dir().join("staging_gap");
13014 dataset
13015 .object_store(None)
13016 .await
13017 .unwrap()
13018 .inner
13019 .put(&staging_path, manifest_data.into())
13020 .await
13021 .unwrap();
13022
13023 let mut req = CreateTableVersionRequest::new(5, staging_path.to_string());
13025 req.id = Some(table_id);
13026 req.naming_scheme = Some("V2".to_string());
13027 let err = namespace
13028 .create_table_version(req)
13029 .await
13030 .expect_err("gap create must fail CAS");
13031 let msg = err.to_string();
13032 assert!(
13033 msg.contains("CAS") || msg.contains("ConcurrentModification"),
13034 "expected CAS ConcurrentModification, got: {msg}"
13035 );
13036 }
13037
13038 #[tokio::test]
13039 async fn test_create_table_version_branch_cas_requires_parent_version() {
13040 use futures::TryStreamExt;
13043 use lance_namespace::models::CreateTableVersionRequest;
13044
13045 let (namespace, _temp_dir) = create_test_namespace().await;
13046 create_scalar_table(&namespace, "users").await;
13047 let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
13048 append_scalar_version(&main_uri, 10).await; let mut main = open_dataset(&namespace, "users").await;
13051 let fork_version = main.version().version;
13052 assert_eq!(fork_version, 2);
13053 let branch_uri = main
13054 .create_branch("exp", fork_version, None)
13055 .await
13056 .unwrap()
13057 .uri()
13058 .to_string();
13059
13060 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
13061 let versions_dir = branch_ds.versions_dir();
13062 let store = branch_ds.object_store(None).await.unwrap();
13063 let manifests: Vec<_> = store
13064 .inner
13065 .list(Some(&versions_dir))
13066 .try_collect()
13067 .await
13068 .unwrap();
13069 for meta in &manifests {
13070 if meta
13071 .location
13072 .filename()
13073 .is_some_and(|f| f.ends_with(".manifest"))
13074 {
13075 store.inner.delete(&meta.location).await.unwrap();
13076 }
13077 }
13078 let remaining_manifests = store
13081 .inner
13082 .list(Some(&versions_dir))
13083 .try_collect::<Vec<_>>()
13084 .await
13085 .unwrap()
13086 .into_iter()
13087 .filter(|m| {
13088 m.location
13089 .filename()
13090 .is_some_and(|f| f.ends_with(".manifest"))
13091 })
13092 .count();
13093 assert_eq!(
13094 remaining_manifests, 0,
13095 "branch version chain should be empty after deleting manifests"
13096 );
13097
13098 let main_ds = open_dataset(&namespace, "users").await;
13100 let main_versions = main_ds.versions_dir();
13101 let main_store = main_ds.object_store(None).await.unwrap();
13102 let source_meta = main_store
13103 .inner
13104 .list(Some(&main_versions))
13105 .try_collect::<Vec<_>>()
13106 .await
13107 .unwrap()
13108 .into_iter()
13109 .find(|m| {
13110 m.location
13111 .filename()
13112 .is_some_and(|f| f.ends_with(".manifest"))
13113 })
13114 .expect("main should have a manifest");
13115 let source_bytes = main_store
13116 .inner
13117 .get(&source_meta.location)
13118 .await
13119 .unwrap()
13120 .bytes()
13121 .await
13122 .unwrap();
13123
13124 let staging_wrong = versions_dir.clone().join("staging_wrong");
13125 store
13126 .inner
13127 .put(&staging_wrong, source_bytes.clone().into())
13128 .await
13129 .unwrap();
13130 let err = namespace
13131 .create_table_version(CreateTableVersionRequest {
13132 id: Some(vec!["users".to_string()]),
13133 version: 1,
13134 manifest_path: staging_wrong.to_string(),
13135 naming_scheme: Some("V2".to_string()),
13136 branch: Some("exp".to_string()),
13137 ..Default::default()
13138 })
13139 .await
13140 .expect_err("bootstrap at v1 must fail when parent_version is 2");
13141 let msg = err.to_string();
13142 assert!(
13143 msg.contains("CAS") || msg.contains("ConcurrentModification"),
13144 "expected CAS ConcurrentModification, got: {msg}"
13145 );
13146
13147 let staging_ok = versions_dir.join("staging_ok");
13148 store
13149 .inner
13150 .put(&staging_ok, source_bytes.into())
13151 .await
13152 .unwrap();
13153 let resp = namespace
13154 .create_table_version(CreateTableVersionRequest {
13155 id: Some(vec!["users".to_string()]),
13156 version: 2,
13157 manifest_path: staging_ok.to_string(),
13158 naming_scheme: Some("V2".to_string()),
13159 branch: Some("exp".to_string()),
13160 ..Default::default()
13161 })
13162 .await
13163 .expect("bootstrap at parent_version must succeed");
13164 assert_eq!(resp.version.as_ref().map(|v| v.version), Some(2));
13165 }
13166
13167 #[tokio::test]
13168 async fn test_create_table_version_table_not_found() {
13169 use lance_namespace::models::CreateTableVersionRequest;
13170
13171 let temp_dir = TempStdDir::default();
13172 let temp_path = temp_dir.to_str().unwrap();
13173
13174 let namespace = DirectoryNamespaceBuilder::new(temp_path)
13175 .table_version_tracking_enabled(true)
13176 .build()
13177 .await
13178 .unwrap();
13179
13180 let mut create_version_req =
13182 CreateTableVersionRequest::new(1, "/some/staging/path".to_string());
13183 create_version_req.id = Some(vec!["non_existent_table".to_string()]);
13184
13185 let result = namespace.create_table_version(create_version_req).await;
13186 assert!(
13187 result.is_err(),
13188 "create_table_version should fail for non-existent table"
13189 );
13190 let err_msg = result.unwrap_err().to_string();
13191 assert!(
13192 err_msg.contains("Table not found"),
13193 "Error should mention table not found, got: {}",
13194 err_msg
13195 );
13196 }
13197
13198 mod e2e_table_version_tracking {
13200 use super::*;
13201 use std::sync::atomic::{AtomicUsize, Ordering};
13202
13203 struct TrackingNamespace {
13205 inner: DirectoryNamespace,
13206 create_table_version_count: AtomicUsize,
13207 describe_table_version_count: AtomicUsize,
13208 list_table_versions_count: AtomicUsize,
13209 }
13210
13211 impl TrackingNamespace {
13212 fn new(inner: DirectoryNamespace) -> Self {
13213 Self {
13214 inner,
13215 create_table_version_count: AtomicUsize::new(0),
13216 describe_table_version_count: AtomicUsize::new(0),
13217 list_table_versions_count: AtomicUsize::new(0),
13218 }
13219 }
13220
13221 fn create_table_version_calls(&self) -> usize {
13222 self.create_table_version_count.load(Ordering::SeqCst)
13223 }
13224
13225 fn describe_table_version_calls(&self) -> usize {
13226 self.describe_table_version_count.load(Ordering::SeqCst)
13227 }
13228
13229 fn list_table_versions_calls(&self) -> usize {
13230 self.list_table_versions_count.load(Ordering::SeqCst)
13231 }
13232 }
13233
13234 impl std::fmt::Debug for TrackingNamespace {
13235 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
13236 f.debug_struct("TrackingNamespace")
13237 .field(
13238 "create_table_version_calls",
13239 &self.create_table_version_calls(),
13240 )
13241 .finish()
13242 }
13243 }
13244
13245 #[async_trait]
13246 impl LanceNamespace for TrackingNamespace {
13247 async fn create_namespace(
13248 &self,
13249 request: CreateNamespaceRequest,
13250 ) -> Result<CreateNamespaceResponse> {
13251 self.inner.create_namespace(request).await
13252 }
13253
13254 async fn describe_namespace(
13255 &self,
13256 request: DescribeNamespaceRequest,
13257 ) -> Result<DescribeNamespaceResponse> {
13258 self.inner.describe_namespace(request).await
13259 }
13260
13261 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
13262 self.inner.namespace_exists(request).await
13263 }
13264
13265 async fn list_namespaces(
13266 &self,
13267 request: ListNamespacesRequest,
13268 ) -> Result<ListNamespacesResponse> {
13269 self.inner.list_namespaces(request).await
13270 }
13271
13272 async fn drop_namespace(
13273 &self,
13274 request: DropNamespaceRequest,
13275 ) -> Result<DropNamespaceResponse> {
13276 self.inner.drop_namespace(request).await
13277 }
13278
13279 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
13280 self.inner.list_tables(request).await
13281 }
13282
13283 async fn describe_table(
13284 &self,
13285 request: DescribeTableRequest,
13286 ) -> Result<DescribeTableResponse> {
13287 self.inner.describe_table(request).await
13288 }
13289
13290 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
13291 self.inner.table_exists(request).await
13292 }
13293
13294 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
13295 self.inner.drop_table(request).await
13296 }
13297
13298 async fn create_table(
13299 &self,
13300 request: CreateTableRequest,
13301 request_data: Bytes,
13302 ) -> Result<CreateTableResponse> {
13303 self.inner.create_table(request, request_data).await
13304 }
13305
13306 async fn declare_table(
13307 &self,
13308 request: DeclareTableRequest,
13309 ) -> Result<DeclareTableResponse> {
13310 self.inner.declare_table(request).await
13311 }
13312
13313 async fn list_table_versions(
13314 &self,
13315 request: ListTableVersionsRequest,
13316 ) -> Result<ListTableVersionsResponse> {
13317 self.list_table_versions_count
13318 .fetch_add(1, Ordering::SeqCst);
13319 self.inner.list_table_versions(request).await
13320 }
13321
13322 async fn create_table_version(
13323 &self,
13324 request: CreateTableVersionRequest,
13325 ) -> Result<CreateTableVersionResponse> {
13326 self.create_table_version_count
13327 .fetch_add(1, Ordering::SeqCst);
13328 self.inner.create_table_version(request).await
13329 }
13330
13331 async fn describe_table_version(
13332 &self,
13333 request: DescribeTableVersionRequest,
13334 ) -> Result<DescribeTableVersionResponse> {
13335 self.describe_table_version_count
13336 .fetch_add(1, Ordering::SeqCst);
13337 self.inner.describe_table_version(request).await
13338 }
13339
13340 async fn batch_delete_table_versions(
13341 &self,
13342 request: BatchDeleteTableVersionsRequest,
13343 ) -> Result<BatchDeleteTableVersionsResponse> {
13344 self.inner.batch_delete_table_versions(request).await
13345 }
13346
13347 fn namespace_id(&self) -> String {
13348 self.inner.namespace_id()
13349 }
13350 }
13351
13352 #[tokio::test]
13353 async fn test_describe_table_returns_managed_versioning() {
13354 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableRequest};
13355
13356 let temp_dir = TempStdDir::default();
13357 let temp_path = temp_dir.to_str().unwrap();
13358
13359 let ns = DirectoryNamespaceBuilder::new(temp_path)
13361 .table_version_tracking_enabled(true)
13362 .manifest_enabled(true)
13363 .build()
13364 .await
13365 .unwrap();
13366
13367 let mut create_ns_req = CreateNamespaceRequest::new();
13369 create_ns_req.id = Some(vec!["workspace".to_string()]);
13370 ns.create_namespace(create_ns_req).await.unwrap();
13371
13372 let schema = create_test_schema();
13374 let ipc_data = create_test_ipc_data(&schema);
13375 let mut create_req = CreateTableRequest::new();
13376 create_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
13377 ns.create_table(create_req, bytes::Bytes::from(ipc_data))
13378 .await
13379 .unwrap();
13380
13381 let mut describe_req = DescribeTableRequest::new();
13383 describe_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
13384 let describe_resp = ns.describe_table(describe_req).await.unwrap();
13385
13386 assert_eq!(
13388 describe_resp.managed_versioning,
13389 Some(true),
13390 "managed_versioning should be true when table_version_tracking_enabled=true"
13391 );
13392 }
13393
13394 #[tokio::test]
13395 async fn test_external_manifest_store_invokes_namespace_apis() {
13396 use arrow::array::{Int32Array, StringArray};
13397 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
13398 use arrow::record_batch::RecordBatch;
13399 use lance::Dataset;
13400 use lance::dataset::builder::DatasetBuilder;
13401 use lance::dataset::{WriteMode, WriteParams};
13402 use lance_namespace::models::CreateNamespaceRequest;
13403
13404 let temp_dir = TempStdDir::default();
13405 let temp_path = temp_dir.to_str().unwrap();
13406
13407 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
13409 .table_version_tracking_enabled(true)
13410 .manifest_enabled(true)
13411 .build()
13412 .await
13413 .unwrap();
13414
13415 let tracking_ns = Arc::new(TrackingNamespace::new(inner_ns));
13416 let ns: Arc<dyn LanceNamespace> = tracking_ns.clone();
13417
13418 let mut create_ns_req = CreateNamespaceRequest::new();
13420 create_ns_req.id = Some(vec!["workspace".to_string()]);
13421 ns.create_namespace(create_ns_req).await.unwrap();
13422
13423 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
13425
13426 let arrow_schema = Arc::new(ArrowSchema::new(vec![
13428 Field::new("id", DataType::Int32, false),
13429 Field::new("name", DataType::Utf8, true),
13430 ]));
13431 let batch = RecordBatch::try_new(
13432 arrow_schema.clone(),
13433 vec![
13434 Arc::new(Int32Array::from(vec![1, 2, 3])),
13435 Arc::new(StringArray::from(vec!["a", "b", "c"])),
13436 ],
13437 )
13438 .unwrap();
13439
13440 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
13442 let write_params = WriteParams {
13443 mode: WriteMode::Create,
13444 ..Default::default()
13445 };
13446 let mut dataset = Dataset::write_into_namespace(
13447 batches,
13448 ns.clone(),
13449 table_id.clone(),
13450 Some(write_params),
13451 )
13452 .await
13453 .unwrap();
13454 assert_eq!(dataset.version().version, 1);
13455
13456 assert_eq!(
13458 tracking_ns.create_table_version_calls(),
13459 1,
13460 "create_table_version should have been called once during initial write_into_namespace"
13461 );
13462
13463 let append_batch = RecordBatch::try_new(
13465 arrow_schema.clone(),
13466 vec![
13467 Arc::new(Int32Array::from(vec![4, 5, 6])),
13468 Arc::new(StringArray::from(vec!["d", "e", "f"])),
13469 ],
13470 )
13471 .unwrap();
13472 let append_batches = RecordBatchIterator::new(vec![Ok(append_batch)], arrow_schema);
13473 dataset.append(append_batches, None).await.unwrap();
13474
13475 assert_eq!(
13476 tracking_ns.create_table_version_calls(),
13477 2,
13478 "create_table_version should have been called twice (once for create, once for append)"
13479 );
13480
13481 let initial_list_calls = tracking_ns.list_table_versions_calls();
13483 let latest_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
13484 .await
13485 .unwrap()
13486 .load()
13487 .await
13488 .unwrap();
13489 assert_eq!(latest_dataset.version().version, 2);
13490 assert_eq!(
13491 tracking_ns.list_table_versions_calls(),
13492 initial_list_calls + 1,
13493 "list_table_versions should have been called exactly once during checkout_latest"
13494 );
13495
13496 let initial_describe_calls = tracking_ns.describe_table_version_calls();
13498 let v1_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
13499 .await
13500 .unwrap()
13501 .with_version(1)
13502 .load()
13503 .await
13504 .unwrap();
13505 assert_eq!(v1_dataset.version().version, 1);
13506 assert_eq!(
13507 tracking_ns.describe_table_version_calls(),
13508 initial_describe_calls + 1,
13509 "describe_table_version should have been called exactly once during checkout to version 1"
13510 );
13511 }
13512
13513 #[tokio::test]
13514 async fn test_dataset_commit_with_external_manifest_store() {
13515 use arrow::array::{Int32Array, StringArray};
13516 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
13517 use arrow::record_batch::RecordBatch;
13518 use futures::TryStreamExt;
13519 use lance::dataset::{Dataset, WriteMode, WriteParams};
13520 use lance_namespace::models::CreateNamespaceRequest;
13521 use lance_table::io::commit::ManifestNamingScheme;
13522
13523 let temp_dir = TempStdDir::default();
13524 let temp_path = temp_dir.to_str().unwrap();
13525
13526 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
13528 .table_version_tracking_enabled(true)
13529 .manifest_enabled(true)
13530 .build()
13531 .await
13532 .unwrap();
13533
13534 let tracking_ns: Arc<dyn LanceNamespace> = Arc::new(TrackingNamespace::new(inner_ns));
13535
13536 let mut create_ns_req = CreateNamespaceRequest::new();
13538 create_ns_req.id = Some(vec!["workspace".to_string()]);
13539 tracking_ns.create_namespace(create_ns_req).await.unwrap();
13540
13541 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
13543 let arrow_schema = Arc::new(ArrowSchema::new(vec![
13544 Field::new("id", DataType::Int32, false),
13545 Field::new("name", DataType::Utf8, true),
13546 ]));
13547 let batch = RecordBatch::try_new(
13548 arrow_schema.clone(),
13549 vec![
13550 Arc::new(Int32Array::from(vec![1, 2, 3])),
13551 Arc::new(StringArray::from(vec!["a", "b", "c"])),
13552 ],
13553 )
13554 .unwrap();
13555 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
13556 let write_params = WriteParams {
13557 mode: WriteMode::Create,
13558 ..Default::default()
13559 };
13560 let dataset = Dataset::write_into_namespace(
13561 batches,
13562 tracking_ns.clone(),
13563 table_id.clone(),
13564 Some(write_params),
13565 )
13566 .await
13567 .unwrap();
13568 assert_eq!(dataset.version().version, 1);
13569
13570 let batch2 = RecordBatch::try_new(
13572 arrow_schema.clone(),
13573 vec![
13574 Arc::new(Int32Array::from(vec![4, 5, 6])),
13575 Arc::new(StringArray::from(vec!["d", "e", "f"])),
13576 ],
13577 )
13578 .unwrap();
13579 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
13580 let write_params = WriteParams {
13581 mode: WriteMode::Append,
13582 ..Default::default()
13583 };
13584 Dataset::write_into_namespace(
13585 batches,
13586 tracking_ns.clone(),
13587 table_id.clone(),
13588 Some(write_params),
13589 )
13590 .await
13591 .unwrap();
13592
13593 let manifest_metas: Vec<_> = dataset
13596 .object_store(None)
13597 .await
13598 .unwrap()
13599 .inner
13600 .list(Some(&dataset.versions_dir()))
13601 .try_collect()
13602 .await
13603 .unwrap();
13604 let version_2_found = manifest_metas.iter().any(|m| {
13605 m.location
13606 .filename()
13607 .map(|f| {
13608 f.ends_with(".manifest")
13609 && ManifestNamingScheme::V2.parse_version(f) == Some(2)
13610 })
13611 .unwrap_or(false)
13612 });
13613 assert!(
13614 version_2_found,
13615 "Version 2 manifest should exist in versions directory"
13616 );
13617 }
13618
13619 async fn create_ns_with_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
13621 use arrow::array::{Int32Array, StringArray};
13622 use arrow::ipc::writer::StreamWriter;
13623
13624 let (namespace, temp_dir) = create_test_namespace().await;
13625
13626 let schema = create_test_schema();
13627 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13628 let arrow_schema = Arc::new(arrow_schema);
13629
13630 let id_array = Int32Array::from(vec![1, 2, 3]);
13631 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
13632 let batch = arrow::record_batch::RecordBatch::try_new(
13633 arrow_schema.clone(),
13634 vec![Arc::new(id_array), Arc::new(name_array)],
13635 )
13636 .unwrap();
13637
13638 let mut buffer = Vec::new();
13639 {
13640 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13641 writer.write(&batch).unwrap();
13642 writer.finish().unwrap();
13643 }
13644
13645 let mut request = CreateTableRequest::new();
13646 let table_id = vec!["test_ops_table".to_string()];
13647 request.id = Some(table_id.clone());
13648
13649 namespace
13650 .create_table(request, Bytes::from(buffer))
13651 .await
13652 .unwrap();
13653
13654 (namespace, temp_dir, table_id)
13655 }
13656
13657 #[tokio::test]
13658 async fn test_count_table_rows_basic() {
13659 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13660
13661 let request = CountTableRowsRequest {
13662 id: Some(table_id),
13663 version: None,
13664 predicate: None,
13665 ..Default::default()
13666 };
13667
13668 let count = namespace.count_table_rows(request).await.unwrap();
13669 assert_eq!(count, 3);
13670 }
13671
13672 #[tokio::test]
13673 async fn test_count_table_rows_with_predicate() {
13674 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13675
13676 let request = CountTableRowsRequest {
13677 id: Some(table_id),
13678 version: None,
13679 predicate: Some("id > 1".to_string()),
13680 ..Default::default()
13681 };
13682
13683 let count = namespace.count_table_rows(request).await.unwrap();
13684 assert_eq!(count, 2);
13685 }
13686
13687 #[tokio::test]
13688 async fn test_query_table_invalid_distance_type() {
13689 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13690
13691 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13692 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13693 multi_vector: None,
13694 });
13695
13696 let request = QueryTableRequest {
13697 id: Some(table_id),
13698 k: 2,
13699 vector,
13700 vector_column: Some("vector".to_string()),
13701 distance_type: Some("invalid_metric".to_string()),
13702 filter: None,
13703 offset: None,
13704 version: None,
13705 ..Default::default()
13706 };
13707
13708 let result = namespace.query_table(request).await;
13709 assert!(result.is_err());
13710 let err_msg = result.unwrap_err().to_string();
13711 assert!(
13712 err_msg.contains("Unknown distance type"),
13713 "Expected error about unknown distance type, got: {}",
13714 err_msg
13715 );
13716 }
13717
13718 #[tokio::test]
13719 async fn test_insert_into_table_append() {
13720 use arrow::array::{Int32Array, StringArray};
13721 use arrow::ipc::writer::StreamWriter;
13722
13723 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13724
13725 let schema = create_test_schema();
13727 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13728 let arrow_schema = Arc::new(arrow_schema);
13729
13730 let id_array = Int32Array::from(vec![4, 5]);
13731 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13732 let batch = arrow::record_batch::RecordBatch::try_new(
13733 arrow_schema.clone(),
13734 vec![Arc::new(id_array), Arc::new(name_array)],
13735 )
13736 .unwrap();
13737
13738 let mut buffer = Vec::new();
13739 {
13740 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13741 writer.write(&batch).unwrap();
13742 writer.finish().unwrap();
13743 }
13744
13745 let request = InsertIntoTableRequest {
13746 id: Some(table_id.clone()),
13747 mode: Some("append".to_string()),
13748 ..Default::default()
13749 };
13750
13751 let response = namespace
13752 .insert_into_table(request, Bytes::from(buffer))
13753 .await
13754 .unwrap();
13755 assert!(response.transaction_id.is_none());
13756
13757 let count_req = CountTableRowsRequest {
13759 id: Some(table_id),
13760 version: None,
13761 predicate: None,
13762 ..Default::default()
13763 };
13764 let count = namespace.count_table_rows(count_req).await.unwrap();
13765 assert_eq!(count, 5);
13766 }
13767
13768 #[tokio::test]
13769 async fn test_insert_into_table_overwrite() {
13770 use arrow::array::{Int32Array, StringArray};
13771 use arrow::ipc::writer::StreamWriter;
13772
13773 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13774
13775 let schema = create_test_schema();
13776 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13777 let arrow_schema = Arc::new(arrow_schema);
13778
13779 let id_array = Int32Array::from(vec![10, 20]);
13780 let name_array = StringArray::from(vec!["X", "Y"]);
13781 let batch = arrow::record_batch::RecordBatch::try_new(
13782 arrow_schema.clone(),
13783 vec![Arc::new(id_array), Arc::new(name_array)],
13784 )
13785 .unwrap();
13786
13787 let mut buffer = Vec::new();
13788 {
13789 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13790 writer.write(&batch).unwrap();
13791 writer.finish().unwrap();
13792 }
13793
13794 let request = InsertIntoTableRequest {
13795 id: Some(table_id.clone()),
13796 mode: Some("overwrite".to_string()),
13797 ..Default::default()
13798 };
13799
13800 namespace
13801 .insert_into_table(request, Bytes::from(buffer))
13802 .await
13803 .unwrap();
13804
13805 let count_req = CountTableRowsRequest {
13807 id: Some(table_id),
13808 version: None,
13809 predicate: None,
13810 ..Default::default()
13811 };
13812 let count = namespace.count_table_rows(count_req).await.unwrap();
13813 assert_eq!(count, 2);
13814 }
13815
13816 #[tokio::test]
13817 async fn test_insert_into_table_empty_data() {
13818 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13819
13820 let request = InsertIntoTableRequest {
13821 id: Some(table_id),
13822 mode: None,
13823 ..Default::default()
13824 };
13825
13826 let result = namespace.insert_into_table(request, Bytes::new()).await;
13827 assert!(result.is_err());
13828 assert!(
13829 result
13830 .unwrap_err()
13831 .to_string()
13832 .contains("Arrow IPC stream) is required")
13833 );
13834 }
13835
13836 #[tokio::test]
13837 async fn test_insert_into_table_with_storage_options() {
13838 use arrow::array::{Int32Array, StringArray};
13839 use arrow::ipc::writer::StreamWriter;
13840
13841 let temp_dir = TempStdDir::default();
13842
13843 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
13845 .storage_option("allow_http", "true")
13846 .build()
13847 .await
13848 .unwrap();
13849
13850 let schema = create_test_schema();
13852 let ipc_data = create_test_ipc_data(&schema);
13853 let mut create_req = CreateTableRequest::new();
13854 let table_id = vec!["so_table".to_string()];
13855 create_req.id = Some(table_id.clone());
13856 namespace
13857 .create_table(create_req, Bytes::from(ipc_data))
13858 .await
13859 .unwrap();
13860
13861 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13863 let arrow_schema = Arc::new(arrow_schema);
13864
13865 let id_array = Int32Array::from(vec![10, 20]);
13866 let name_array = StringArray::from(vec!["X", "Y"]);
13867 let batch = arrow::record_batch::RecordBatch::try_new(
13868 arrow_schema.clone(),
13869 vec![Arc::new(id_array), Arc::new(name_array)],
13870 )
13871 .unwrap();
13872
13873 let mut buffer = Vec::new();
13874 {
13875 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13876 writer.write(&batch).unwrap();
13877 writer.finish().unwrap();
13878 }
13879
13880 let request = InsertIntoTableRequest {
13881 id: Some(table_id.clone()),
13882 mode: Some("append".to_string()),
13883 ..Default::default()
13884 };
13885
13886 let response = namespace
13887 .insert_into_table(request, Bytes::from(buffer))
13888 .await
13889 .unwrap();
13890 assert!(response.transaction_id.is_none());
13891
13892 let count_req = CountTableRowsRequest {
13894 id: Some(table_id),
13895 version: None,
13896 predicate: None,
13897 ..Default::default()
13898 };
13899 let count = namespace.count_table_rows(count_req).await.unwrap();
13900 assert_eq!(count, 2);
13901 }
13902
13903 #[tokio::test]
13904 async fn test_query_table_basic() {
13905 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13906
13907 let request = QueryTableRequest {
13908 id: Some(table_id),
13909 k: 10,
13910 filter: None,
13911 offset: None,
13912 version: None,
13913 ..Default::default()
13914 };
13915
13916 let bytes = namespace.query_table(request).await.unwrap();
13917
13918 let cursor = Cursor::new(bytes.to_vec());
13920 let reader = FileReader::try_new(cursor, None).unwrap();
13921 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13922 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13923 assert_eq!(total_rows, 3);
13924 }
13925
13926 #[tokio::test]
13927 async fn test_query_table_with_filter() {
13928 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13929
13930 let request = QueryTableRequest {
13931 id: Some(table_id),
13932 k: 10,
13933 filter: Some("id <= 2".to_string()),
13934 offset: None,
13935 version: None,
13936 ..Default::default()
13937 };
13938
13939 let bytes = namespace.query_table(request).await.unwrap();
13940
13941 let cursor = Cursor::new(bytes.to_vec());
13942 let reader = FileReader::try_new(cursor, None).unwrap();
13943 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13944 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13945 assert_eq!(total_rows, 2);
13946 }
13947
13948 #[tokio::test]
13949 async fn test_query_table_with_limit_and_offset() {
13950 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13951
13952 let request = QueryTableRequest {
13953 id: Some(table_id),
13954 k: 2,
13955 filter: None,
13956 offset: Some(1),
13957 version: None,
13958 ..Default::default()
13959 };
13960
13961 let bytes = namespace.query_table(request).await.unwrap();
13962
13963 let cursor = Cursor::new(bytes.to_vec());
13964 let reader = FileReader::try_new(cursor, None).unwrap();
13965 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13966 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13967 assert_eq!(total_rows, 2);
13968 }
13969
13970 #[tokio::test]
13971 async fn test_query_table_no_limit() {
13972 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13973
13974 let request = QueryTableRequest {
13976 id: Some(table_id),
13977 k: 0,
13978 filter: None,
13979 offset: None,
13980 version: None,
13981 ..Default::default()
13982 };
13983
13984 let bytes = namespace.query_table(request).await.unwrap();
13985
13986 let cursor = Cursor::new(bytes.to_vec());
13987 let reader = FileReader::try_new(cursor, None).unwrap();
13988 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13989 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13990 assert_eq!(total_rows, 3);
13991 }
13992
13993 #[tokio::test]
13994 async fn test_query_table_with_columns() {
13995 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13996
13997 let columns = Box::new(lance_namespace::models::QueryTableRequestColumns {
13998 column_names: Some(vec!["id".to_string()]),
13999 column_aliases: None,
14000 });
14001
14002 let request = QueryTableRequest {
14003 id: Some(table_id),
14004 k: 10,
14005 filter: None,
14006 offset: None,
14007 version: None,
14008 columns: Some(columns),
14009 ..Default::default()
14010 };
14011
14012 let bytes = namespace.query_table(request).await.unwrap();
14013
14014 let cursor = Cursor::new(bytes.to_vec());
14015 let reader = FileReader::try_new(cursor, None).unwrap();
14016 let schema = reader.schema();
14017 assert_eq!(schema.fields().len(), 1);
14018 assert_eq!(schema.field(0).name(), "id");
14019 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14020 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14021 assert_eq!(total_rows, 3);
14022 }
14023
14024 #[tokio::test]
14025 async fn test_count_table_rows_with_version() {
14026 use arrow::array::{Int32Array, StringArray};
14027 use arrow::ipc::writer::StreamWriter;
14028
14029 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14030
14031 let schema = create_test_schema();
14033 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
14034 let arrow_schema = Arc::new(arrow_schema);
14035
14036 let id_array = Int32Array::from(vec![4, 5]);
14037 let name_array = StringArray::from(vec!["Dave", "Eve"]);
14038 let batch = arrow::record_batch::RecordBatch::try_new(
14039 arrow_schema.clone(),
14040 vec![Arc::new(id_array), Arc::new(name_array)],
14041 )
14042 .unwrap();
14043
14044 let mut buffer = Vec::new();
14045 {
14046 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
14047 writer.write(&batch).unwrap();
14048 writer.finish().unwrap();
14049 }
14050
14051 let request = InsertIntoTableRequest {
14052 id: Some(table_id.clone()),
14053 mode: None,
14054 ..Default::default()
14055 };
14056 namespace
14057 .insert_into_table(request, Bytes::from(buffer))
14058 .await
14059 .unwrap();
14060
14061 let count_req = CountTableRowsRequest {
14063 id: Some(table_id.clone()),
14064 version: Some(1),
14065 predicate: None,
14066 ..Default::default()
14067 };
14068 let count = namespace.count_table_rows(count_req).await.unwrap();
14069 assert_eq!(count, 3);
14070
14071 let count_req = CountTableRowsRequest {
14073 id: Some(table_id),
14074 version: None,
14075 predicate: None,
14076 ..Default::default()
14077 };
14078 let count = namespace.count_table_rows(count_req).await.unwrap();
14079 assert_eq!(count, 5);
14080 }
14081
14082 #[tokio::test]
14083 async fn test_query_table_with_version() {
14084 use arrow::array::{Int32Array, StringArray};
14085 use arrow::ipc::writer::StreamWriter;
14086
14087 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14088
14089 let schema = create_test_schema();
14091 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
14092 let arrow_schema = Arc::new(arrow_schema);
14093
14094 let id_array = Int32Array::from(vec![4, 5]);
14095 let name_array = StringArray::from(vec!["Dave", "Eve"]);
14096 let batch = arrow::record_batch::RecordBatch::try_new(
14097 arrow_schema.clone(),
14098 vec![Arc::new(id_array), Arc::new(name_array)],
14099 )
14100 .unwrap();
14101
14102 let mut buffer = Vec::new();
14103 {
14104 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
14105 writer.write(&batch).unwrap();
14106 writer.finish().unwrap();
14107 }
14108
14109 let request = InsertIntoTableRequest {
14110 id: Some(table_id.clone()),
14111 mode: None,
14112 ..Default::default()
14113 };
14114 namespace
14115 .insert_into_table(request, Bytes::from(buffer))
14116 .await
14117 .unwrap();
14118
14119 let request = QueryTableRequest {
14121 id: Some(table_id.clone()),
14122 k: 100,
14123 filter: None,
14124 offset: None,
14125 version: Some(1),
14126 ..Default::default()
14127 };
14128
14129 let bytes = namespace.query_table(request).await.unwrap();
14130 let cursor = Cursor::new(bytes.to_vec());
14131 let reader = FileReader::try_new(cursor, None).unwrap();
14132 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14133 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14134 assert_eq!(total_rows, 3);
14135
14136 let request = QueryTableRequest {
14138 id: Some(table_id),
14139 k: 100,
14140 filter: None,
14141 offset: None,
14142 version: None,
14143 ..Default::default()
14144 };
14145
14146 let bytes = namespace.query_table(request).await.unwrap();
14147 let cursor = Cursor::new(bytes.to_vec());
14148 let reader = FileReader::try_new(cursor, None).unwrap();
14149 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14150 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14151 assert_eq!(total_rows, 5);
14152 }
14153
14154 async fn create_ns_with_vector_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
14157 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
14158 use arrow::ipc::writer::StreamWriter;
14159
14160 let (namespace, temp_dir) = create_test_namespace().await;
14161
14162 let arrow_schema = Arc::new(arrow::datatypes::Schema::new(vec![
14164 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
14165 arrow::datatypes::Field::new(
14166 "vector",
14167 arrow::datatypes::DataType::FixedSizeList(
14168 Arc::new(arrow::datatypes::Field::new(
14169 "item",
14170 arrow::datatypes::DataType::Float32,
14171 true,
14172 )),
14173 4,
14174 ),
14175 true,
14176 ),
14177 ]));
14178
14179 let id_array = Int32Array::from(vec![1, 2, 3]);
14180 let values = Float32Array::from(vec![
14181 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ]);
14185 let vector_array = FixedSizeListArray::try_new(
14186 Arc::new(arrow::datatypes::Field::new(
14187 "item",
14188 arrow::datatypes::DataType::Float32,
14189 true,
14190 )),
14191 4,
14192 Arc::new(values),
14193 None,
14194 )
14195 .unwrap();
14196
14197 let batch = arrow::record_batch::RecordBatch::try_new(
14198 arrow_schema.clone(),
14199 vec![Arc::new(id_array), Arc::new(vector_array)],
14200 )
14201 .unwrap();
14202
14203 let mut buffer = Vec::new();
14204 {
14205 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
14206 writer.write(&batch).unwrap();
14207 writer.finish().unwrap();
14208 }
14209
14210 let table_name = "vector_table";
14212 let table_uri = format!("{}/{}.lance", temp_dir.to_str().unwrap(), table_name);
14213 let reader = arrow::record_batch::RecordBatchIterator::new(
14214 vec![Ok(batch)],
14215 arrow_schema.clone(),
14216 );
14217 Dataset::write(reader, &table_uri, None).await.unwrap();
14218
14219 let table_id = vec![table_name.to_string()];
14220 (namespace, temp_dir, table_id)
14221 }
14222
14223 #[tokio::test]
14224 async fn test_query_table_vector_search() {
14225 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14226
14227 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14228 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
14229 multi_vector: None,
14230 });
14231
14232 let request = QueryTableRequest {
14233 id: Some(table_id),
14234 k: 2,
14235 vector,
14236 filter: None,
14237 offset: None,
14238 version: None,
14239 ..Default::default()
14240 };
14241
14242 let bytes = namespace.query_table(request).await.unwrap();
14243
14244 let cursor = Cursor::new(bytes.to_vec());
14245 let reader = FileReader::try_new(cursor, None).unwrap();
14246 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14247 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14248 assert_eq!(total_rows, 2);
14249 }
14250
14251 #[tokio::test]
14252 async fn test_query_table_vector_search_with_distance_type() {
14253 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14254
14255 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14256 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
14257 multi_vector: None,
14258 });
14259
14260 let request = QueryTableRequest {
14261 id: Some(table_id),
14262 k: 3,
14263 vector,
14264 filter: None,
14265 offset: None,
14266 version: None,
14267 distance_type: Some("cosine".to_string()),
14268 ..Default::default()
14269 };
14270
14271 let bytes = namespace.query_table(request).await.unwrap();
14272
14273 let cursor = Cursor::new(bytes.to_vec());
14274 let reader = FileReader::try_new(cursor, None).unwrap();
14275 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14276 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14277 assert_eq!(total_rows, 3);
14278 }
14279
14280 #[tokio::test]
14281 async fn test_query_table_vector_search_with_filter() {
14282 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14283
14284 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14285 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
14286 multi_vector: None,
14287 });
14288
14289 let request = QueryTableRequest {
14290 id: Some(table_id),
14291 k: 10,
14292 vector,
14293 filter: Some("id <= 2".to_string()),
14294 offset: None,
14295 version: None,
14296 ..Default::default()
14297 };
14298
14299 let bytes = namespace.query_table(request).await.unwrap();
14300
14301 let cursor = Cursor::new(bytes.to_vec());
14302 let reader = FileReader::try_new(cursor, None).unwrap();
14303 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14304 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14305 assert!(total_rows <= 2);
14306 }
14307
14308 #[tokio::test]
14309 async fn test_query_table_vector_search_with_nprobes_and_refine() {
14310 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14311
14312 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14313 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
14314 multi_vector: None,
14315 });
14316
14317 let request = QueryTableRequest {
14318 id: Some(table_id),
14319 k: 2,
14320 vector,
14321 filter: None,
14322 offset: None,
14323 version: None,
14324 nprobes: Some(1),
14325 refine_factor: Some(1),
14326 prefilter: Some(true),
14327 ..Default::default()
14328 };
14329
14330 let bytes = namespace.query_table(request).await.unwrap();
14331
14332 let cursor = Cursor::new(bytes.to_vec());
14333 let reader = FileReader::try_new(cursor, None).unwrap();
14334 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14335 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14336 assert_eq!(total_rows, 2);
14337 }
14338
14339 #[tokio::test]
14340 async fn test_namespace_id() {
14341 let (namespace, _temp_dir) = create_test_namespace().await;
14342 let id = namespace.namespace_id();
14343 assert!(id.contains("DirectoryNamespace"));
14344 assert!(id.contains("root"));
14345 }
14346
14347 #[tokio::test]
14348 async fn test_query_table_empty_table() {
14349 let (namespace, _temp_dir) = create_test_namespace().await;
14350
14351 let schema = create_test_schema();
14353 let ipc_data = create_test_ipc_data(&schema);
14354 let mut create_request = CreateTableRequest::new();
14355 create_request.id = Some(vec!["empty_table".to_string()]);
14356 namespace
14357 .create_table(create_request, bytes::Bytes::from(ipc_data))
14358 .await
14359 .unwrap();
14360
14361 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14363 single_vector: None,
14364 multi_vector: None,
14365 });
14366 let request = QueryTableRequest {
14367 id: Some(vec!["empty_table".to_string()]),
14368 k: 10,
14369 vector,
14370 ..Default::default()
14371 };
14372 let bytes = namespace.query_table(request).await.unwrap();
14373
14374 let cursor = Cursor::new(bytes.to_vec());
14375 let reader = FileReader::try_new(cursor, None).unwrap();
14376 let batches: Vec<_> = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
14377 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14378 assert_eq!(total_rows, 0, "empty table should yield no rows");
14379 }
14380
14381 #[tokio::test]
14382 async fn test_query_table_with_plain_filter_no_vector() {
14383 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14384
14385 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14387 single_vector: None,
14388 multi_vector: None,
14389 });
14390 let request = QueryTableRequest {
14391 id: Some(table_id),
14392 k: 0,
14393 vector,
14394 filter: Some("id > 1".to_string()),
14395 ..Default::default()
14396 };
14397 let bytes = namespace.query_table(request).await.unwrap();
14398
14399 let cursor = Cursor::new(bytes.to_vec());
14400 let reader = FileReader::try_new(cursor, None).unwrap();
14401 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14402 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14403 assert!(total_rows > 0);
14404 assert!(total_rows < 3);
14405 }
14406
14407 #[tokio::test]
14410 async fn test_update_full_table() {
14411 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14412
14413 let base_version = open_dataset(&namespace, &table_id[0])
14415 .await
14416 .version()
14417 .version;
14418
14419 let request = UpdateTableRequest {
14420 id: Some(table_id.clone()),
14421 updates: vec![vec!["name".to_string(), "'updated'".to_string()]],
14422 predicate: None,
14423 ..Default::default()
14424 };
14425
14426 let response = namespace.update_table(request).await.unwrap();
14427 assert_eq!(response.updated_rows, 3);
14428 assert!(response.version as u64 > base_version);
14429
14430 let count_req = CountTableRowsRequest {
14432 id: Some(table_id),
14433 version: None,
14434 predicate: Some("name = 'updated'".to_string()),
14435 ..Default::default()
14436 };
14437 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 3);
14438 }
14439
14440 #[tokio::test]
14441 async fn test_update_with_predicate() {
14442 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14443
14444 let request = UpdateTableRequest {
14445 id: Some(table_id.clone()),
14446 updates: vec![vec!["name".to_string(), "'matched'".to_string()]],
14447 predicate: Some("id > 1".to_string()),
14448 ..Default::default()
14449 };
14450
14451 let response = namespace.update_table(request).await.unwrap();
14452 assert_eq!(response.updated_rows, 2);
14453
14454 let untouched = CountTableRowsRequest {
14456 id: Some(table_id.clone()),
14457 version: None,
14458 predicate: Some("name = 'Alice'".to_string()),
14459 ..Default::default()
14460 };
14461 assert_eq!(namespace.count_table_rows(untouched).await.unwrap(), 1);
14462
14463 let touched = CountTableRowsRequest {
14464 id: Some(table_id),
14465 version: None,
14466 predicate: Some("name = 'matched'".to_string()),
14467 ..Default::default()
14468 };
14469 assert_eq!(namespace.count_table_rows(touched).await.unwrap(), 2);
14470 }
14471
14472 #[tokio::test]
14473 async fn test_update_invalid_expression_returns_invalid_input() {
14474 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14475
14476 let request = UpdateTableRequest {
14477 id: Some(table_id),
14478 updates: vec![vec!["name".to_string(), "no_such_column + 1".to_string()]],
14480 predicate: None,
14481 ..Default::default()
14482 };
14483
14484 let err = namespace.update_table(request).await.unwrap_err();
14485 let msg = err.to_string();
14486 assert!(
14487 msg.contains("Invalid input"),
14488 "expected Invalid input error, got: {}",
14489 msg
14490 );
14491 }
14492
14493 #[tokio::test]
14494 async fn test_update_rejects_duplicate_columns() {
14495 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14496
14497 let request = UpdateTableRequest {
14498 id: Some(table_id),
14499 updates: vec![
14500 vec!["name".to_string(), "'a'".to_string()],
14501 vec!["name".to_string(), "'b'".to_string()],
14502 ],
14503 predicate: None,
14504 ..Default::default()
14505 };
14506
14507 let err = namespace.update_table(request).await.unwrap_err();
14508 let msg = err.to_string();
14509 assert!(
14510 msg.contains("Invalid input") && msg.contains("more than once"),
14511 "expected duplicate column InvalidInput error, got: {}",
14512 msg
14513 );
14514 }
14515
14516 #[tokio::test]
14517 async fn test_delete_with_predicate() {
14518 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14519
14520 let request = DeleteFromTableRequest {
14521 id: Some(table_id.clone()),
14522 predicate: "id > 1".to_string(),
14523 ..Default::default()
14524 };
14525
14526 let response = namespace.delete_from_table(request).await.unwrap();
14527 assert!(response.version.is_some());
14528
14529 let count_req = CountTableRowsRequest {
14530 id: Some(table_id),
14531 version: None,
14532 predicate: None,
14533 ..Default::default()
14534 };
14535 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 1);
14537 }
14538
14539 #[tokio::test]
14540 async fn test_delete_empty_predicate_returns_invalid_input() {
14541 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14542
14543 let request = DeleteFromTableRequest {
14544 id: Some(table_id),
14545 predicate: " ".to_string(),
14546 ..Default::default()
14547 };
14548
14549 let err = namespace.delete_from_table(request).await.unwrap_err();
14550 let msg = err.to_string();
14551 assert!(
14552 msg.contains("Invalid input") && msg.contains("non-empty predicate"),
14553 "expected non-empty predicate InvalidInput error, got: {}",
14554 msg
14555 );
14556 }
14557
14558 #[tokio::test]
14559 async fn test_delete_table_not_found() {
14560 let (namespace, _temp_dir) = create_test_namespace().await;
14561
14562 let request = DeleteFromTableRequest {
14563 id: Some(vec!["does_not_exist".to_string()]),
14564 predicate: "id = 1".to_string(),
14565 ..Default::default()
14566 };
14567
14568 let err = namespace.delete_from_table(request).await.unwrap_err();
14569 let msg = err.to_string();
14570 assert!(
14571 msg.contains("Table not found"),
14572 "expected TableNotFound for missing table, got: {}",
14573 msg
14574 );
14575 }
14576
14577 #[tokio::test]
14578 async fn test_delete_invalid_predicate_returns_invalid_input() {
14579 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14580
14581 let request = DeleteFromTableRequest {
14585 id: Some(table_id),
14586 predicate: "no_such_column = 1".to_string(),
14587 ..Default::default()
14588 };
14589
14590 let err = namespace.delete_from_table(request).await.unwrap_err();
14591 let lance_core::Error::Namespace { source, .. } = &err else {
14592 panic!("expected a Namespace error, got: {}", err);
14593 };
14594 let ns_err = source
14595 .downcast_ref::<NamespaceError>()
14596 .expect("expected a NamespaceError source");
14597 assert_eq!(
14598 ns_err.code(),
14599 lance_namespace::ErrorCode::InvalidInput,
14600 "expected InvalidInput for an invalid delete predicate, got: {}",
14601 err
14602 );
14603 }
14604 }
14605
14606 #[tokio::test]
14607 async fn test_list_all_tables() {
14608 use lance_namespace::models::ListTablesRequest;
14609
14610 let (namespace, _temp_dir) = create_test_namespace().await;
14611 create_scalar_table(&namespace, "alpha").await;
14612 create_scalar_table(&namespace, "beta").await;
14613
14614 let request = ListTablesRequest {
14615 id: Some(vec![]),
14616 page_token: None,
14617 limit: None,
14618 ..Default::default()
14619 };
14620 let response = namespace.list_all_tables(request).await.unwrap();
14621 let mut tables = response.tables;
14622 tables.sort();
14623 assert_eq!(tables, vec!["alpha", "beta"]);
14624 }
14625
14626 #[tokio::test]
14627 async fn test_restore_table() {
14628 use lance_namespace::models::RestoreTableRequest;
14629
14630 let (namespace, _temp_dir) = create_test_namespace().await;
14631 create_scalar_table(&namespace, "users").await;
14632
14633 create_scalar_index(&namespace, "users", "users_id_idx").await;
14635
14636 let dataset = open_dataset(&namespace, "users").await;
14637 let current_version = dataset.version().version;
14638 assert!(current_version >= 2, "Should have at least 2 versions");
14639
14640 let mut restore_req = RestoreTableRequest::new(1);
14642 restore_req.id = Some(vec!["users".to_string()]);
14643 let response = namespace.restore_table(restore_req).await.unwrap();
14644
14645 assert!(
14647 response.transaction_id.is_some(),
14648 "restore_table should return a transaction_id"
14649 );
14650
14651 let dataset_after = open_dataset(&namespace, "users").await;
14653 assert!(
14654 dataset_after.version().version > current_version,
14655 "Restore should create a new version"
14656 );
14657 }
14658
14659 #[tokio::test]
14660 async fn test_alter_table_add_columns() {
14661 use lance_namespace::models::{
14662 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
14663 };
14664
14665 let (namespace, _temp_dir) = create_test_namespace().await;
14666
14667 let schema = create_test_schema();
14669 let ipc_data = create_test_ipc_data(&schema);
14670 let mut create_request = CreateTableRequest::new();
14671 create_request.id = Some(vec!["test_table".to_string()]);
14672 namespace
14673 .create_table(create_request, bytes::Bytes::from(ipc_data))
14674 .await
14675 .unwrap();
14676
14677 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
14679 new_col.expression = Some(Some("id * 2".to_string()));
14680 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
14681 add_request.id = Some(vec!["test_table".to_string()]);
14682
14683 let response = namespace
14684 .alter_table_add_columns(add_request)
14685 .await
14686 .unwrap();
14687 assert!(
14688 response.version > 1,
14689 "Version should increment after adding columns"
14690 );
14691
14692 let mut describe_request = DescribeTableRequest::new();
14694 describe_request.id = Some(vec!["test_table".to_string()]);
14695 describe_request.load_detailed_metadata = Some(true);
14696 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14697 assert!(describe_response.schema.is_some());
14698
14699 let resp_schema = describe_response.schema.unwrap();
14700 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14701 assert!(
14702 field_names.contains(&"doubled_id"),
14703 "Column 'doubled_id' should exist, got: {:?}",
14704 field_names
14705 );
14706 }
14707
14708 #[tokio::test]
14709 async fn test_update_table_schema_metadata() {
14710 use lance_namespace::models::UpdateTableSchemaMetadataRequest;
14711
14712 let (namespace, _temp_dir) = create_test_namespace().await;
14713 create_scalar_table(&namespace, "products").await;
14714
14715 let mut metadata = HashMap::new();
14716 metadata.insert("owner".to_string(), "team_a".to_string());
14717 metadata.insert("version".to_string(), "1.0".to_string());
14718
14719 let mut req = UpdateTableSchemaMetadataRequest::new();
14720 req.id = Some(vec!["products".to_string()]);
14721 req.metadata = Some(metadata.clone());
14722
14723 let response = namespace.update_table_schema_metadata(req).await.unwrap();
14724
14725 assert!(response.metadata.is_some());
14726 let returned = response.metadata.unwrap();
14727 assert_eq!(returned.get("owner"), Some(&"team_a".to_string()));
14728 assert_eq!(returned.get("version"), Some(&"1.0".to_string()));
14729 assert!(
14730 response.transaction_id.is_some(),
14731 "update_table_schema_metadata should return a transaction_id"
14732 );
14733 }
14734
14735 #[tokio::test]
14736 async fn test_alter_table_add_columns_missing_id() {
14737 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
14738
14739 let (namespace, _temp_dir) = create_test_namespace().await;
14740
14741 let new_col = AddColumnsEntry::new("col".to_string());
14742 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
14743 let result = namespace.alter_table_add_columns(request).await;
14744 assert!(result.is_err(), "Should fail when table ID is missing");
14745 }
14746
14747 #[tokio::test]
14748 async fn test_alter_table_alter_columns_rename() {
14749 use lance_namespace::models::{
14750 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
14751 };
14752
14753 let (namespace, _temp_dir) = create_test_namespace().await;
14754
14755 let schema = create_test_schema();
14757 let ipc_data = create_test_ipc_data(&schema);
14758 let mut create_request = CreateTableRequest::new();
14759 create_request.id = Some(vec!["test_table".to_string()]);
14760 namespace
14761 .create_table(create_request, bytes::Bytes::from(ipc_data))
14762 .await
14763 .unwrap();
14764
14765 let mut entry = AlterColumnsEntry::new("name".to_string());
14767 entry.rename = Some(Some("full_name".to_string()));
14768 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
14769 alter_request.id = Some(vec!["test_table".to_string()]);
14770
14771 let response = namespace
14772 .alter_table_alter_columns(alter_request)
14773 .await
14774 .unwrap();
14775 assert!(
14776 response.version > 1,
14777 "Version should increment after altering columns"
14778 );
14779
14780 let mut describe_request = DescribeTableRequest::new();
14782 describe_request.id = Some(vec!["test_table".to_string()]);
14783 describe_request.load_detailed_metadata = Some(true);
14784 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14785 assert!(describe_response.schema.is_some());
14786
14787 let resp_schema = describe_response.schema.unwrap();
14788 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14789 assert!(
14790 field_names.contains(&"full_name"),
14791 "Column should be renamed to 'full_name', got: {:?}",
14792 field_names
14793 );
14794 assert!(
14795 !field_names.contains(&"name"),
14796 "Old column 'name' should not exist, got: {:?}",
14797 field_names
14798 );
14799 }
14800
14801 #[tokio::test]
14802 async fn test_get_table_stats() {
14803 use lance_namespace::models::GetTableStatsRequest;
14804
14805 let (namespace, _temp_dir) = create_test_namespace().await;
14806 create_scalar_table(&namespace, "items").await;
14807 create_scalar_index(&namespace, "items", "items_id_idx").await;
14808
14809 let mut req = GetTableStatsRequest::new();
14810 req.id = Some(vec!["items".to_string()]);
14811
14812 let response = namespace.get_table_stats(req).await.unwrap();
14813 assert_eq!(response.num_rows, 3);
14814 assert_eq!(response.num_indices, 1);
14815 }
14816
14817 #[tokio::test]
14818 async fn test_explain_table_query_plan() {
14819 use lance_namespace::models::QueryTableRequestVector;
14820 use lance_namespace::models::{ExplainTableQueryPlanRequest, QueryTableRequest};
14821
14822 let (namespace, _temp_dir) = create_test_namespace().await;
14823 create_scalar_table(&namespace, "catalog").await;
14824
14825 let mut query = QueryTableRequest::new(1, QueryTableRequestVector::new());
14826 query.filter = Some("id > 1".to_string());
14827 query.columns = Some(Box::new(QueryTableRequestColumns {
14828 column_names: Some(vec!["id".to_string(), "name".to_string()]),
14829 column_aliases: None,
14830 }));
14831 query.with_row_id = Some(true);
14832
14833 let mut req = ExplainTableQueryPlanRequest::new(query);
14834 req.id = Some(vec!["catalog".to_string()]);
14835
14836 let plan_str = namespace.explain_table_query_plan(req).await.unwrap();
14837 assert_plan_contains_all(
14838 &plan_str,
14839 &[
14840 "ProjectionExec: expr=[id@0 as id, name@2 as name",
14841 "projection=[name], source=stream(_rowid)",
14842 "LanceRead: uri=",
14843 "projection=[id]",
14844 "row_id=true, row_addr=false",
14845 "full_filter=id > Int32(1)",
14846 "refine_filter=id > Int32(1)",
14847 ],
14848 "Filtered explain plan should preserve late materialization and filter pushdown",
14849 );
14850 }
14851
14852 #[tokio::test]
14853 async fn test_alter_table_alter_columns_missing_id() {
14854 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
14855
14856 let (namespace, _temp_dir) = create_test_namespace().await;
14857
14858 let entry = AlterColumnsEntry::new("name".to_string());
14859 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
14860 let result = namespace.alter_table_alter_columns(request).await;
14861 assert!(result.is_err(), "Should fail when table ID is missing");
14862 }
14863
14864 #[tokio::test]
14865 async fn test_alter_table_drop_columns() {
14866 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
14867
14868 let (namespace, _temp_dir) = create_test_namespace().await;
14869
14870 let schema = create_test_schema();
14872 let ipc_data = create_test_ipc_data(&schema);
14873 let mut create_request = CreateTableRequest::new();
14874 create_request.id = Some(vec!["test_table".to_string()]);
14875 namespace
14876 .create_table(create_request, bytes::Bytes::from(ipc_data))
14877 .await
14878 .unwrap();
14879
14880 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
14882 drop_request.id = Some(vec!["test_table".to_string()]);
14883
14884 let response = namespace
14885 .alter_table_drop_columns(drop_request)
14886 .await
14887 .unwrap();
14888 assert!(
14889 response.version > 1,
14890 "Version should increment after dropping columns"
14891 );
14892
14893 let mut describe_request = DescribeTableRequest::new();
14895 describe_request.id = Some(vec!["test_table".to_string()]);
14896 describe_request.load_detailed_metadata = Some(true);
14897 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14898 assert!(describe_response.schema.is_some());
14899
14900 let resp_schema = describe_response.schema.unwrap();
14901 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14902 assert!(
14903 !field_names.contains(&"name"),
14904 "Column 'name' should be dropped, got: {:?}",
14905 field_names
14906 );
14907 assert!(
14908 field_names.contains(&"id"),
14909 "Column 'id' should still exist, got: {:?}",
14910 field_names
14911 );
14912 }
14913
14914 #[tokio::test]
14915 async fn test_analyze_table_query_plan() {
14916 use lance_namespace::models::AnalyzeTableQueryPlanRequest;
14917 use lance_namespace::models::QueryTableRequestVector;
14918
14919 let (namespace, _temp_dir) = create_test_namespace().await;
14920 create_scalar_table(&namespace, "catalog").await;
14921
14922 let mut req = AnalyzeTableQueryPlanRequest::new(1, QueryTableRequestVector::new());
14923 req.id = Some(vec!["catalog".to_string()]);
14924 req.filter = Some("id > 0".to_string());
14925 req.columns = Some(Box::new(QueryTableRequestColumns {
14926 column_names: Some(vec!["id".to_string(), "name".to_string()]),
14927 column_aliases: None,
14928 }));
14929 req.with_row_id = Some(true);
14930
14931 let analysis_str = namespace.analyze_table_query_plan(req).await.unwrap();
14932 assert_plan_contains_all(
14933 &analysis_str,
14934 &[
14935 "AnalyzeExec verbose=true",
14936 "ProjectionExec: elapsed=",
14937 "expr=[id@0 as id, name@2 as name",
14938 "projection=[name], source=stream(_rowid)",
14939 "LanceRead: elapsed=",
14940 "projection=[id]",
14941 "row_id=true, row_addr=false",
14942 "full_filter=id > Int32(0)",
14943 "refine_filter=id > Int32(0)",
14944 "metrics=[output_rows=",
14945 ],
14946 "Filtered analyze plan should preserve late materialization and filter pushdown",
14947 );
14948 }
14949
14950 #[tokio::test]
14951 async fn test_dir_listing_no_extra_calls_without_migration() {
14952 let temp_dir = TempStdDir::default();
14953 let temp_path = temp_dir.to_str().unwrap();
14954 let root_uri = file_object_store_uri(temp_path);
14955 let listing_count = Arc::new(AtomicUsize::new(0));
14956 let session = build_listing_counting_session(listing_count.clone());
14957
14958 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
14960 .session(session.clone())
14961 .manifest_enabled(false)
14962 .dir_listing_enabled(true)
14963 .build()
14964 .await
14965 .unwrap();
14966
14967 let schema = create_test_schema();
14968 let ipc_data = create_test_ipc_data(&schema);
14969 let mut create_req = CreateTableRequest::new();
14970 create_req.id = Some(vec!["test_table".to_string()]);
14971 dir_only_ns
14972 .create_table(create_req, Bytes::from(ipc_data))
14973 .await
14974 .unwrap();
14975
14976 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
14978 .session(session)
14979 .manifest_enabled(true)
14980 .dir_listing_enabled(true)
14981 .dir_listing_to_manifest_migration_enabled(false)
14982 .build()
14983 .await
14984 .unwrap();
14985
14986 listing_count.store(0, Ordering::SeqCst);
14988
14989 let mut exists_req = TableExistsRequest::new();
14991 exists_req.id = Some(vec!["test_table".to_string()]);
14992 hybrid_ns.table_exists(exists_req).await.unwrap();
14993
14994 let count = listing_count.load(Ordering::SeqCst);
14995 assert_eq!(
14996 count, 1,
14997 "Expected exactly 1 listing call for table_exists \
14998 without migration mode, but got {}",
14999 count
15000 );
15001
15002 listing_count.store(0, Ordering::SeqCst);
15004
15005 let mut describe_req = DescribeTableRequest::new();
15006 describe_req.id = Some(vec!["test_table".to_string()]);
15007 hybrid_ns.describe_table(describe_req).await.unwrap();
15008
15009 let count = listing_count.load(Ordering::SeqCst);
15010 assert_eq!(
15011 count, 1,
15012 "Expected exactly 1 listing call for describe_table \
15013 without migration mode, but got {}",
15014 count
15015 );
15016 }
15017
15018 #[tokio::test]
15019 async fn test_build_and_root_reads_do_not_create_manifest() {
15020 let temp_dir = TempStdDir::default();
15021 let temp_path = temp_dir.to_str().unwrap();
15022 let manifest_path = std::path::Path::new(temp_path).join("__manifest");
15023
15024 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
15025 .manifest_enabled(false)
15026 .dir_listing_enabled(true)
15027 .build()
15028 .await
15029 .unwrap();
15030 create_scalar_table(&dir_only_ns, "catalog").await;
15031 assert!(!manifest_path.exists());
15032
15033 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15034 .manifest_enabled(true)
15035 .dir_listing_enabled(true)
15036 .build()
15037 .await
15038 .unwrap();
15039 assert!(!manifest_path.exists());
15040
15041 let mut exists_req = TableExistsRequest::new();
15042 exists_req.id = Some(vec!["catalog".to_string()]);
15043 namespace.table_exists(exists_req).await.unwrap();
15044 assert!(!manifest_path.exists());
15045
15046 let mut describe_req = DescribeTableRequest::new();
15047 describe_req.id = Some(vec!["catalog".to_string()]);
15048 namespace.describe_table(describe_req).await.unwrap();
15049 assert!(!manifest_path.exists());
15050
15051 let list_response = namespace
15052 .list_tables(ListTablesRequest {
15053 id: Some(vec![]),
15054 ..Default::default()
15055 })
15056 .await
15057 .unwrap();
15058 assert_eq!(list_response.tables, vec!["catalog".to_string()]);
15059 assert!(!manifest_path.exists());
15060
15061 let mut list_namespaces_req = ListNamespacesRequest::new();
15062 list_namespaces_req.id = Some(vec!["workspace".to_string()]);
15063 let err = namespace
15064 .list_namespaces(list_namespaces_req)
15065 .await
15066 .unwrap_err();
15067 assert!(err.to_string().contains("__manifest"));
15068 assert!(!manifest_path.exists());
15069
15070 let err = namespace
15071 .list_tables(ListTablesRequest {
15072 id: Some(vec!["workspace".to_string()]),
15073 ..Default::default()
15074 })
15075 .await
15076 .unwrap_err();
15077 assert!(err.to_string().contains("__manifest"));
15078 assert!(!manifest_path.exists());
15079
15080 let mut child_describe_req = DescribeTableRequest::new();
15081 child_describe_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
15082 let err = namespace
15083 .describe_table(child_describe_req)
15084 .await
15085 .unwrap_err();
15086 assert!(err.to_string().contains("__manifest"));
15087 assert!(!manifest_path.exists());
15088
15089 let mut child_exists_req = TableExistsRequest::new();
15090 child_exists_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
15091 let err = namespace.table_exists(child_exists_req).await.unwrap_err();
15092 assert!(err.to_string().contains("__manifest"));
15093 assert!(!manifest_path.exists());
15094
15095 let mut create_ns_req = CreateNamespaceRequest::new();
15096 create_ns_req.id = Some(vec!["workspace".to_string()]);
15097 namespace.create_namespace(create_ns_req).await.unwrap();
15098 assert!(manifest_path.exists());
15099 }
15100
15101 #[tokio::test]
15102 async fn test_migrate_updates_read_opened_legacy_manifest() {
15103 let temp_dir = TempStdDir::default();
15104 let temp_path = temp_dir.to_str().unwrap();
15105 create_legacy_manifest_without_primary_key_metadata(temp_path).await;
15106 assert!(!manifest_has_primary_key_metadata(temp_path).await);
15107
15108 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15109 .manifest_enabled(true)
15110 .dir_listing_enabled(true)
15111 .build()
15112 .await
15113 .unwrap();
15114 assert!(!manifest_has_primary_key_metadata(temp_path).await);
15115
15116 let migrated = namespace.migrate().await.unwrap();
15117 assert_eq!(migrated, 0);
15118 assert!(manifest_has_primary_key_metadata(temp_path).await);
15119 }
15120
15121 #[tokio::test]
15122 async fn test_describe_declared_table_checks_versions_only_when_requested() {
15123 let temp_dir = TempStdDir::default();
15124 let temp_path = temp_dir.to_str().unwrap();
15125 let root_uri = file_object_store_uri(temp_path);
15126 let listing_count = Arc::new(AtomicUsize::new(0));
15127 let session = build_listing_counting_session(listing_count.clone());
15128
15129 let namespace = DirectoryNamespaceBuilder::new(root_uri)
15130 .session(session)
15131 .manifest_enabled(false)
15132 .dir_listing_enabled(true)
15133 .build()
15134 .await
15135 .unwrap();
15136
15137 let mut declare_req = DeclareTableRequest::new();
15138 declare_req.id = Some(vec!["test_table".to_string()]);
15139 namespace.declare_table(declare_req).await.unwrap();
15140
15141 listing_count.store(0, Ordering::SeqCst);
15142
15143 let mut describe_req = DescribeTableRequest::new();
15144 describe_req.id = Some(vec!["test_table".to_string()]);
15145 let describe_response = namespace.describe_table(describe_req).await.unwrap();
15146
15147 assert_eq!(describe_response.is_only_declared, None);
15148 assert_eq!(
15149 listing_count.load(Ordering::SeqCst),
15150 1,
15151 "Default describe_table should only list the table directory"
15152 );
15153
15154 listing_count.store(0, Ordering::SeqCst);
15155
15156 let mut describe_req = DescribeTableRequest::new();
15157 describe_req.id = Some(vec!["test_table".to_string()]);
15158 describe_req.check_declared = Some(true);
15159 let describe_response = namespace.describe_table(describe_req).await.unwrap();
15160
15161 assert_eq!(describe_response.is_only_declared, Some(true));
15162 assert_eq!(
15163 listing_count.load(Ordering::SeqCst),
15164 2,
15165 "check_declared describe_table should list the table directory and _versions"
15166 );
15167 }
15168
15169 #[tokio::test]
15170 async fn test_dir_listing_extra_calls_with_migration() {
15171 let temp_dir = TempStdDir::default();
15172 let temp_path = temp_dir.to_str().unwrap();
15173 let root_uri = file_object_store_uri(temp_path);
15174 let listing_count = Arc::new(AtomicUsize::new(0));
15175 let session = build_listing_counting_session(listing_count.clone());
15176
15177 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
15179 .session(session.clone())
15180 .manifest_enabled(false)
15181 .dir_listing_enabled(true)
15182 .build()
15183 .await
15184 .unwrap();
15185
15186 let schema = create_test_schema();
15187 let ipc_data = create_test_ipc_data(&schema);
15188 let mut create_req = CreateTableRequest::new();
15189 create_req.id = Some(vec!["test_table".to_string()]);
15190 dir_only_ns
15191 .create_table(create_req, Bytes::from(ipc_data))
15192 .await
15193 .unwrap();
15194
15195 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
15196 .session(session)
15197 .manifest_enabled(true)
15198 .dir_listing_enabled(true)
15199 .dir_listing_to_manifest_migration_enabled(true)
15200 .build()
15201 .await
15202 .unwrap();
15203
15204 listing_count.store(0, Ordering::SeqCst);
15212
15213 let mut exists_req = TableExistsRequest::new();
15214 exists_req.id = Some(vec!["test_table".to_string()]);
15215 hybrid_ns.table_exists(exists_req).await.unwrap();
15216
15217 let count = listing_count.load(Ordering::SeqCst);
15218 assert_eq!(
15219 count, 2,
15220 "Expected 2 listing calls for table_exists with migration mode \
15221 (absent-__manifest probe + table directory fallback), but got {}",
15222 count
15223 );
15224
15225 listing_count.store(0, Ordering::SeqCst);
15228
15229 let mut describe_req = DescribeTableRequest::new();
15230 describe_req.id = Some(vec!["test_table".to_string()]);
15231 hybrid_ns.describe_table(describe_req).await.unwrap();
15232
15233 let count = listing_count.load(Ordering::SeqCst);
15234 assert_eq!(
15235 count, 2,
15236 "Expected 2 listing calls for describe_table with migration mode \
15237 (absent-__manifest probe + table directory fallback), but got {}",
15238 count
15239 );
15240 }
15241
15242 #[tokio::test]
15243 async fn test_manifest_reload_observes_new_version_from_other_namespace() {
15244 let temp_dir = TempStdDir::default();
15245 let temp_path = temp_dir.to_str().unwrap();
15246
15247 let namespace_a = DirectoryNamespaceBuilder::new(temp_path)
15248 .manifest_enabled(true)
15249 .dir_listing_enabled(false)
15250 .build()
15251 .await
15252 .unwrap();
15253 create_scalar_table(&namespace_a, "alpha").await;
15254
15255 let namespace_b = DirectoryNamespaceBuilder::new(temp_path)
15256 .manifest_enabled(true)
15257 .dir_listing_enabled(false)
15258 .build()
15259 .await
15260 .unwrap();
15261 create_scalar_table(&namespace_b, "beta").await;
15262
15263 let response = namespace_a
15264 .list_tables(ListTablesRequest {
15265 id: Some(vec![]),
15266 ..Default::default()
15267 })
15268 .await
15269 .unwrap();
15270
15271 let mut tables = response.tables;
15272 tables.sort();
15273 assert_eq!(tables, vec!["alpha", "beta"]);
15274 }
15275
15276 #[tokio::test]
15277 async fn test_migration_not_found_errors_include_table_id() {
15278 let temp_dir = TempStdDir::default();
15279 let temp_path = temp_dir.to_str().unwrap();
15280
15281 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15282 .manifest_enabled(true)
15283 .dir_listing_enabled(true)
15284 .dir_listing_to_manifest_migration_enabled(true)
15285 .build()
15286 .await
15287 .unwrap();
15288
15289 let mut exists_req = TableExistsRequest::new();
15290 exists_req.id = Some(vec!["missing_table".to_string()]);
15291 let err = namespace.table_exists(exists_req).await.unwrap_err();
15292 assert!(matches!(err, Error::Namespace { .. }));
15293 let err_msg = err.to_string();
15294 assert!(err_msg.contains("Table not found"));
15295 assert!(err_msg.contains("table id 'missing_table'"));
15296
15297 let mut describe_req = DescribeTableRequest::new();
15298 describe_req.id = Some(vec!["missing_table".to_string()]);
15299 let err = namespace.describe_table(describe_req).await.unwrap_err();
15300 assert!(matches!(err, Error::Namespace { .. }));
15301 let err_msg = err.to_string();
15302 assert!(err_msg.contains("Table not found"));
15303 assert!(err_msg.contains("table id 'missing_table'"));
15304 }
15305
15306 #[tokio::test]
15307 async fn test_manifest_not_found_errors_include_full_table_id() {
15308 use lance_namespace::models::CreateNamespaceRequest;
15309
15310 let temp_dir = TempStdDir::default();
15311 let temp_path = temp_dir.to_str().unwrap();
15312
15313 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15314 .manifest_enabled(true)
15315 .dir_listing_enabled(true)
15316 .build()
15317 .await
15318 .unwrap();
15319
15320 let mut create_ns_req = CreateNamespaceRequest::new();
15321 create_ns_req.id = Some(vec!["workspace".to_string()]);
15322 namespace.create_namespace(create_ns_req).await.unwrap();
15323
15324 let missing_table_id = vec!["workspace".to_string(), "missing_table".to_string()];
15325
15326 let mut exists_req = TableExistsRequest::new();
15327 exists_req.id = Some(missing_table_id.clone());
15328 let err = namespace.table_exists(exists_req).await.unwrap_err();
15329 assert!(matches!(err, Error::Namespace { .. }));
15330 let err_msg = err.to_string();
15331 assert!(err_msg.contains("Table not found"));
15332 assert!(err_msg.contains("table id 'workspace$missing_table'"));
15333
15334 let mut describe_req = DescribeTableRequest::new();
15335 describe_req.id = Some(missing_table_id);
15336 let err = namespace.describe_table(describe_req).await.unwrap_err();
15337 assert!(matches!(err, Error::Namespace { .. }));
15338 let err_msg = err.to_string();
15339 assert!(err_msg.contains("Table not found"));
15340 assert!(err_msg.contains("table id 'workspace$missing_table'"));
15341 }
15342
15343 async fn create_tagged_test_table(
15346 versions: u32,
15347 ) -> (Arc<DirectoryNamespace>, TempStdDir, Vec<String>) {
15348 use arrow::array::{Int32Array, RecordBatchIterator};
15349 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
15350 use arrow::record_batch::RecordBatch;
15351 use lance::dataset::{Dataset, WriteMode, WriteParams};
15352
15353 assert!(versions >= 1, "versions must be at least 1");
15354
15355 let temp_dir = TempStdDir::default();
15356 let temp_path = temp_dir.to_str().unwrap();
15357
15358 let namespace = Arc::new(
15359 DirectoryNamespaceBuilder::new(temp_path)
15360 .build()
15361 .await
15362 .unwrap(),
15363 );
15364 let table_id = vec!["tag_table".to_string()];
15365 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
15366 "id",
15367 DataType::Int32,
15368 false,
15369 )]));
15370 let initial_batch = RecordBatch::try_new(
15371 arrow_schema.clone(),
15372 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
15373 )
15374 .unwrap();
15375 let batches = RecordBatchIterator::new(vec![Ok(initial_batch)], arrow_schema.clone());
15376 let write_params = WriteParams {
15377 mode: WriteMode::Create,
15378 ..Default::default()
15379 };
15380
15381 let mut dataset = Dataset::write_into_namespace(
15382 batches,
15383 namespace.clone() as Arc<dyn LanceNamespace>,
15384 table_id.clone(),
15385 Some(write_params),
15386 )
15387 .await
15388 .unwrap();
15389
15390 for i in 1..versions {
15391 let value_start = (i as i32) * 10;
15392 let batch = RecordBatch::try_new(
15393 arrow_schema.clone(),
15394 vec![Arc::new(Int32Array::from(vec![
15395 value_start,
15396 value_start + 1,
15397 ]))],
15398 )
15399 .unwrap();
15400 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
15401 dataset.append(batches, None).await.unwrap();
15402 }
15403
15404 (namespace, temp_dir, table_id)
15405 }
15406
15407 fn namespace_code(err: &Error) -> Option<ErrorCode> {
15409 match err {
15410 Error::Namespace { source, .. } => {
15411 source.downcast_ref::<NamespaceError>().map(|e| e.code())
15412 }
15413 _ => None,
15414 }
15415 }
15416
15417 #[tokio::test]
15418 async fn test_create_and_list_branches() {
15419 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15420
15421 namespace
15422 .create_table_branch(CreateTableBranchRequest {
15423 id: Some(table_id.clone()),
15424 name: "dev".to_string(),
15425 ..Default::default()
15426 })
15427 .await
15428 .unwrap();
15429 namespace
15430 .create_table_branch(CreateTableBranchRequest {
15431 id: Some(table_id.clone()),
15432 name: "staging".to_string(),
15433 ..Default::default()
15434 })
15435 .await
15436 .unwrap();
15437
15438 let resp = namespace
15439 .list_table_branches(ListTableBranchesRequest {
15440 id: Some(table_id.clone()),
15441 ..Default::default()
15442 })
15443 .await
15444 .unwrap();
15445 assert_eq!(
15446 resp.branches.len(),
15447 2,
15448 "expected 2 branches, got: {:?}",
15449 resp.branches
15450 );
15451 assert!(resp.branches.contains_key("dev"));
15452 assert!(resp.branches.contains_key("staging"));
15453 assert!(resp.page_token.is_none());
15454
15455 namespace
15457 .delete_table_branch(DeleteTableBranchRequest {
15458 id: Some(table_id.clone()),
15459 name: "dev".to_string(),
15460 ..Default::default()
15461 })
15462 .await
15463 .unwrap();
15464
15465 let resp = namespace
15466 .list_table_branches(ListTableBranchesRequest {
15467 id: Some(table_id),
15468 ..Default::default()
15469 })
15470 .await
15471 .unwrap();
15472 assert_eq!(resp.branches.len(), 1, "expected 1 branch after delete");
15473 assert!(!resp.branches.contains_key("dev"));
15474 assert!(resp.branches.contains_key("staging"));
15475 }
15476
15477 #[tokio::test]
15478 async fn test_create_branch_from_version() {
15479 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15480
15481 namespace
15483 .create_table_branch(CreateTableBranchRequest {
15484 id: Some(table_id.clone()),
15485 name: "from-v1".to_string(),
15486 from_version: Some(1),
15487 ..Default::default()
15488 })
15489 .await
15490 .unwrap();
15491
15492 let resp = namespace
15493 .list_table_branches(ListTableBranchesRequest {
15494 id: Some(table_id),
15495 ..Default::default()
15496 })
15497 .await
15498 .unwrap();
15499 let branch = resp
15500 .branches
15501 .get("from-v1")
15502 .expect("forked branch should be listed");
15503 assert_eq!(
15504 branch.parent_version, 1,
15505 "branch should fork from version 1"
15506 );
15507 assert!(
15508 branch.parent_branch.is_none(),
15509 "a branch forked from main has no parent branch"
15510 );
15511 }
15512
15513 #[tokio::test]
15518 async fn test_create_branch_from_other_branch() {
15519 use lance::dataset::builder::DatasetBuilder;
15520
15521 let (namespace, _temp_dir) = create_test_namespace().await;
15522 create_scalar_table(&namespace, "users").await; create_branch_with_commits(&namespace, "users", "dev", 1).await;
15525 let main_ds = open_dataset(&namespace, "users").await;
15527 append_scalar_version(main_ds.uri(), 500).await; namespace
15530 .create_table_branch(CreateTableBranchRequest {
15531 id: Some(vec!["users".to_string()]),
15532 name: "child".to_string(),
15533 from_branch: Some("dev".to_string()),
15534 from_version: Some(2),
15535 ..Default::default()
15536 })
15537 .await
15538 .unwrap();
15539
15540 let child_ds = DatasetBuilder::from_uri(main_ds.uri())
15541 .with_branch("child", None)
15542 .load()
15543 .await
15544 .unwrap();
15545 let ids = scan_id_column(&child_ds).await;
15546 assert!(
15547 ids.contains(&100) && ids.contains(&101),
15548 "child must contain dev's appended rows, got: {:?}",
15549 ids
15550 );
15551 assert!(
15552 !ids.contains(&500),
15553 "child must not contain main's diverged rows, got: {:?}",
15554 ids
15555 );
15556
15557 let listed = namespace
15559 .list_table_branches(ListTableBranchesRequest {
15560 id: Some(vec!["users".to_string()]),
15561 ..Default::default()
15562 })
15563 .await
15564 .unwrap();
15565 assert_eq!(
15566 listed
15567 .branches
15568 .get("child")
15569 .unwrap()
15570 .parent_branch
15571 .as_deref(),
15572 Some("dev")
15573 );
15574 }
15575
15576 #[tokio::test]
15577 async fn test_create_existing_branch_conflict() {
15578 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15579
15580 namespace
15581 .create_table_branch(CreateTableBranchRequest {
15582 id: Some(table_id.clone()),
15583 name: "dev".to_string(),
15584 ..Default::default()
15585 })
15586 .await
15587 .unwrap();
15588
15589 let err = namespace
15590 .create_table_branch(CreateTableBranchRequest {
15591 id: Some(table_id),
15592 name: "dev".to_string(),
15593 ..Default::default()
15594 })
15595 .await
15596 .unwrap_err();
15597 assert_eq!(
15598 namespace_code(&err),
15599 Some(ErrorCode::TableBranchAlreadyExists),
15600 "expected TableBranchAlreadyExists, got: {}",
15601 err
15602 );
15603 assert!(
15604 err.to_string().to_lowercase().contains("already exists"),
15605 "expected already-exists message, got: {}",
15606 err
15607 );
15608 }
15609
15610 #[tokio::test]
15611 async fn test_delete_unknown_branch() {
15612 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15613
15614 let err = namespace
15615 .delete_table_branch(DeleteTableBranchRequest {
15616 id: Some(table_id),
15617 name: "does-not-exist".to_string(),
15618 ..Default::default()
15619 })
15620 .await
15621 .unwrap_err();
15622 assert_eq!(
15623 namespace_code(&err),
15624 Some(ErrorCode::TableBranchNotFound),
15625 "expected TableBranchNotFound, got: {}",
15626 err
15627 );
15628 assert!(
15629 err.to_string().to_lowercase().contains("not found"),
15630 "expected not-found message, got: {}",
15631 err
15632 );
15633 }
15634
15635 #[tokio::test]
15636 async fn test_delete_referenced_branch_conflict() {
15637 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15638
15639 namespace
15641 .create_table_branch(CreateTableBranchRequest {
15642 id: Some(table_id.clone()),
15643 name: "parent".to_string(),
15644 ..Default::default()
15645 })
15646 .await
15647 .unwrap();
15648 namespace
15649 .create_table_branch(CreateTableBranchRequest {
15650 id: Some(table_id.clone()),
15651 name: "child".to_string(),
15652 from_branch: Some("parent".to_string()),
15653 ..Default::default()
15654 })
15655 .await
15656 .unwrap();
15657
15658 let listed = namespace
15660 .list_table_branches(ListTableBranchesRequest {
15661 id: Some(table_id.clone()),
15662 ..Default::default()
15663 })
15664 .await
15665 .unwrap();
15666 let child = listed
15667 .branches
15668 .get("child")
15669 .expect("child branch should be listed");
15670 assert_eq!(
15671 child.parent_branch.as_deref(),
15672 Some("parent"),
15673 "child should record parent branch as its fork point"
15674 );
15675 assert!(
15676 child.parent_version >= 1,
15677 "child should record the parent version it forked from, got {}",
15678 child.parent_version
15679 );
15680
15681 let err = namespace
15684 .delete_table_branch(DeleteTableBranchRequest {
15685 id: Some(table_id),
15686 name: "parent".to_string(),
15687 ..Default::default()
15688 })
15689 .await
15690 .unwrap_err();
15691 assert_eq!(
15692 namespace_code(&err),
15693 Some(ErrorCode::InvalidInput),
15694 "expected InvalidInput for deleting a referenced branch, got: {}",
15695 err
15696 );
15697 assert!(
15698 err.to_string().to_lowercase().contains("referenced"),
15699 "error should explain the branch is still referenced, got: {}",
15700 err
15701 );
15702 }
15703
15704 #[tokio::test]
15705 async fn test_branch_name_required() {
15706 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15707
15708 let create_err = namespace
15709 .create_table_branch(CreateTableBranchRequest {
15710 id: Some(table_id.clone()),
15711 name: String::new(),
15712 ..Default::default()
15713 })
15714 .await
15715 .unwrap_err();
15716 assert_eq!(
15717 namespace_code(&create_err),
15718 Some(ErrorCode::InvalidInput),
15719 "empty name on create should be InvalidInput, got: {}",
15720 create_err
15721 );
15722 assert!(
15723 create_err
15724 .to_string()
15725 .to_lowercase()
15726 .contains("must not be empty")
15727 );
15728
15729 let delete_err = namespace
15730 .delete_table_branch(DeleteTableBranchRequest {
15731 id: Some(table_id),
15732 name: String::new(),
15733 ..Default::default()
15734 })
15735 .await
15736 .unwrap_err();
15737 assert_eq!(
15738 namespace_code(&delete_err),
15739 Some(ErrorCode::InvalidInput),
15740 "empty name on delete should be InvalidInput, got: {}",
15741 delete_err
15742 );
15743 assert!(
15744 delete_err
15745 .to_string()
15746 .to_lowercase()
15747 .contains("must not be empty")
15748 );
15749 }
15750
15751 #[tokio::test]
15752 async fn test_create_branch_rejects_negative_from_version() {
15753 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15754
15755 let err = namespace
15756 .create_table_branch(CreateTableBranchRequest {
15757 id: Some(table_id),
15758 name: "dev".to_string(),
15759 from_version: Some(-1),
15760 ..Default::default()
15761 })
15762 .await
15763 .unwrap_err();
15764 assert_eq!(
15765 namespace_code(&err),
15766 Some(ErrorCode::InvalidInput),
15767 "negative from_version should be InvalidInput, got: {}",
15768 err
15769 );
15770 assert!(err.to_string().to_lowercase().contains("from_version"));
15771 }
15772
15773 #[tokio::test]
15774 async fn test_create_branch_nonexistent_from_version() {
15775 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15776
15777 let err = namespace
15780 .create_table_branch(CreateTableBranchRequest {
15781 id: Some(table_id),
15782 name: "dev".to_string(),
15783 from_version: Some(999),
15784 ..Default::default()
15785 })
15786 .await
15787 .unwrap_err();
15788 assert_eq!(
15789 namespace_code(&err),
15790 Some(ErrorCode::InvalidInput),
15791 "non-existent from_version should map to InvalidInput, got: {}",
15792 err
15793 );
15794 assert!(
15795 err.to_string().to_lowercase().contains("does not exist"),
15796 "error should name the missing source, got: {}",
15797 err
15798 );
15799 }
15800
15801 #[tokio::test]
15802 async fn test_create_and_list_tags() {
15803 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15804
15805 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15806 req.id = Some(table_id.clone());
15807 namespace.create_table_tag(req).await.unwrap();
15808
15809 let mut req = CreateTableTagRequest::new("v2".to_string(), 2);
15810 req.id = Some(table_id.clone());
15811 namespace.create_table_tag(req).await.unwrap();
15812
15813 let mut list_req = ListTableTagsRequest::new();
15814 list_req.id = Some(table_id);
15815 let resp = namespace.list_table_tags(list_req).await.unwrap();
15816
15817 assert_eq!(resp.tags.len(), 2, "expected 2 tags, got: {:?}", resp.tags);
15818 assert_eq!(resp.tags.get("v1").unwrap().version, 1);
15819 assert_eq!(resp.tags.get("v2").unwrap().version, 2);
15820 assert!(resp.page_token.is_none());
15821 }
15822
15823 #[tokio::test]
15824 async fn test_create_existing_tag_conflict() {
15825 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15826
15827 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15828 req.id = Some(table_id.clone());
15829 namespace.create_table_tag(req).await.unwrap();
15830
15831 let mut req = CreateTableTagRequest::new("v1".to_string(), 2);
15832 req.id = Some(table_id);
15833 let err = namespace.create_table_tag(req).await.unwrap_err();
15834 assert!(
15835 err.to_string().to_lowercase().contains("already exists"),
15836 "expected already-exists error, got: {}",
15837 err
15838 );
15839 }
15840
15841 #[tokio::test]
15842 async fn test_get_tag_version() {
15843 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15844
15845 let mut req = CreateTableTagRequest::new("release".to_string(), 2);
15846 req.id = Some(table_id.clone());
15847 namespace.create_table_tag(req).await.unwrap();
15848
15849 let mut get_req = GetTableTagVersionRequest::new("release".to_string());
15850 get_req.id = Some(table_id);
15851 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15852 assert_eq!(resp.version, 2);
15853 assert_eq!(resp.branch, None);
15854 }
15855
15856 #[tokio::test]
15857 async fn test_get_unknown_tag() {
15858 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15859
15860 let mut get_req = GetTableTagVersionRequest::new("does-not-exist".to_string());
15861 get_req.id = Some(table_id);
15862 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
15863 assert!(
15864 err.to_string().to_lowercase().contains("not found"),
15865 "expected not-found error, got: {}",
15866 err
15867 );
15868 }
15869
15870 #[tokio::test]
15871 async fn test_update_tag_to_new_version() {
15872 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15873
15874 let mut req = CreateTableTagRequest::new("rolling".to_string(), 1);
15875 req.id = Some(table_id.clone());
15876 namespace.create_table_tag(req).await.unwrap();
15877
15878 let mut update_req = UpdateTableTagRequest::new("rolling".to_string(), 3);
15879 update_req.id = Some(table_id.clone());
15880 namespace.update_table_tag(update_req).await.unwrap();
15881
15882 let mut get_req = GetTableTagVersionRequest::new("rolling".to_string());
15883 get_req.id = Some(table_id);
15884 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15885 assert_eq!(resp.version, 3);
15886 }
15887
15888 #[tokio::test]
15889 async fn test_update_unknown_tag() {
15890 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15891
15892 let mut update_req = UpdateTableTagRequest::new("ghost".to_string(), 1);
15893 update_req.id = Some(table_id);
15894 let err = namespace.update_table_tag(update_req).await.unwrap_err();
15895 assert!(
15896 err.to_string().to_lowercase().contains("not found"),
15897 "expected not-found error, got: {}",
15898 err
15899 );
15900 }
15901
15902 #[tokio::test]
15903 async fn test_delete_tag() {
15904 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15905
15906 let mut req = CreateTableTagRequest::new("doomed".to_string(), 1);
15907 req.id = Some(table_id.clone());
15908 namespace.create_table_tag(req).await.unwrap();
15909
15910 let mut delete_req = DeleteTableTagRequest::new("doomed".to_string());
15911 delete_req.id = Some(table_id.clone());
15912 namespace.delete_table_tag(delete_req).await.unwrap();
15913
15914 let mut list_req = ListTableTagsRequest::new();
15915 list_req.id = Some(table_id.clone());
15916 let resp = namespace.list_table_tags(list_req).await.unwrap();
15917 assert!(resp.tags.is_empty(), "tag should be removed after delete");
15918
15919 let mut get_req = GetTableTagVersionRequest::new("doomed".to_string());
15921 get_req.id = Some(table_id);
15922 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
15923 assert!(err.to_string().to_lowercase().contains("not found"));
15924 }
15925
15926 #[tokio::test]
15927 async fn test_delete_unknown_tag() {
15928 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15929
15930 let mut delete_req = DeleteTableTagRequest::new("nope".to_string());
15931 delete_req.id = Some(table_id);
15932 let err = namespace.delete_table_tag(delete_req).await.unwrap_err();
15933 assert!(
15934 err.to_string().to_lowercase().contains("not found"),
15935 "expected not-found error, got: {}",
15936 err
15937 );
15938 }
15939
15940 #[tokio::test]
15941 async fn test_create_tag_invalid_version() {
15942 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15943
15944 let mut req = CreateTableTagRequest::new("v0".to_string(), 0);
15946 req.id = Some(table_id.clone());
15947 let err = namespace.create_table_tag(req).await.unwrap_err();
15948 assert!(
15949 err.to_string().to_lowercase().contains("positive"),
15950 "expected positive-version error, got: {}",
15951 err
15952 );
15953
15954 let mut req = CreateTableTagRequest::new(String::new(), 1);
15956 req.id = Some(table_id);
15957 let err = namespace.create_table_tag(req).await.unwrap_err();
15958 assert!(
15959 err.to_string().to_lowercase().contains("must not be empty"),
15960 "expected empty-tag-name error, got: {}",
15961 err
15962 );
15963 }
15964
15965 #[tokio::test]
15966 async fn test_create_tag_table_not_found() {
15967 let (namespace, _temp_dir) = create_test_namespace().await;
15968
15969 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15970 req.id = Some(vec!["does_not_exist".to_string()]);
15971 let err = namespace.create_table_tag(req).await.unwrap_err();
15972 let msg = err.to_string();
15973 assert!(
15974 msg.contains("Table") && msg.to_lowercase().contains("not found"),
15975 "expected TableNotFound error, got: {}",
15976 err
15977 );
15978 }
15979 #[tokio::test]
15980 async fn test_alter_table_drop_columns_missing_id() {
15981 use lance_namespace::models::AlterTableDropColumnsRequest;
15982
15983 let (namespace, _temp_dir) = create_test_namespace().await;
15984
15985 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
15986 let result = namespace.alter_table_drop_columns(request).await;
15987 assert!(result.is_err(), "Should fail when table ID is missing");
15988 }
15989
15990 #[tokio::test]
15991 async fn test_alter_table_drop_columns_nonexistent_table() {
15992 use lance_namespace::models::AlterTableDropColumnsRequest;
15993
15994 let (namespace, _temp_dir) = create_test_namespace().await;
15995
15996 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
15997 request.id = Some(vec!["nonexistent".to_string()]);
15998 let result = namespace.alter_table_drop_columns(request).await;
15999 assert!(result.is_err(), "Should fail when table does not exist");
16000 }
16001
16002 #[tokio::test]
16003 async fn test_create_branch_on_managed_dataset_succeeds() {
16004 use lance::dataset::builder::DatasetBuilder;
16005
16006 let temp = TempStdDir::default();
16007 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
16008 let table_id = vec!["t".to_string()];
16009 let mut main = create_managed_table(&ns, &table_id).await;
16010
16011 let fork_version = main.version().version;
16012 let branch = main
16013 .create_branch("exp", fork_version, None)
16014 .await
16015 .expect("create_branch failed");
16016 assert_eq!(branch.manifest.branch.as_deref(), Some("exp"));
16017 assert_eq!(scan_id_column(&branch).await, vec![1, 2]);
16018
16019 let reopened = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
16020 .await
16021 .unwrap()
16022 .with_branch("exp", None)
16023 .load()
16024 .await
16025 .expect("reopen branch failed");
16026 assert_eq!(scan_id_column(&reopened).await, vec![1, 2]);
16027 }
16028
16029 #[tokio::test]
16030 async fn test_alter_transaction_set_status() {
16031 use lance_namespace::models::{
16032 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16033 DescribeTransactionRequest,
16034 };
16035
16036 let (namespace, _temp_dir) = create_test_namespace().await;
16037 create_scalar_table(&namespace, "users").await;
16038 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16039 .await
16040 .expect("create_scalar_index should return a transaction id");
16041
16042 let describe_resp = namespace
16044 .describe_transaction(DescribeTransactionRequest {
16045 id: Some(vec!["users".to_string(), txn_id.clone()]),
16046 ..Default::default()
16047 })
16048 .await
16049 .unwrap();
16050 assert_eq!(describe_resp.status, "SUCCEEDED");
16051
16052 let response = namespace
16054 .alter_transaction(AlterTransactionRequest {
16055 id: Some(vec!["users".to_string(), txn_id.clone()]),
16056 actions: vec![AlterTransactionAction {
16057 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16058 status: Some("Canceled".to_string()),
16059 })),
16060 set_property_action: None,
16061 unset_property_action: None,
16062 }],
16063 ..Default::default()
16064 })
16065 .await
16066 .unwrap();
16067 assert_eq!(response.status, "Canceled");
16068 assert!(response.properties.is_some());
16069 let props = response.properties.unwrap();
16070 assert_eq!(props.get("uuid"), Some(&txn_id));
16071 assert_eq!(props.get("operation"), Some(&"CreateIndex".to_string()));
16072 }
16073
16074 #[tokio::test]
16075 async fn test_alter_transaction_set_property() {
16076 use lance_namespace::models::{
16077 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16078 };
16079
16080 let (namespace, _temp_dir) = create_test_namespace().await;
16081 create_scalar_table(&namespace, "users").await;
16082 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16083 .await
16084 .expect("create_scalar_index should return a transaction id");
16085
16086 let response = namespace
16087 .alter_transaction(AlterTransactionRequest {
16088 id: Some(vec!["users".to_string(), txn_id.clone()]),
16089 actions: vec![AlterTransactionAction {
16090 set_status_action: None,
16091 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16092 key: Some("custom_key".to_string()),
16093 value: Some("custom_value".to_string()),
16094 mode: None,
16095 })),
16096 unset_property_action: None,
16097 }],
16098 ..Default::default()
16099 })
16100 .await
16101 .unwrap();
16102 assert_eq!(response.status, "SUCCEEDED");
16103 let props = response.properties.unwrap();
16104 assert_eq!(props.get("custom_key"), Some(&"custom_value".to_string()));
16105 }
16106
16107 #[tokio::test]
16108 async fn test_alter_transaction_set_property_fail_mode() {
16109 use lance_namespace::models::{
16110 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16111 };
16112
16113 let (namespace, _temp_dir) = create_test_namespace().await;
16114 create_scalar_table(&namespace, "users").await;
16115 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16116 .await
16117 .expect("create_scalar_index should return a transaction id");
16118
16119 namespace
16121 .alter_transaction(AlterTransactionRequest {
16122 id: Some(vec!["users".to_string(), txn_id.clone()]),
16123 actions: vec![AlterTransactionAction {
16124 set_status_action: None,
16125 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16126 key: Some("custom_key".to_string()),
16127 value: Some("initial_value".to_string()),
16128 mode: None,
16129 })),
16130 unset_property_action: None,
16131 }],
16132 ..Default::default()
16133 })
16134 .await
16135 .unwrap();
16136
16137 let result = namespace
16140 .alter_transaction(AlterTransactionRequest {
16141 id: Some(vec!["users".to_string(), txn_id.clone()]),
16142 actions: vec![AlterTransactionAction {
16143 set_status_action: None,
16144 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16145 key: Some("custom_key".to_string()),
16146 value: Some("new_value".to_string()),
16147 mode: Some("Fail".to_string()),
16148 })),
16149 unset_property_action: None,
16150 }],
16151 ..Default::default()
16152 })
16153 .await;
16154 assert!(result.is_err());
16155 }
16156
16157 #[tokio::test]
16158 async fn test_alter_transaction_unset_property() {
16159 use lance_namespace::models::{
16160 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16161 AlterTransactionUnsetProperty,
16162 };
16163
16164 let (namespace, _temp_dir) = create_test_namespace().await;
16165 create_scalar_table(&namespace, "users").await;
16166 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16167 .await
16168 .expect("create_scalar_index should return a transaction id");
16169
16170 let response = namespace
16172 .alter_transaction(AlterTransactionRequest {
16173 id: Some(vec!["users".to_string(), txn_id.clone()]),
16174 actions: vec![
16175 AlterTransactionAction {
16176 set_status_action: None,
16177 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16178 key: Some("temp_key".to_string()),
16179 value: Some("temp_value".to_string()),
16180 mode: None,
16181 })),
16182 unset_property_action: None,
16183 },
16184 AlterTransactionAction {
16185 set_status_action: None,
16186 set_property_action: None,
16187 unset_property_action: Some(Box::new(AlterTransactionUnsetProperty {
16188 key: Some("temp_key".to_string()),
16189 mode: None,
16190 })),
16191 },
16192 ],
16193 ..Default::default()
16194 })
16195 .await
16196 .unwrap();
16197 assert_eq!(response.status, "SUCCEEDED");
16198 let props = response.properties.unwrap();
16199 assert!(!props.contains_key("temp_key"));
16200 }
16201
16202 #[tokio::test]
16203 async fn test_alter_transaction_invalid_status() {
16204 use lance_namespace::models::{
16205 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16206 };
16207
16208 let (namespace, _temp_dir) = create_test_namespace().await;
16209 create_scalar_table(&namespace, "users").await;
16210 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16211 .await
16212 .expect("create_scalar_index should return a transaction id");
16213
16214 let result = namespace
16215 .alter_transaction(AlterTransactionRequest {
16216 id: Some(vec!["users".to_string(), txn_id.clone()]),
16217 actions: vec![AlterTransactionAction {
16218 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16219 status: Some("InvalidStatus".to_string()),
16220 })),
16221 set_property_action: None,
16222 unset_property_action: None,
16223 }],
16224 ..Default::default()
16225 })
16226 .await;
16227 assert!(result.is_err());
16228 }
16229
16230 #[tokio::test]
16231 async fn test_alter_transaction_not_found() {
16232 use lance_namespace::models::{
16233 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16234 };
16235
16236 let (namespace, _temp_dir) = create_test_namespace().await;
16237 create_scalar_table(&namespace, "users").await;
16238
16239 let result = namespace
16241 .alter_transaction(AlterTransactionRequest {
16242 id: Some(vec!["users".to_string(), "non_existent_txn".to_string()]),
16243 actions: vec![AlterTransactionAction {
16244 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16245 status: Some("Canceled".to_string()),
16246 })),
16247 set_property_action: None,
16248 unset_property_action: None,
16249 }],
16250 ..Default::default()
16251 })
16252 .await;
16253 assert!(result.is_err());
16254 }
16255
16256 #[tokio::test]
16257 async fn test_alter_transaction_missing_id() {
16258 use lance_namespace::models::{
16259 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16260 };
16261
16262 let (namespace, _temp_dir) = create_test_namespace().await;
16263
16264 let result = namespace
16266 .alter_transaction(AlterTransactionRequest {
16267 id: None,
16268 actions: vec![AlterTransactionAction {
16269 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16270 status: Some("Canceled".to_string()),
16271 })),
16272 set_property_action: None,
16273 unset_property_action: None,
16274 }],
16275 ..Default::default()
16276 })
16277 .await;
16278 assert!(result.is_err());
16279
16280 let result = namespace
16282 .alter_transaction(AlterTransactionRequest {
16283 id: Some(vec!["users".to_string()]),
16284 actions: vec![AlterTransactionAction {
16285 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16286 status: Some("Canceled".to_string()),
16287 })),
16288 set_property_action: None,
16289 unset_property_action: None,
16290 }],
16291 ..Default::default()
16292 })
16293 .await;
16294 assert!(result.is_err());
16295 }
16296
16297 #[tokio::test]
16298 async fn test_alter_transaction_persists_changes() {
16299 use lance_namespace::models::{
16300 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16301 AlterTransactionSetStatus, DescribeTransactionRequest,
16302 };
16303
16304 let (namespace, _temp_dir) = create_test_namespace().await;
16305 create_scalar_table(&namespace, "users").await;
16306 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
16307
16308 let txn_id = transaction_id.expect("scalar index should produce a transaction id");
16309
16310 namespace
16312 .alter_transaction(AlterTransactionRequest {
16313 id: Some(vec!["users".to_string(), txn_id.clone()]),
16314 actions: vec![
16315 AlterTransactionAction {
16316 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16317 status: Some("Canceled".to_string()),
16318 })),
16319 set_property_action: None,
16320 unset_property_action: None,
16321 },
16322 AlterTransactionAction {
16323 set_status_action: None,
16324 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16325 key: Some("owner".to_string()),
16326 value: Some("alice".to_string()),
16327 mode: None,
16328 })),
16329 unset_property_action: None,
16330 },
16331 ],
16332 ..Default::default()
16333 })
16334 .await
16335 .unwrap();
16336
16337 let describe_resp = namespace
16340 .describe_transaction(DescribeTransactionRequest {
16341 id: Some(vec!["users".to_string(), txn_id.clone()]),
16342 ..Default::default()
16343 })
16344 .await
16345 .unwrap();
16346 let props = describe_resp.properties.expect("properties should be set");
16347 assert_eq!(props.get("owner"), Some(&"alice".to_string()));
16348 assert!(!props.contains_key("_status"));
16352
16353 let follow_up = namespace
16354 .alter_transaction(AlterTransactionRequest {
16355 id: Some(vec!["users".to_string(), txn_id.clone()]),
16356 actions: vec![],
16357 ..Default::default()
16358 })
16359 .await
16360 .unwrap();
16361 assert_eq!(follow_up.status, "Canceled");
16362 let follow_up_props = follow_up.properties.unwrap();
16363 assert_eq!(follow_up_props.get("owner"), Some(&"alice".to_string()));
16364 }
16365}