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::RQBuildParams, hnsw::builder::HnswBuildParams, ivf::IvfBuildParams, pq::PQBuildParams,
34 sq::builder::SQBuildParams,
35};
36use lance_index::{IndexType, is_system_index};
37use lance_io::object_store::throttle::is_throttle_error;
38use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
39use lance_linalg::distance::MetricType;
40use lance_table::io::commit::{ManifestNamingScheme, VERSIONS_DIR};
41use object_store::ObjectStoreExt;
42use object_store::path::Path;
43use object_store::{
44 Error as ObjectStoreError, ObjectMeta, ObjectStore as OSObjectStore, PutMode, PutOptions,
45};
46use std::collections::HashMap;
47use std::io::Cursor;
48use std::sync::{Arc, Mutex};
49use tokio::sync::OnceCell;
50
51use crate::context::DynamicContextProvider;
52use lance_namespace::models::{
53 AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest,
54 AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse,
55 AlterTransactionRequest, AlterTransactionResponse, AnalyzeTableQueryPlanRequest,
56 BatchDeleteTableVersionsRequest, BatchDeleteTableVersionsResponse,
57 BranchContents as ModelBranchContents, CountTableRowsRequest, CreateNamespaceRequest,
58 CreateNamespaceResponse, CreateTableBranchRequest, CreateTableBranchResponse,
59 CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest, CreateTableResponse,
60 CreateTableScalarIndexResponse, CreateTableTagRequest, CreateTableTagResponse,
61 CreateTableVersionRequest, CreateTableVersionResponse, DeclareTableRequest,
62 DeclareTableResponse, DeleteFromTableRequest, DeleteFromTableResponse,
63 DeleteTableBranchRequest, DeleteTableBranchResponse, DeleteTableTagRequest,
64 DeleteTableTagResponse, DescribeNamespaceRequest, DescribeNamespaceResponse,
65 DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, DescribeTableRequest,
66 DescribeTableResponse, DescribeTableVersionRequest, DescribeTableVersionResponse,
67 DescribeTransactionRequest, DescribeTransactionResponse, DropNamespaceRequest,
68 DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, DropTableRequest,
69 DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary,
70 GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest,
71 GetTableTagVersionResponse, Identity, IndexContent, InsertIntoTableRequest,
72 InsertIntoTableResponse, ListNamespacesRequest, ListNamespacesResponse,
73 ListTableBranchesRequest, ListTableBranchesResponse, ListTableIndicesRequest,
74 ListTableIndicesResponse, ListTableTagsRequest, ListTableTagsResponse,
75 ListTableVersionsRequest, ListTableVersionsResponse, ListTablesRequest, ListTablesResponse,
76 MergeInsertIntoTableRequest, MergeInsertIntoTableResponse, NamespaceExistsRequest,
77 QueryTableRequest, QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest,
78 RestoreTableResponse, TableExistsRequest, TableVersion, TagContents as ModelTagContents,
79 UpdateTableRequest, UpdateTableResponse, UpdateTableSchemaMetadataRequest,
80 UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, UpdateTableTagResponse,
81};
82
83use lance_core::{Error, Result, box_error};
84use lance_index::scalar::inverted::query::{
85 BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery,
86};
87use lance_namespace::LanceNamespace;
88use lance_namespace::error::NamespaceError;
89use lance_namespace::schema::arrow_schema_to_json;
90
91use crate::credentials::{
92 CredentialVendor, create_credential_vendor_for_location, has_credential_vendor_config,
93};
94
95#[derive(Debug, Default)]
100pub struct OpsMetrics {
101 counters: Mutex<HashMap<String, u64>>,
102}
103
104impl OpsMetrics {
105 pub fn increment(&self, operation: &str) {
107 if let Ok(mut counters) = self.counters.lock() {
108 *counters.entry(operation.to_string()).or_insert(0) += 1;
109 }
110 }
111
112 pub fn retrieve(&self) -> HashMap<String, u64> {
114 self.counters.lock().map(|c| c.clone()).unwrap_or_default()
115 }
116
117 pub fn reset(&self) {
119 if let Ok(mut counters) = self.counters.lock() {
120 counters.clear();
121 }
122 }
123}
124
125pub(crate) fn build_sql_expressions(
128 new_columns: &[lance_namespace::models::AddColumnsEntry],
129) -> Result<Vec<(String, String)>> {
130 new_columns
131 .iter()
132 .map(|col| {
133 let expression = col.expression.clone().and_then(|opt| opt).ok_or_else(|| {
136 Error::invalid_input(format!(
137 "Expression is required for new column '{}'",
138 col.name
139 ))
140 })?;
141 Ok((col.name.clone(), expression))
142 })
143 .collect()
144}
145
146pub(crate) fn build_column_alterations(
149 alterations: &[lance_namespace::models::AlterColumnsEntry],
150) -> Result<Vec<lance::dataset::ColumnAlteration>> {
151 alterations
152 .iter()
153 .map(|entry| {
154 let mut alteration = lance::dataset::ColumnAlteration::new(entry.path.clone());
155 if let Some(Some(rename)) = &entry.rename {
157 alteration = alteration.rename(rename.clone());
158 }
159 if let Some(Some(nullable)) = entry.nullable {
161 alteration = alteration.set_nullable(nullable);
162 }
163 if let Some(data_type) = &entry.data_type
165 && !data_type.is_null()
166 {
167 let type_str = data_type.as_str().ok_or_else(|| {
168 Error::invalid_input(format!(
169 "data_type for column '{}' must be a JSON string, got: {}",
170 entry.path, data_type
171 ))
172 })?;
173 let json_type =
174 lance_namespace::models::JsonArrowDataType::new(type_str.to_string());
175 let dt =
176 lance_namespace::schema::convert_json_arrow_type(&json_type).map_err(|e| {
177 Error::invalid_input(format!(
178 "Failed to parse data_type '{}' for column '{}': {}",
179 type_str, entry.path, e
180 ))
181 })?;
182 alteration = alteration.cast_to(dt);
183 }
184 Ok(alteration)
185 })
186 .collect()
187}
188
189pub(crate) struct TableStatus {
194 pub(crate) exists: bool,
196 pub(crate) is_deregistered: bool,
198 pub(crate) has_reserved_file: bool,
200}
201
202enum DirectoryIndexParams {
203 Scalar {
204 index_type: IndexType,
205 params: ScalarIndexParams,
206 },
207 Inverted(InvertedIndexParams),
208 Vector {
209 index_type: IndexType,
210 params: VectorIndexParams,
211 },
212}
213
214impl DirectoryIndexParams {
215 fn index_type(&self) -> IndexType {
216 match self {
217 Self::Scalar { index_type, .. } | Self::Vector { index_type, .. } => *index_type,
218 Self::Inverted(_) => IndexType::Inverted,
219 }
220 }
221
222 fn params(&self) -> &dyn IndexParams {
223 match self {
224 Self::Scalar { params, .. } => params,
225 Self::Inverted(params) => params,
226 Self::Vector { params, .. } => params,
227 }
228 }
229}
230
231#[derive(Clone)]
266pub struct DirectoryNamespaceBuilder {
267 root: String,
268 storage_options: Option<HashMap<String, String>>,
269 session: Option<Arc<Session>>,
270 manifest_enabled: bool,
271 dir_listing_enabled: bool,
272 inline_optimization_enabled: bool,
273 table_version_tracking_enabled: bool,
274 dir_listing_to_manifest_migration_enabled: bool,
279 credential_vendor_properties: HashMap<String, String>,
280 context_provider: Option<Arc<dyn DynamicContextProvider>>,
281 commit_retries: Option<u32>,
282 vend_input_storage_options: bool,
285 vend_input_storage_options_refresh_interval_millis: Option<u64>,
290 ops_metrics_enabled: bool,
292}
293
294impl std::fmt::Debug for DirectoryNamespaceBuilder {
295 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296 f.debug_struct("DirectoryNamespaceBuilder")
297 .field("root", &self.root)
298 .field("storage_options", &self.storage_options)
299 .field("manifest_enabled", &self.manifest_enabled)
300 .field("dir_listing_enabled", &self.dir_listing_enabled)
301 .field(
302 "inline_optimization_enabled",
303 &self.inline_optimization_enabled,
304 )
305 .field(
306 "table_version_tracking_enabled",
307 &self.table_version_tracking_enabled,
308 )
309 .field(
310 "dir_listing_to_manifest_migration_enabled",
311 &self.dir_listing_to_manifest_migration_enabled,
312 )
313 .field(
314 "context_provider",
315 &self.context_provider.as_ref().map(|_| "Some(...)"),
316 )
317 .field(
318 "vend_input_storage_options",
319 &self.vend_input_storage_options,
320 )
321 .field(
322 "vend_input_storage_options_refresh_interval_millis",
323 &self.vend_input_storage_options_refresh_interval_millis,
324 )
325 .field("ops_metrics_enabled", &self.ops_metrics_enabled)
326 .finish()
327 }
328}
329
330impl DirectoryNamespaceBuilder {
331 pub fn new(root: impl Into<String>) -> Self {
337 Self {
338 root: root.into().trim_end_matches('/').to_string(),
339 storage_options: None,
340 session: None,
341 manifest_enabled: true,
342 dir_listing_enabled: true, inline_optimization_enabled: true,
344 table_version_tracking_enabled: false, dir_listing_to_manifest_migration_enabled: false, credential_vendor_properties: HashMap::new(),
347 context_provider: None,
348 commit_retries: None,
349 vend_input_storage_options: false,
350 vend_input_storage_options_refresh_interval_millis: None,
351 ops_metrics_enabled: false,
352 }
353 }
354
355 pub fn manifest_enabled(mut self, enabled: bool) -> Self {
360 self.manifest_enabled = enabled;
361 self
362 }
363
364 pub fn dir_listing_enabled(mut self, enabled: bool) -> Self {
369 self.dir_listing_enabled = enabled;
370 self
371 }
372
373 pub fn dir_listing_to_manifest_migration_enabled(mut self, enabled: bool) -> Self {
380 self.dir_listing_to_manifest_migration_enabled = enabled;
381 self
382 }
383
384 pub fn inline_optimization_enabled(mut self, enabled: bool) -> Self {
389 self.inline_optimization_enabled = enabled;
390 self
391 }
392
393 pub fn table_version_tracking_enabled(mut self, enabled: bool) -> Self {
401 self.table_version_tracking_enabled = enabled;
402 self
403 }
404
405 pub fn from_properties(
473 properties: HashMap<String, String>,
474 session: Option<Arc<Session>>,
475 ) -> Result<Self> {
476 let root = properties.get("root").cloned().ok_or_else(|| {
478 lance_core::Error::from(NamespaceError::InvalidInput {
479 message: "Missing required property 'root' for directory namespace".to_string(),
480 })
481 })?;
482
483 let storage_options: HashMap<String, String> = properties
485 .iter()
486 .filter_map(|(k, v)| {
487 k.strip_prefix("storage.")
488 .map(|key| (key.to_string(), v.clone()))
489 })
490 .collect();
491
492 let storage_options = if storage_options.is_empty() {
493 None
494 } else {
495 Some(storage_options)
496 };
497
498 let manifest_enabled = properties
500 .get("manifest_enabled")
501 .and_then(|v| v.parse::<bool>().ok())
502 .unwrap_or(true);
503
504 let dir_listing_enabled = properties
506 .get("dir_listing_enabled")
507 .and_then(|v| v.parse::<bool>().ok())
508 .unwrap_or(true);
509
510 let inline_optimization_enabled = properties
512 .get("inline_optimization_enabled")
513 .and_then(|v| v.parse::<bool>().ok())
514 .unwrap_or(true);
515
516 let table_version_tracking_enabled = properties
518 .get("table_version_tracking_enabled")
519 .and_then(|v| v.parse::<bool>().ok())
520 .unwrap_or(false);
521
522 let dir_listing_to_manifest_migration_enabled = properties
524 .get("dir_listing_to_manifest_migration_enabled")
525 .and_then(|v| v.parse::<bool>().ok())
526 .unwrap_or(false);
527
528 let credential_vendor_properties: HashMap<String, String> = properties
532 .iter()
533 .filter_map(|(k, v)| {
534 k.strip_prefix("credential_vendor.")
535 .map(|key| (key.to_string(), v.clone()))
536 })
537 .collect();
538
539 let commit_retries = properties
540 .get("commit_retries")
541 .and_then(|v| v.parse::<u32>().ok());
542
543 let vend_input_storage_options = properties
545 .get("vend_input_storage_options")
546 .and_then(|v| v.parse::<bool>().ok())
547 .unwrap_or(false);
548
549 let vend_input_storage_options_refresh_interval_millis = properties
551 .get("vend_input_storage_options_refresh_interval_millis")
552 .and_then(|v| v.parse::<u64>().ok());
553
554 let ops_metrics_enabled = properties
556 .get("ops_metrics_enabled")
557 .and_then(|v| v.parse::<bool>().ok())
558 .unwrap_or(false);
559
560 Ok(Self {
561 root: root.trim_end_matches('/').to_string(),
562 storage_options,
563 session,
564 manifest_enabled,
565 dir_listing_enabled,
566 inline_optimization_enabled,
567 table_version_tracking_enabled,
568 dir_listing_to_manifest_migration_enabled,
569 credential_vendor_properties,
570 context_provider: None,
571 commit_retries,
572 vend_input_storage_options,
573 vend_input_storage_options_refresh_interval_millis,
574 ops_metrics_enabled,
575 })
576 }
577
578 pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
585 self.storage_options
586 .get_or_insert_with(HashMap::new)
587 .insert(key.into(), value.into());
588 self
589 }
590
591 pub fn storage_options(mut self, options: HashMap<String, String>) -> Self {
597 self.storage_options
598 .get_or_insert_with(HashMap::new)
599 .extend(options);
600 self
601 }
602
603 pub fn session(mut self, session: Arc<Session>) -> Self {
613 self.session = Some(session);
614 self
615 }
616
617 pub fn commit_retries(mut self, retries: u32) -> Self {
620 self.commit_retries = Some(retries);
621 self
622 }
623
624 pub fn credential_vendor_property(
652 mut self,
653 key: impl Into<String>,
654 value: impl Into<String>,
655 ) -> Self {
656 self.credential_vendor_properties
657 .insert(key.into(), value.into());
658 self
659 }
660
661 pub fn credential_vendor_properties(mut self, properties: HashMap<String, String>) -> Self {
669 self.credential_vendor_properties.extend(properties);
670 self
671 }
672
673 pub fn context_provider(mut self, provider: Arc<dyn DynamicContextProvider>) -> Self {
683 self.context_provider = Some(provider);
684 self
685 }
686
687 pub fn vend_input_storage_options(mut self, enabled: bool) -> Self {
696 self.vend_input_storage_options = enabled;
697 self
698 }
699
700 pub fn vend_input_storage_options_refresh_interval_millis(
712 mut self,
713 interval_millis: u64,
714 ) -> Self {
715 self.vend_input_storage_options_refresh_interval_millis = Some(interval_millis);
716 self
717 }
718
719 pub fn ops_metrics_enabled(mut self, enabled: bool) -> Self {
727 self.ops_metrics_enabled = enabled;
728 self
729 }
730
731 pub async fn build(self) -> Result<DirectoryNamespace> {
744 let (object_store, base_path) =
745 Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?;
746
747 let manifest_ns = if self.manifest_enabled {
748 match manifest::ManifestNamespace::open_from_directory(
749 self.root.clone(),
750 self.storage_options.clone(),
751 self.session.clone(),
752 object_store.clone(),
753 base_path.clone(),
754 self.dir_listing_enabled,
755 self.inline_optimization_enabled,
756 self.commit_retries,
757 )
758 .await
759 {
760 Ok(ns) => Some(Arc::new(ns)),
761 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
762 return Err(e);
766 }
767 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => {
768 log::debug!("Manifest namespace does not exist yet: {}", e);
769 None
770 }
771 Err(e) => return Err(e),
772 }
773 } else {
774 None
775 };
776 let manifest_cell = OnceCell::new();
777 if let Some(manifest_ns) = manifest_ns {
778 let _ = manifest_cell.set(manifest_ns);
779 }
780
781 let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties)
783 {
784 create_credential_vendor_for_location(&self.root, &self.credential_vendor_properties)
785 .await?
786 .map(Arc::from)
787 } else {
788 None
789 };
790
791 let ops_metrics = if self.ops_metrics_enabled {
792 Some(Arc::new(OpsMetrics::default()))
793 } else {
794 None
795 };
796
797 Ok(DirectoryNamespace {
798 root: self.root,
799 storage_options: self.storage_options,
800 session: self.session,
801 object_store,
802 base_path,
803 manifest_ns: manifest_cell,
804 write_manifest_ns: OnceCell::new(),
805 manifest_enabled: self.manifest_enabled,
806 dir_listing_enabled: self.dir_listing_enabled,
807 inline_optimization_enabled: self.inline_optimization_enabled,
808 commit_retries: self.commit_retries,
809 dir_listing_to_manifest_migration_enabled: self
810 .dir_listing_to_manifest_migration_enabled,
811 table_version_tracking_enabled: self.table_version_tracking_enabled,
812 credential_vendor,
813 context_provider: self.context_provider,
814 vend_input_storage_options: self.vend_input_storage_options,
815 vend_input_storage_options_refresh_interval_millis: self
816 .vend_input_storage_options_refresh_interval_millis,
817 ops_metrics,
818 })
819 }
820
821 async fn initialize_object_store(
823 root: &str,
824 storage_options: &Option<HashMap<String, String>>,
825 session: &Option<Arc<Session>>,
826 ) -> Result<(Arc<ObjectStore>, Path)> {
827 let accessor = storage_options.clone().map(|opts| {
829 Arc::new(lance_io::object_store::StorageOptionsAccessor::with_static_options(opts))
830 });
831 let params = ObjectStoreParams {
832 storage_options_accessor: accessor,
833 ..Default::default()
834 };
835
836 let registry = if let Some(session) = session {
838 session.store_registry()
839 } else {
840 Arc::new(ObjectStoreRegistry::default())
841 };
842
843 let (object_store, base_path) = ObjectStore::from_uri_and_params(registry, root, ¶ms)
845 .await
846 .map_err(|e| {
847 lance_core::Error::from(NamespaceError::Internal {
848 message: format!("Failed to create object store: {:?}", e),
849 })
850 })?;
851
852 Ok((object_store, base_path))
853 }
854}
855
856pub struct DirectoryNamespace {
880 root: String,
881 storage_options: Option<HashMap<String, String>>,
882 session: Option<Arc<Session>>,
883 object_store: Arc<ObjectStore>,
884 base_path: Path,
885 manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
886 write_manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
887 manifest_enabled: bool,
888 dir_listing_enabled: bool,
889 inline_optimization_enabled: bool,
890 commit_retries: Option<u32>,
891 dir_listing_to_manifest_migration_enabled: bool,
895 table_version_tracking_enabled: bool,
898 credential_vendor: Option<Arc<dyn CredentialVendor>>,
901 #[allow(dead_code)]
904 context_provider: Option<Arc<dyn DynamicContextProvider>>,
905 vend_input_storage_options: bool,
907 vend_input_storage_options_refresh_interval_millis: Option<u64>,
910 ops_metrics: Option<Arc<OpsMetrics>>,
912}
913
914impl std::fmt::Debug for DirectoryNamespace {
915 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
916 write!(f, "{}", self.namespace_id())
917 }
918}
919
920impl std::fmt::Display for DirectoryNamespace {
921 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
922 write!(f, "{}", self.namespace_id())
923 }
924}
925
926struct ExistingTableVersionResolve<'a> {
928 staging_path: &'a Path,
929 final_path: &'a Path,
930 version: u64,
931 table_uri: &'a str,
932 final_meta: &'a ObjectMeta,
933 request_manifest_size: Option<i64>,
934}
935
936struct TableDeleteEntry {
939 table_id: Option<Vec<String>>,
940 ranges: Vec<(i64, i64)>,
941}
942
943#[derive(Debug, Clone, Default)]
953struct TransactionAlteration {
954 status: Option<String>,
956 properties: HashMap<String, String>,
960 removed_properties: std::collections::HashSet<String>,
964}
965
966impl TransactionAlteration {
967 const F_STATUS: &'static str = "status";
969 const F_PROPERTIES: &'static str = "properties";
970 const F_REMOVED_PROPERTIES: &'static str = "removed_properties";
971
972 fn to_json_bytes(&self) -> serde_json::Result<Vec<u8>> {
979 serde_json::to_vec(&serde_json::json!({
980 Self::F_STATUS: self.status,
981 Self::F_PROPERTIES: self.properties,
982 Self::F_REMOVED_PROPERTIES: self.removed_properties,
983 }))
984 }
985
986 fn from_json_slice(bytes: &[u8]) -> serde_json::Result<Self> {
991 let mut obj: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes)?;
992 Ok(Self {
993 status: serde_json::from_value(
994 obj.remove(Self::F_STATUS)
995 .unwrap_or(serde_json::Value::Null),
996 )?,
997 properties: serde_json::from_value(
998 obj.remove(Self::F_PROPERTIES)
999 .unwrap_or(serde_json::Value::Null),
1000 )
1001 .unwrap_or_default(),
1002 removed_properties: serde_json::from_value(
1003 obj.remove(Self::F_REMOVED_PROPERTIES)
1004 .unwrap_or(serde_json::Value::Null),
1005 )
1006 .unwrap_or_default(),
1007 })
1008 }
1009}
1010
1011impl DirectoryNamespace {
1012 fn manifest_ns_for_read(&self) -> Option<&Arc<manifest::ManifestNamespace>> {
1013 self.write_manifest_ns
1014 .get()
1015 .or_else(|| self.manifest_ns.get())
1016 }
1017
1018 async fn manifest_ns_for_write(&self) -> Result<Option<Arc<manifest::ManifestNamespace>>> {
1019 if !self.manifest_enabled {
1020 return Ok(None);
1021 }
1022
1023 let manifest_ns = self
1024 .write_manifest_ns
1025 .get_or_try_init(|| async {
1026 manifest::ManifestNamespace::from_directory(
1027 self.root.clone(),
1028 self.storage_options.clone(),
1029 self.session.clone(),
1030 self.object_store.clone(),
1031 self.base_path.clone(),
1032 self.dir_listing_enabled,
1033 self.inline_optimization_enabled,
1034 self.commit_retries,
1035 )
1036 .await
1037 .map(Arc::new)
1038 })
1039 .await?;
1040 Ok(Some(manifest_ns.clone()))
1041 }
1042
1043 async fn ensure_read_manifest(&self) -> Result<()> {
1055 if !self.manifest_enabled
1056 || self.manifest_ns.get().is_some()
1057 || self.write_manifest_ns.get().is_some()
1058 {
1059 return Ok(());
1060 }
1061 match self
1062 .manifest_ns
1063 .get_or_try_init(|| async {
1064 manifest::ManifestNamespace::open_from_directory(
1065 self.root.clone(),
1066 self.storage_options.clone(),
1067 self.session.clone(),
1068 self.object_store.clone(),
1069 self.base_path.clone(),
1070 self.dir_listing_enabled,
1071 self.inline_optimization_enabled,
1072 self.commit_retries,
1073 )
1074 .await
1075 .map(Arc::new)
1076 })
1077 .await
1078 {
1079 Ok(_) => Ok(()),
1080 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(()),
1085 Err(e) => Err(e),
1086 }
1087 }
1088
1089 fn child_namespace_requires_manifest_error(&self) -> Error {
1090 if self.manifest_enabled {
1091 NamespaceError::NamespaceNotFound {
1092 message: "Child namespace reads require an existing __manifest dataset".to_string(),
1093 }
1094 .into()
1095 } else {
1096 NamespaceError::Unsupported {
1097 message: "Child namespaces are only supported when manifest mode is enabled"
1098 .to_string(),
1099 }
1100 .into()
1101 }
1102 }
1103
1104 fn apply_pagination(
1117 names: &mut Vec<String>,
1118 page_token: Option<String>,
1119 limit: Option<i32>,
1120 ) -> Option<String> {
1121 names.sort();
1123
1124 if let Some(start_after) = page_token {
1126 if let Some(index) = names
1127 .iter()
1128 .position(|name| name.as_str() > start_after.as_str())
1129 {
1130 names.drain(0..index);
1131 } else {
1132 names.clear();
1133 }
1134 }
1135
1136 if let Some(limit) = limit
1138 && limit >= 0
1139 {
1140 let limit = limit as usize;
1141 if names.len() > limit {
1142 let next_page_token = if limit > 0 {
1143 Some(names[limit - 1].clone())
1144 } else {
1145 None
1146 };
1147 names.truncate(limit);
1148 return next_page_token;
1149 }
1150 }
1151
1152 None
1153 }
1154
1155 async fn list_directory_tables(&self) -> Result<Vec<String>> {
1157 let mut tables = Vec::new();
1158 let entries = self
1159 .object_store
1160 .read_dir(self.base_path.clone())
1161 .await
1162 .map_err(|e| {
1163 lance_core::Error::from(NamespaceError::Internal {
1164 message: format!("Failed to list directory: {:?}", e),
1165 })
1166 })?;
1167
1168 for entry in entries {
1169 let path = entry.trim_end_matches('/');
1170 if !path.ends_with(".lance") {
1171 continue;
1172 }
1173
1174 let table_name = &path[..path.len() - 6];
1175
1176 let status = self.check_table_status(table_name).await?;
1178 if status.is_deregistered {
1179 continue;
1180 }
1181
1182 tables.push(table_name.to_string());
1183 }
1184
1185 Ok(tables)
1186 }
1187
1188 fn validate_root_namespace_id(id: &Option<Vec<String>>) -> Result<()> {
1190 if let Some(id) = id
1191 && !id.is_empty()
1192 {
1193 return Err(NamespaceError::Unsupported {
1194 message: format!(
1195 "Directory namespace only supports root namespace operations, but got namespace ID: {:?}. Expected empty ID.",
1196 id
1197 ),
1198 }
1199 .into());
1200 }
1201 Ok(())
1202 }
1203
1204 fn table_name_from_id(id: &Option<Vec<String>>) -> Result<String> {
1206 let id = id.as_ref().ok_or_else(|| {
1207 lance_core::Error::from(NamespaceError::InvalidInput {
1208 message: "Directory namespace table ID cannot be empty".to_string(),
1209 })
1210 })?;
1211
1212 if id.len() != 1 {
1213 return Err(NamespaceError::Unsupported {
1214 message: format!(
1215 "Multi-level table IDs are only supported when manifest mode is enabled, but got: {:?}",
1216 id
1217 ),
1218 }
1219 .into());
1220 }
1221
1222 Ok(id[0].clone())
1223 }
1224
1225 fn format_table_id(table_id: &[String]) -> String {
1226 format!(
1227 "table id '{}'",
1228 manifest::ManifestNamespace::str_object_id(table_id)
1229 )
1230 }
1231
1232 fn format_table_id_from_request(id: &Option<Vec<String>>) -> String {
1233 id.as_ref()
1234 .map(|table_id| Self::format_table_id(table_id))
1235 .unwrap_or_else(|| "table id '<unknown>'".to_string())
1236 }
1237
1238 async fn resolve_table_location(&self, id: &Option<Vec<String>>) -> Result<String> {
1239 let mut describe_req = DescribeTableRequest::new();
1240 describe_req.id = id.clone();
1241 describe_req.load_detailed_metadata = Some(false);
1242
1243 let describe_resp = self.describe_table_impl(describe_req).await?;
1245
1246 describe_resp.location.ok_or_else(|| {
1247 lance_core::Error::from(NamespaceError::TableNotFound {
1248 message: format!("Table location not found for: {:?}", id),
1249 })
1250 })
1251 }
1252
1253 fn map_tag_error(err: lance_core::Error, tag: &str, table_uri: &str) -> lance_core::Error {
1256 match err {
1257 lance_core::Error::RefNotFound { .. } => NamespaceError::TableTagNotFound {
1258 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1259 }
1260 .into(),
1261 lance_core::Error::RefConflict { .. } => NamespaceError::TableTagAlreadyExists {
1262 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1263 }
1264 .into(),
1265 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1266 message: format!("invalid tag '{}': {}", tag, message),
1267 }
1268 .into(),
1269 lance_core::Error::VersionNotFound { message } => {
1270 NamespaceError::TableVersionNotFound {
1271 message: format!(
1272 "version referenced by tag '{}' not found for table at '{}': {}",
1273 tag, table_uri, message
1274 ),
1275 }
1276 .into()
1277 }
1278 other => NamespaceError::Internal {
1279 message: format!(
1280 "tag operation failed for tag '{}' on table at '{}': {}",
1281 tag, table_uri, other
1282 ),
1283 }
1284 .into(),
1285 }
1286 }
1287
1288 fn map_branch_error(
1294 err: lance_core::Error,
1295 branch: &str,
1296 table_uri: &str,
1297 ) -> lance_core::Error {
1298 match err {
1299 lance_core::Error::RefNotFound { .. } => NamespaceError::TableBranchNotFound {
1300 message: format!("branch '{}' for table at '{}'", branch, table_uri),
1301 }
1302 .into(),
1303 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1304 message: format!("invalid branch '{}': {}", branch, message),
1305 }
1306 .into(),
1307 lance_core::Error::VersionNotFound { message } => {
1308 NamespaceError::TableVersionNotFound {
1309 message: format!(
1310 "source version for branch '{}' not found for table at '{}': {}",
1311 branch, table_uri, message
1312 ),
1313 }
1314 .into()
1315 }
1316 other => NamespaceError::Internal {
1317 message: format!(
1318 "branch operation failed for branch '{}' on table at '{}': {}",
1319 branch, table_uri, other
1320 ),
1321 }
1322 .into(),
1323 }
1324 }
1325
1326 fn map_mutation_error(
1338 err: lance_core::Error,
1339 operation: &str,
1340 table_uri: &str,
1341 ) -> lance_core::Error {
1342 let detail = err.to_string();
1343 let ns_err = match &err {
1344 lance_core::Error::InvalidInput { .. }
1345 | lance_core::Error::Unprocessable { .. }
1346 | lance_core::Error::InvalidRef { .. } => NamespaceError::InvalidInput {
1347 message: format!(
1348 "Invalid input for {} on table at '{}': {}",
1349 operation, table_uri, detail
1350 ),
1351 },
1352 lance_core::Error::NotFound { .. } | lance_core::Error::DatasetNotFound { .. } => {
1353 NamespaceError::TableNotFound {
1354 message: format!(
1355 "Table at '{}' not found while running {}: {}",
1356 table_uri, operation, detail
1357 ),
1358 }
1359 }
1360 lance_core::Error::SchemaMismatch { .. } | lance_core::Error::Schema { .. } => {
1361 NamespaceError::TableSchemaValidationError {
1362 message: format!(
1363 "Schema validation failed for {} on table at '{}': {}",
1364 operation, table_uri, detail
1365 ),
1366 }
1367 }
1368 lance_core::Error::CommitConflict { .. } => NamespaceError::Throttling {
1372 message: format!(
1373 "Too many concurrent writes for {} on table at '{}', please retry later: {}",
1374 operation, table_uri, detail
1375 ),
1376 },
1377 lance_core::Error::TooMuchWriteContention { .. }
1381 | lance_core::Error::RetryableCommitConflict { .. }
1382 | lance_core::Error::IncompatibleTransaction { .. }
1383 | lance_core::Error::VersionConflict { .. } => NamespaceError::ConcurrentModification {
1384 message: format!(
1385 "Concurrent modification detected for {} on table at '{}': {}",
1386 operation, table_uri, detail
1387 ),
1388 },
1389 lance_core::Error::NotSupported { .. } => NamespaceError::Unsupported {
1390 message: format!(
1391 "{} is not supported on table at '{}': {}",
1392 operation, table_uri, detail
1393 ),
1394 },
1395 _ => NamespaceError::Internal {
1396 message: format!(
1397 "Failed to run {} on table at '{}': {}",
1398 operation, table_uri, detail
1399 ),
1400 },
1401 };
1402 ns_err.into()
1403 }
1404
1405 async fn table_has_actual_manifests(&self, table_name: &str) -> Result<bool> {
1406 manifest::ManifestNamespace::path_has_actual_manifests(
1407 &self.object_store,
1408 &self.table_path(table_name),
1409 )
1410 .await
1411 }
1412
1413 async fn filter_declared_tables(
1414 &self,
1415 tables: Vec<String>,
1416 include_declared: bool,
1417 ) -> Result<Vec<String>> {
1418 if include_declared {
1419 return Ok(tables);
1420 }
1421
1422 let mut stream = futures::stream::iter(tables.into_iter().map(|table_name| async move {
1423 if self.table_has_actual_manifests(&table_name).await? {
1427 Ok::<Option<String>, Error>(Some(table_name))
1428 } else {
1429 Ok::<Option<String>, Error>(None)
1430 }
1431 }))
1432 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1433
1434 let mut filtered = Vec::new();
1435 while let Some(result) = stream.next().await {
1436 if let Some(table_name) = result? {
1437 filtered.push(table_name);
1438 }
1439 }
1440 Ok(filtered)
1441 }
1442
1443 fn ipc_reader_from_request_data(
1444 request_data: &Bytes,
1445 operation: &str,
1446 ) -> Result<(
1447 Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1448 usize,
1449 )> {
1450 if request_data.is_empty() {
1451 return Err(NamespaceError::InvalidInput {
1452 message: format!(
1453 "Request data (Arrow IPC stream) is required for {}",
1454 operation
1455 ),
1456 }
1457 .into());
1458 }
1459
1460 let cursor = Cursor::new(request_data.as_ref());
1461 let stream_reader =
1462 StreamReader::try_new(cursor, None).map_err(|e| NamespaceError::InvalidInput {
1463 message: format!("Invalid Arrow IPC stream: {}", e),
1464 })?;
1465 let arrow_schema = stream_reader.schema();
1466
1467 let mut num_rows = 0usize;
1468 let mut batches = Vec::new();
1469 for batch_result in stream_reader {
1470 let batch = batch_result.map_err(|e| NamespaceError::Internal {
1471 message: format!("Failed to read batch from IPC stream: {}", e),
1472 })?;
1473 num_rows += batch.num_rows();
1474 batches.push(batch);
1475 }
1476
1477 let reader: Box<dyn arrow::record_batch::RecordBatchReader + Send> = if batches.is_empty() {
1478 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
1479 Box::new(RecordBatchIterator::new(vec![Ok(batch)], arrow_schema))
1480 } else {
1481 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
1482 Box::new(RecordBatchIterator::new(batch_results, arrow_schema))
1483 };
1484
1485 Ok((reader, num_rows))
1486 }
1487
1488 async fn table_uri_has_actual_manifests(&self, table_uri: &str) -> Result<bool> {
1489 let table_path = self.object_store_path_from_uri(table_uri)?;
1490 manifest::ManifestNamespace::path_has_actual_manifests(&self.object_store, &table_path)
1491 .await
1492 }
1493
1494 fn object_store_path_from_uri(&self, uri: &str) -> Result<Path> {
1495 let registry = self
1496 .session
1497 .as_ref()
1498 .map(|session| session.store_registry())
1499 .unwrap_or_else(|| Arc::new(ObjectStoreRegistry::default()));
1500 ObjectStore::extract_path_from_uri(registry, uri)
1501 }
1502
1503 fn normalized_branch(branch: Option<&str>) -> Result<Option<&str>> {
1508 match branch.filter(|b| !b.is_empty() && *b != "main") {
1509 Some(branch) => {
1510 check_valid_branch(branch).map_err(|e| {
1511 lance_core::Error::from(NamespaceError::InvalidInput {
1512 message: format!("invalid branch name '{}': {}", branch, e),
1513 })
1514 })?;
1515 Ok(Some(branch))
1516 }
1517 None => Ok(None),
1518 }
1519 }
1520
1521 async fn open_validated_branch(&self, table_uri: &str, branch: &str) -> Result<Dataset> {
1522 let dataset = self
1523 .configured_builder(table_uri)
1524 .with_branch(branch, None)
1525 .load()
1526 .await
1527 .map_err(|e| {
1528 let message = format!(
1529 "branch '{}' not found for table at '{}': {}",
1530 branch, table_uri, e
1531 );
1532 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1533 })?;
1534 dataset.branches().get(branch).await.map_err(|e| {
1535 Self::map_open_error(
1536 e,
1537 NamespaceError::TableNotFound {
1538 message: format!("branch '{}' not found for table at '{}'", branch, table_uri),
1539 },
1540 )
1541 })?;
1542 Ok(dataset)
1543 }
1544
1545 async fn resolve_branch_location(&self, table_uri: &str, branch: &str) -> Result<String> {
1546 Ok(self
1547 .open_validated_branch(table_uri, branch)
1548 .await?
1549 .branch_location()
1550 .uri)
1551 }
1552
1553 async fn resolve_branch_for_commit(
1563 &self,
1564 table_uri: &str,
1565 branch: &str,
1566 ) -> Result<(String, Path, Option<u64>)> {
1567 let main = self
1568 .configured_builder(table_uri)
1569 .load()
1570 .await
1571 .map_err(|e| {
1572 let message = format!("table at '{}' not found: {}", table_uri, e);
1573 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1574 })?;
1575 let branch_location = main.branch_location().find_branch(Some(branch))?;
1576 match main.branches().get(branch).await {
1577 Ok(contents) => Ok((
1578 branch_location.uri,
1579 branch_location.path,
1580 Some(contents.parent_version),
1581 )),
1582 Err(lance_core::Error::RefNotFound { .. }) => {
1583 if self
1584 .branch_has_committed_versions(&branch_location.path)
1585 .await?
1586 {
1587 return Err(NamespaceError::TableNotFound {
1588 message: format!(
1589 "branch '{}' not found for table at '{}'",
1590 branch, table_uri
1591 ),
1592 }
1593 .into());
1594 }
1595 Ok((branch_location.uri, branch_location.path, None))
1596 }
1597 Err(e) => Err(e),
1598 }
1599 }
1600
1601 async fn branch_has_committed_versions(&self, branch_path: &Path) -> Result<bool> {
1602 Ok(!self
1603 .list_versions_under(branch_path, false, Some(1))
1604 .await?
1605 .is_empty())
1606 }
1607
1608 fn validate_dir_only_properties(
1609 properties: Option<&HashMap<String, String>>,
1610 operation: &str,
1611 ) -> Result<()> {
1612 if properties.is_some_and(|properties| !properties.is_empty()) {
1616 return Err(NamespaceError::Unsupported {
1617 message: format!(
1618 "{} with non-empty table properties requires manifest_enabled=true",
1619 operation
1620 ),
1621 }
1622 .into());
1623 }
1624 Ok(())
1625 }
1626
1627 async fn write_reader_to_table(
1628 &self,
1629 table_uri: &str,
1630 reader: Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1631 mode: WriteMode,
1632 extra_storage_options: Option<HashMap<String, String>>,
1633 ) -> Result<Dataset> {
1634 let mut merged_storage_options = self.storage_options.clone().unwrap_or_default();
1637 if let Some(extra_storage_options) = extra_storage_options {
1638 merged_storage_options.extend(extra_storage_options);
1639 }
1640 let store_params = (!merged_storage_options.is_empty()).then(|| ObjectStoreParams {
1641 storage_options_accessor: Some(Arc::new(
1642 lance_io::object_store::StorageOptionsAccessor::with_static_options(
1643 merged_storage_options,
1644 ),
1645 )),
1646 ..Default::default()
1647 });
1648
1649 let write_params = WriteParams {
1650 mode,
1651 store_params,
1652 session: self.session.clone(),
1653 ..Default::default()
1654 };
1655
1656 let dataset = Dataset::write(reader, table_uri, Some(write_params))
1657 .await
1658 .map_err(|e| NamespaceError::Internal {
1659 message: format!("Failed to write table at '{}': {}", table_uri, e),
1660 })?;
1661
1662 Ok(dataset)
1663 }
1664
1665 fn manifest_version_from_filename(filename: &str) -> Option<u64> {
1669 ManifestNamingScheme::detect_scheme(filename)?.parse_version(filename)
1670 }
1671
1672 fn create_table_version_response(
1674 version: u64,
1675 final_path: &Path,
1676 final_meta: &ObjectMeta,
1677 ) -> CreateTableVersionResponse {
1678 CreateTableVersionResponse {
1679 transaction_id: None,
1680 version: Some(Box::new(TableVersion {
1681 version: version as i64,
1682 manifest_path: final_path.to_string(),
1683 manifest_size: Some(final_meta.size as i64),
1684 e_tag: final_meta.e_tag.clone(),
1685 timestamp_millis: None,
1686 metadata: None,
1687 })),
1688 ..Default::default()
1689 }
1690 }
1691
1692 async fn staging_matches_final_manifest(
1700 &self,
1701 staging_path: &Path,
1702 final_path: &Path,
1703 final_meta: &ObjectMeta,
1704 request_manifest_size: Option<i64>,
1705 ) -> Result<bool> {
1706 if let Some(size) = request_manifest_size
1707 && size != final_meta.size as i64
1708 {
1709 return Ok(false);
1710 }
1711
1712 let staging_bytes = match self.object_store.inner.get(staging_path).await {
1713 Ok(r) => r.bytes().await.map_err(|e| {
1714 lance_core::Error::from(NamespaceError::Internal {
1715 message: format!(
1716 "Failed to read staging manifest at '{}': {}",
1717 staging_path, e
1718 ),
1719 })
1720 })?,
1721 Err(ObjectStoreError::NotFound { .. }) => return Ok(false),
1722 Err(e) => {
1723 return Err(lance_core::Error::from(NamespaceError::Internal {
1724 message: format!(
1725 "Failed to read staging manifest at '{}': {}",
1726 staging_path, e
1727 ),
1728 }));
1729 }
1730 };
1731
1732 let final_bytes = self
1733 .object_store
1734 .inner
1735 .get(final_path)
1736 .await
1737 .map_err(|e| {
1738 lance_core::Error::from(NamespaceError::Internal {
1739 message: format!(
1740 "Failed to read existing version manifest at '{}': {}",
1741 final_path, e
1742 ),
1743 })
1744 })?
1745 .bytes()
1746 .await
1747 .map_err(|e| {
1748 lance_core::Error::from(NamespaceError::Internal {
1749 message: format!(
1750 "Failed to read existing version manifest bytes at '{}': {}",
1751 final_path, e
1752 ),
1753 })
1754 })?;
1755
1756 Ok(staging_bytes.as_ref() == final_bytes.as_ref())
1757 }
1758
1759 async fn resolve_existing_table_version(
1761 &self,
1762 args: ExistingTableVersionResolve<'_>,
1763 ) -> Result<CreateTableVersionResponse> {
1764 if self
1765 .staging_matches_final_manifest(
1766 args.staging_path,
1767 args.final_path,
1768 args.final_meta,
1769 args.request_manifest_size,
1770 )
1771 .await?
1772 {
1773 if let Err(e) = self.object_store.inner.delete(args.staging_path).await {
1775 log::warn!(
1776 "Failed to delete staging manifest at '{}': {:?}",
1777 args.staging_path,
1778 e
1779 );
1780 }
1781 return Ok(Self::create_table_version_response(
1782 args.version,
1783 args.final_path,
1784 args.final_meta,
1785 ));
1786 }
1787
1788 Err(lance_core::Error::from(
1789 NamespaceError::ConcurrentModification {
1790 message: format!(
1791 "Version {} already exists for table at '{}' with different content",
1792 args.version, args.table_uri
1793 ),
1794 },
1795 ))
1796 }
1797
1798 async fn enforce_create_table_version_cas(
1807 &self,
1808 table_path: &Path,
1809 version: u64,
1810 table_uri: &str,
1811 is_branch: bool,
1812 branch_parent_version: Option<u64>,
1813 ) -> Result<()> {
1814 let latest = self.list_versions_under(table_path, true, Some(1)).await?;
1815 let expected = match latest.first() {
1816 Some(v) => (v.version as u64).checked_add(1).ok_or_else(|| {
1817 lance_core::Error::from(NamespaceError::ConcurrentModification {
1818 message: format!(
1819 "Version overflow computing next version for table at '{}': \
1820 latest version {} cannot advance",
1821 table_uri, v.version
1822 ),
1823 })
1824 })?,
1825 None => {
1826 if is_branch {
1827 match branch_parent_version {
1830 Some(parent_version) => parent_version,
1831 None => version,
1832 }
1833 } else {
1834 1
1835 }
1836 }
1837 };
1838 if version != expected {
1839 let latest_display = latest
1840 .first()
1841 .map(|v| v.version.to_string())
1842 .unwrap_or_else(|| "none".to_string());
1843 return Err(lance_core::Error::from(
1844 NamespaceError::ConcurrentModification {
1845 message: format!(
1846 "Version CAS failed for table at '{}': requested {}, expected {} (latest {})",
1847 table_uri, version, expected, latest_display
1848 ),
1849 },
1850 ));
1851 }
1852 Ok(())
1853 }
1854
1855 async fn materialize_version_manifest_create(
1857 &self,
1858 staging_path: &Path,
1859 final_path: &Path,
1860 staging_manifest_path: &str,
1861 ) -> std::result::Result<(), ObjectStoreError> {
1862 match self
1863 .object_store
1864 .inner
1865 .copy_if_not_exists(staging_path, final_path)
1866 .await
1867 {
1868 Ok(()) => Ok(()),
1869 Err(ObjectStoreError::NotImplemented { .. })
1870 | Err(ObjectStoreError::NotSupported { .. }) => {
1871 let manifest_data = self
1872 .object_store
1873 .inner
1874 .get(staging_path)
1875 .await?
1876 .bytes()
1877 .await
1878 .map_err(|e| ObjectStoreError::Generic {
1879 store: "DirectoryNamespace",
1880 source: Box::new(std::io::Error::other(format!(
1881 "Failed to read staging manifest bytes at '{}': {}",
1882 staging_manifest_path, e
1883 ))),
1884 })?;
1885 self.object_store
1886 .inner
1887 .put_opts(
1888 final_path,
1889 manifest_data.into(),
1890 PutOptions {
1891 mode: PutMode::Create,
1892 ..Default::default()
1893 },
1894 )
1895 .await
1896 .map(|_| ())
1897 }
1898 Err(e) => Err(e),
1899 }
1900 }
1901
1902 async fn list_table_versions_from_storage(
1903 &self,
1904 table_uri: &str,
1905 descending: bool,
1906 limit: Option<i32>,
1907 ) -> Result<Vec<TableVersion>> {
1908 let table_path = self.object_store_path_from_uri(table_uri)?;
1909 self.list_versions_under(&table_path, descending, limit)
1910 .await
1911 }
1912
1913 async fn list_versions_under(
1917 &self,
1918 table_path: &Path,
1919 descending: bool,
1920 limit: Option<i32>,
1921 ) -> Result<Vec<TableVersion>> {
1922 let versions_dir = table_path.clone().join(VERSIONS_DIR);
1923 let manifest_metas: Vec<_> = self
1924 .object_store
1925 .read_dir_all(&versions_dir, None)
1926 .try_collect()
1927 .await
1928 .map_err(|e| {
1929 lance_core::Error::from(NamespaceError::Internal {
1930 message: format!(
1931 "Failed to list manifest files under '{}': {}",
1932 versions_dir, e
1933 ),
1934 })
1935 })?;
1936
1937 let is_v2_naming = manifest_metas
1938 .first()
1939 .is_some_and(|meta| meta.location.filename().is_some_and(|f| f.len() == 29));
1940
1941 let mut table_versions: Vec<TableVersion> = manifest_metas
1942 .into_iter()
1943 .filter_map(|meta| {
1944 let filename = meta.location.filename()?;
1945 let actual_version = Self::manifest_version_from_filename(filename)?;
1946
1947 Some(TableVersion {
1948 version: actual_version as i64,
1949 manifest_path: meta.location.to_string(),
1950 manifest_size: Some(meta.size as i64),
1951 e_tag: meta.e_tag,
1952 timestamp_millis: Some(meta.last_modified.timestamp_millis()),
1953 metadata: None,
1954 })
1955 })
1956 .collect();
1957
1958 let list_is_ordered = self.object_store.list_is_lexically_ordered;
1959
1960 let needs_sort = if list_is_ordered {
1961 if is_v2_naming {
1962 !descending
1963 } else {
1964 descending
1965 }
1966 } else {
1967 true
1968 };
1969
1970 if needs_sort {
1971 if descending {
1972 table_versions.sort_by_key(|v| std::cmp::Reverse(v.version));
1973 } else {
1974 table_versions.sort_by_key(|v| v.version);
1975 }
1976 }
1977
1978 if let Some(limit) = limit {
1979 table_versions.truncate(limit as usize);
1980 }
1981
1982 Ok(table_versions)
1983 }
1984
1985 async fn describe_table_impl(
1989 &self,
1990 request: DescribeTableRequest,
1991 ) -> Result<DescribeTableResponse> {
1992 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
1993 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
1994 let skip_manifest_for_root = self.dir_listing_enabled
1995 && is_root_level
1996 && !self.dir_listing_to_manifest_migration_enabled;
1997 if is_child_table
2005 || !self.dir_listing_enabled
2006 || self.dir_listing_to_manifest_migration_enabled
2007 {
2008 self.ensure_read_manifest().await?;
2009 }
2010 if let Some(manifest_ns) = self.manifest_ns_for_read()
2011 && !skip_manifest_for_root
2012 {
2013 match manifest_ns.describe_table(request.clone()).await {
2014 Ok(mut response) => {
2015 if let Some(ref table_uri) = response.table_uri {
2016 let vend = request.vend_credentials.unwrap_or(true);
2018 let identity = request.identity.as_deref();
2019 response.storage_options = self
2020 .get_storage_options_for_table(table_uri, vend, identity)
2021 .await?;
2022 }
2023 if self.table_version_tracking_enabled {
2025 response.managed_versioning = Some(true);
2026 }
2027 return Ok(response);
2028 }
2029 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
2030 return Err(e);
2033 }
2034 Err(e) if self.dir_listing_enabled && is_root_level => {
2035 if !Self::is_manifest_table_absent_error(&e) {
2039 return Err(Self::classify_storage_error(e));
2040 }
2041 }
2042 Err(e) => return Err(e),
2043 }
2044 }
2045 if is_child_table {
2046 return Err(self.child_namespace_requires_manifest_error());
2047 }
2048
2049 let table_name = Self::table_name_from_id(&request.id)?;
2050 let table_id = Self::format_table_id_from_request(&request.id);
2051 if !self.dir_listing_enabled {
2052 return Err(NamespaceError::TableNotFound { message: table_id }.into());
2053 }
2054
2055 let table_uri = self.table_full_uri(&table_name);
2056
2057 let status = self.check_table_status(&table_name).await?;
2059
2060 if !status.exists {
2061 return Err(NamespaceError::TableNotFound {
2062 message: table_id.clone(),
2063 }
2064 .into());
2065 }
2066
2067 if status.is_deregistered {
2068 return Err(NamespaceError::TableNotFound {
2069 message: format!("Table is deregistered: {}", table_id),
2070 }
2071 .into());
2072 }
2073
2074 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
2075 let should_check_declared =
2076 load_detailed_metadata || request.check_declared.unwrap_or(false);
2077 let vend_credentials = request.vend_credentials.unwrap_or(true);
2079 let identity = request.identity.as_deref();
2080 let is_only_declared = if should_check_declared {
2081 if status.has_reserved_file {
2082 Some(!self.table_has_actual_manifests(&table_name).await?)
2083 } else {
2084 Some(false)
2085 }
2086 } else {
2087 None
2088 };
2089
2090 if !load_detailed_metadata {
2091 let storage_options = self
2092 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2093 .await?;
2094 return Ok(DescribeTableResponse {
2095 table: Some(table_name),
2096 namespace: request.id.as_ref().map(|id| {
2097 if id.len() > 1 {
2098 id[..id.len() - 1].to_vec()
2099 } else {
2100 vec![]
2101 }
2102 }),
2103 location: Some(table_uri.clone()),
2104 table_uri: Some(table_uri),
2105 storage_options,
2106 is_only_declared,
2107 managed_versioning: if self.table_version_tracking_enabled {
2108 Some(true)
2109 } else {
2110 None
2111 },
2112 ..Default::default()
2113 });
2114 }
2115
2116 if is_only_declared == Some(true) {
2117 let storage_options = self
2118 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2119 .await?;
2120 return Ok(DescribeTableResponse {
2121 table: Some(table_name),
2122 namespace: request.id.as_ref().map(|id| {
2123 if id.len() > 1 {
2124 id[..id.len() - 1].to_vec()
2125 } else {
2126 vec![]
2127 }
2128 }),
2129 location: Some(table_uri.clone()),
2130 table_uri: Some(table_uri),
2131 storage_options,
2132 is_only_declared,
2133 managed_versioning: if self.table_version_tracking_enabled {
2134 Some(true)
2135 } else {
2136 None
2137 },
2138 ..Default::default()
2139 });
2140 }
2141
2142 let mut builder = DatasetBuilder::from_uri(&table_uri);
2145 if let Some(opts) = &self.storage_options {
2146 builder = builder.with_storage_options(opts.clone());
2147 }
2148 if let Some(sess) = &self.session {
2149 builder = builder.with_session(sess.clone());
2150 }
2151 match builder.load().await {
2152 Ok(mut dataset) => {
2153 if let Some(requested_version) = request.version {
2155 dataset = dataset
2156 .checkout_version(requested_version as u64)
2157 .await
2158 .map_err(|e| {
2159 let message = format!(
2160 "Version {} not found for table '{}': {}",
2161 requested_version, table_name, e
2162 );
2163 Self::map_open_error(
2164 e,
2165 NamespaceError::TableVersionNotFound { message },
2166 )
2167 })?;
2168 }
2169
2170 let version_info = dataset.version();
2171 let lance_schema = dataset.schema();
2172 let arrow_schema: arrow_schema::Schema = lance_schema.into();
2173 let json_schema = arrow_schema_to_json(&arrow_schema)?;
2174 let storage_options = self
2175 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2176 .await?;
2177
2178 let metadata: std::collections::HashMap<String, String> =
2180 version_info.metadata.into_iter().collect();
2181
2182 Ok(DescribeTableResponse {
2183 table: Some(table_name),
2184 namespace: request.id.as_ref().map(|id| {
2185 if id.len() > 1 {
2186 id[..id.len() - 1].to_vec()
2187 } else {
2188 vec![]
2189 }
2190 }),
2191 version: Some(version_info.version as i64),
2192 location: Some(table_uri.clone()),
2193 table_uri: Some(table_uri),
2194 schema: Some(Box::new(json_schema)),
2195 storage_options,
2196 metadata: Some(metadata),
2197 is_only_declared,
2198 managed_versioning: if self.table_version_tracking_enabled {
2199 Some(true)
2200 } else {
2201 None
2202 },
2203 ..Default::default()
2204 })
2205 }
2206 Err(err) => {
2207 if manifest::ManifestNamespace::is_not_found_load_error(&err)
2208 && is_only_declared == Some(true)
2209 {
2210 let storage_options = self
2211 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2212 .await?;
2213 Ok(DescribeTableResponse {
2214 table: Some(table_name),
2215 namespace: request.id.as_ref().map(|id| {
2216 if id.len() > 1 {
2217 id[..id.len() - 1].to_vec()
2218 } else {
2219 vec![]
2220 }
2221 }),
2222 location: Some(table_uri.clone()),
2223 table_uri: Some(table_uri),
2224 storage_options,
2225 is_only_declared,
2226 managed_versioning: if self.table_version_tracking_enabled {
2227 Some(true)
2228 } else {
2229 None
2230 },
2231 ..Default::default()
2232 })
2233 } else {
2234 Err(NamespaceError::Internal {
2235 message: format!(
2236 "Table directory exists but cannot load dataset {}: {:?}",
2237 table_name, err
2238 ),
2239 }
2240 .into())
2241 }
2242 }
2243 }
2244 }
2245
2246 fn configured_builder(&self, table_uri: &str) -> DatasetBuilder {
2249 let mut builder = DatasetBuilder::from_uri(table_uri);
2250 if let Some(opts) = &self.storage_options {
2251 builder = builder.with_storage_options(opts.clone());
2252 }
2253 if let Some(sess) = &self.session {
2254 builder = builder.with_session(sess.clone());
2255 }
2256 builder
2257 }
2258
2259 async fn load_dataset(
2260 &self,
2261 table_uri: &str,
2262 version: Option<i64>,
2263 operation: &str,
2264 ) -> Result<Dataset> {
2265 if let Some(version) = version
2266 && version < 0
2267 {
2268 return Err(NamespaceError::InvalidInput {
2269 message: format!(
2270 "Table version for {} must be non-negative, got {}",
2271 operation, version
2272 ),
2273 }
2274 .into());
2275 }
2276
2277 let builder = self.configured_builder(table_uri);
2278
2279 let dataset = builder.load().await.map_err(|e| {
2280 let message = format!(
2281 "Failed to open table at '{}' for {}: {}",
2282 table_uri, operation, e
2283 );
2284 Self::map_open_error(e, NamespaceError::TableNotFound { message })
2285 })?;
2286
2287 if let Some(version) = version {
2288 return dataset.checkout_version(version as u64).await.map_err(|e| {
2289 let message = format!(
2290 "Failed to checkout version {} for table at '{}' during {}: {}",
2291 version, table_uri, operation, e
2292 );
2293 Self::map_open_error(e, NamespaceError::TableVersionNotFound { message })
2294 });
2295 }
2296
2297 Ok(dataset)
2298 }
2299
2300 fn parse_index_type(index_type: &str) -> Result<IndexType> {
2301 match index_type.trim().to_ascii_uppercase().as_str() {
2302 "SCALAR" | "BTREE" => Ok(IndexType::BTree),
2303 "BITMAP" => Ok(IndexType::Bitmap),
2304 "LABEL_LIST" | "LABELLIST" => Ok(IndexType::LabelList),
2305 "INVERTED" | "FTS" => Ok(IndexType::Inverted),
2306 "NGRAM" => Ok(IndexType::NGram),
2307 "ZONEMAP" | "ZONE_MAP" => Ok(IndexType::ZoneMap),
2308 "BLOOMFILTER" | "BLOOM_FILTER" => Ok(IndexType::BloomFilter),
2309 "RTREE" | "R_TREE" => Ok(IndexType::RTree),
2310 "VECTOR" | "IVF_PQ" => Ok(IndexType::IvfPq),
2311 "IVF_FLAT" => Ok(IndexType::IvfFlat),
2312 "IVF_SQ" => Ok(IndexType::IvfSq),
2313 "IVF_RQ" => Ok(IndexType::IvfRq),
2314 "IVF_HNSW_FLAT" => Ok(IndexType::IvfHnswFlat),
2315 "IVF_HNSW_SQ" => Ok(IndexType::IvfHnswSq),
2316 "IVF_HNSW_PQ" => Ok(IndexType::IvfHnswPq),
2317 other => Err(NamespaceError::InvalidInput {
2318 message: format!("Unsupported index_type '{}'", other),
2319 }
2320 .into()),
2321 }
2322 }
2323
2324 fn parse_metric_type(distance_type: Option<&str>) -> Result<MetricType> {
2325 let distance_type = distance_type.unwrap_or("l2");
2326 MetricType::try_from(distance_type).map_err(|e| {
2327 lance_core::Error::from(NamespaceError::InvalidInput {
2328 message: format!(
2329 "Unsupported distance_type '{}' for vector index: {}",
2330 distance_type, e
2331 ),
2332 })
2333 })
2334 }
2335
2336 fn build_index_params(request: &CreateTableIndexRequest) -> Result<DirectoryIndexParams> {
2337 let index_type = Self::parse_index_type(&request.index_type)?;
2338 Ok(match index_type {
2339 IndexType::BTree => DirectoryIndexParams::Scalar {
2340 index_type,
2341 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
2342 },
2343 IndexType::Bitmap => DirectoryIndexParams::Scalar {
2344 index_type,
2345 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
2346 },
2347 IndexType::LabelList => DirectoryIndexParams::Scalar {
2348 index_type,
2349 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
2350 },
2351 IndexType::NGram => DirectoryIndexParams::Scalar {
2352 index_type,
2353 params: ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
2354 },
2355 IndexType::ZoneMap => DirectoryIndexParams::Scalar {
2356 index_type,
2357 params: ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
2358 },
2359 IndexType::BloomFilter => DirectoryIndexParams::Scalar {
2360 index_type,
2361 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter),
2362 },
2363 IndexType::RTree => DirectoryIndexParams::Scalar {
2364 index_type,
2365 params: ScalarIndexParams::for_builtin(BuiltinIndexType::RTree),
2366 },
2367 IndexType::Inverted => {
2368 let mut params = InvertedIndexParams::default();
2369 if let Some(with_position) = request.with_position {
2370 params = params.with_position(with_position);
2371 }
2372 if let Some(base_tokenizer) = &request.base_tokenizer {
2373 params = params.base_tokenizer(base_tokenizer.clone());
2374 }
2375 if let Some(language) = &request.language {
2376 params = params.language(language)?;
2377 }
2378 if let Some(max_token_length) = request.max_token_length {
2379 if max_token_length < 0 {
2380 return Err(NamespaceError::InvalidInput {
2381 message: format!(
2382 "FTS max_token_length must be non-negative, got {}",
2383 max_token_length
2384 ),
2385 }
2386 .into());
2387 }
2388 params = params.max_token_length(Some(max_token_length as usize));
2389 }
2390 if let Some(lower_case) = request.lower_case {
2391 params = params.lower_case(lower_case);
2392 }
2393 if let Some(stem) = request.stem {
2394 params = params.stem(stem);
2395 }
2396 if let Some(remove_stop_words) = request.remove_stop_words {
2397 params = params.remove_stop_words(remove_stop_words);
2398 }
2399 if let Some(ascii_folding) = request.ascii_folding {
2400 params = params.ascii_folding(ascii_folding);
2401 }
2402 DirectoryIndexParams::Inverted(params)
2403 }
2404 IndexType::IvfFlat => DirectoryIndexParams::Vector {
2405 index_type,
2406 params: VectorIndexParams::with_ivf_flat_params(
2407 Self::parse_metric_type(request.distance_type.as_deref())?,
2408 IvfBuildParams::default(),
2409 ),
2410 },
2411 IndexType::IvfPq => DirectoryIndexParams::Vector {
2412 index_type,
2413 params: VectorIndexParams::with_ivf_pq_params(
2414 Self::parse_metric_type(request.distance_type.as_deref())?,
2415 IvfBuildParams::default(),
2416 PQBuildParams::default(),
2417 ),
2418 },
2419 IndexType::IvfSq => DirectoryIndexParams::Vector {
2420 index_type,
2421 params: VectorIndexParams::with_ivf_sq_params(
2422 Self::parse_metric_type(request.distance_type.as_deref())?,
2423 IvfBuildParams::default(),
2424 SQBuildParams::default(),
2425 ),
2426 },
2427 IndexType::IvfRq => DirectoryIndexParams::Vector {
2428 index_type,
2429 params: VectorIndexParams::with_ivf_rq_params(
2430 Self::parse_metric_type(request.distance_type.as_deref())?,
2431 IvfBuildParams::default(),
2432 RQBuildParams::default(),
2433 ),
2434 },
2435 IndexType::IvfHnswFlat => DirectoryIndexParams::Vector {
2436 index_type,
2437 params: VectorIndexParams::ivf_hnsw(
2438 Self::parse_metric_type(request.distance_type.as_deref())?,
2439 IvfBuildParams::default(),
2440 HnswBuildParams::default(),
2441 ),
2442 },
2443 IndexType::IvfHnswSq => DirectoryIndexParams::Vector {
2444 index_type,
2445 params: VectorIndexParams::with_ivf_hnsw_sq_params(
2446 Self::parse_metric_type(request.distance_type.as_deref())?,
2447 IvfBuildParams::default(),
2448 HnswBuildParams::default(),
2449 SQBuildParams::default(),
2450 ),
2451 },
2452 IndexType::IvfHnswPq => DirectoryIndexParams::Vector {
2453 index_type,
2454 params: VectorIndexParams::with_ivf_hnsw_pq_params(
2455 Self::parse_metric_type(request.distance_type.as_deref())?,
2456 IvfBuildParams::default(),
2457 HnswBuildParams::default(),
2458 PQBuildParams::default(),
2459 ),
2460 },
2461 other => {
2462 return Err(NamespaceError::InvalidInput {
2463 message: format!("Unsupported index type for namespace API: {}", other),
2464 }
2465 .into());
2466 }
2467 })
2468 }
2469
2470 fn paginate_indices(
2471 indices: &mut Vec<IndexContent>,
2472 page_token: Option<String>,
2473 limit: Option<i32>,
2474 ) -> Option<String> {
2475 indices.sort_by(|a, b| a.index_name.cmp(&b.index_name));
2476
2477 if let Some(start_after) = page_token {
2478 if let Some(index) = indices
2479 .iter()
2480 .position(|index| index.index_name.as_str() > start_after.as_str())
2481 {
2482 indices.drain(0..index);
2483 } else {
2484 indices.clear();
2485 }
2486 }
2487
2488 let mut next_page_token = None;
2489 if let Some(limit) = limit
2490 && limit >= 0
2491 {
2492 let limit = limit as usize;
2493 if limit > 0 && indices.len() > limit {
2494 next_page_token = Some(indices[limit - 1].index_name.clone());
2495 }
2496 indices.truncate(limit);
2497 }
2498 if indices.is_empty() {
2499 None
2500 } else {
2501 next_page_token
2502 }
2503 }
2504
2505 fn transaction_operation_name(transaction: &Transaction) -> String {
2506 match &transaction.operation {
2507 Operation::CreateIndex {
2508 new_indices,
2509 removed_indices,
2510 ..
2511 } if new_indices.is_empty() && !removed_indices.is_empty() => "DropIndex".to_string(),
2512 _ => transaction.operation.to_string(),
2513 }
2514 }
2515
2516 fn transaction_response(
2517 version: u64,
2518 transaction: &Transaction,
2519 alteration: Option<TransactionAlteration>,
2520 ) -> DescribeTransactionResponse {
2521 let mut properties = transaction
2522 .transaction_properties
2523 .as_ref()
2524 .map(|properties| (**properties).clone())
2525 .unwrap_or_default();
2526
2527 let mut effective_status = "SUCCEEDED".to_string();
2530 if let Some(alteration) = alteration {
2531 for key in &alteration.removed_properties {
2532 properties.remove(key);
2533 }
2534 for (key, value) in alteration.properties {
2535 properties.insert(key, value);
2536 }
2537 if let Some(status) = alteration.status {
2538 effective_status = status;
2539 }
2540 }
2541
2542 properties.insert("uuid".to_string(), transaction.uuid.clone());
2543 properties.insert("version".to_string(), version.to_string());
2544 properties.insert(
2545 "read_version".to_string(),
2546 transaction.read_version.to_string(),
2547 );
2548 properties.insert(
2549 "operation".to_string(),
2550 Self::transaction_operation_name(transaction),
2551 );
2552 if let Some(tag) = &transaction.tag {
2553 properties.insert("tag".to_string(), tag.clone());
2554 }
2555
2556 DescribeTransactionResponse {
2557 status: effective_status,
2558 properties: Some(properties),
2559 ..Default::default()
2560 }
2561 }
2562
2563 fn describe_table_index_stats_response(
2564 stats: &serde_json::Value,
2565 ) -> DescribeTableIndexStatsResponse {
2566 let get_i64 = |key: &str| {
2567 stats.get(key).and_then(|value| {
2568 value
2569 .as_i64()
2570 .or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
2571 })
2572 };
2573
2574 DescribeTableIndexStatsResponse {
2575 distance_type: stats
2576 .get("distance_type")
2577 .and_then(|value| value.as_str())
2578 .map(str::to_string),
2579 index_type: stats
2580 .get("index_type")
2581 .and_then(|value| value.as_str())
2582 .map(str::to_string),
2583 num_indexed_rows: get_i64("num_indexed_rows"),
2584 num_unindexed_rows: get_i64("num_unindexed_rows"),
2585 num_indices: get_i64("num_indices").and_then(|value| i32::try_from(value).ok()),
2586 ..Default::default()
2587 }
2588 }
2589
2590 async fn find_transaction(&self, dataset: &Dataset, id: &str) -> Result<(u64, Transaction)> {
2595 if let Ok(version) = id.parse::<u64>() {
2596 let transaction = dataset
2597 .read_transaction_by_version(version)
2598 .await
2599 .map_err(|e| {
2600 lance_core::Error::from(NamespaceError::TransactionNotFound {
2601 message: format!(
2602 "Failed to read transaction for version {}: {}",
2603 version, e
2604 ),
2605 })
2606 })?
2607 .ok_or_else(|| {
2608 lance_core::Error::from(NamespaceError::TransactionNotFound {
2609 message: format!("version {}", version),
2610 })
2611 })?;
2612 return Ok((version, transaction));
2613 }
2614
2615 let versions = dataset.versions().await.map_err(|e| {
2616 lance_core::Error::from(NamespaceError::Internal {
2617 message: format!(
2618 "Failed to list table versions while resolving transaction '{}': {}",
2619 id, e
2620 ),
2621 })
2622 })?;
2623
2624 for version in versions.into_iter().rev() {
2625 if let Some(transaction) = dataset
2626 .read_transaction_by_version(version.version)
2627 .await
2628 .map_err(|e| {
2629 lance_core::Error::from(NamespaceError::Internal {
2630 message: format!(
2631 "Failed to read transaction for version {} while resolving '{}': {}",
2632 version.version, id, e
2633 ),
2634 })
2635 })?
2636 && transaction.uuid == id
2637 {
2638 return Ok((version.version, transaction));
2639 }
2640 }
2641
2642 Err(NamespaceError::TransactionNotFound {
2643 message: id.to_string(),
2644 }
2645 .into())
2646 }
2647
2648 const TRANSACTION_ALTERATIONS_DIR: &'static str = "_alter_transactions";
2652
2653 fn transaction_alteration_path(&self, table_uri: &str, txn_uuid: &str) -> Result<Path> {
2654 let table_path = self.object_store_path_from_uri(table_uri)?;
2655 Ok(table_path
2656 .join(Self::TRANSACTION_ALTERATIONS_DIR)
2657 .join(format!("{}.json", txn_uuid).as_str()))
2658 }
2659
2660 async fn load_transaction_alteration(
2661 &self,
2662 table_uri: &str,
2663 txn_uuid: &str,
2664 ) -> Result<Option<TransactionAlteration>> {
2665 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2666 match self.object_store.inner.get(&path).await {
2667 Ok(get_result) => {
2668 let bytes = get_result.bytes().await.map_err(|e| {
2669 lance_core::Error::from(NamespaceError::Internal {
2670 message: format!(
2671 "Failed to read alter_transaction sidecar for '{}': {}",
2672 txn_uuid, e
2673 ),
2674 })
2675 })?;
2676 let alteration = TransactionAlteration::from_json_slice(&bytes).map_err(|e| {
2677 lance_core::Error::from(NamespaceError::Internal {
2678 message: format!(
2679 "Failed to parse alter_transaction sidecar for '{}': {}",
2680 txn_uuid, e
2681 ),
2682 })
2683 })?;
2684 Ok(Some(alteration))
2685 }
2686 Err(ObjectStoreError::NotFound { .. }) => Ok(None),
2687 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2688 message: format!(
2689 "Failed to load alter_transaction sidecar for '{}': {}",
2690 txn_uuid, e
2691 ),
2692 })),
2693 }
2694 }
2695
2696 async fn save_transaction_alteration(
2697 &self,
2698 table_uri: &str,
2699 txn_uuid: &str,
2700 alteration: &TransactionAlteration,
2701 ) -> Result<()> {
2702 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2703 let bytes = alteration.to_json_bytes().map_err(|e| {
2704 lance_core::Error::from(NamespaceError::Internal {
2705 message: format!(
2706 "Failed to serialize alter_transaction sidecar for '{}': {}",
2707 txn_uuid, e
2708 ),
2709 })
2710 })?;
2711 self.object_store
2712 .inner
2713 .put(&path, bytes.into())
2714 .await
2715 .map_err(|e| {
2716 lance_core::Error::from(NamespaceError::Internal {
2717 message: format!(
2718 "Failed to persist alter_transaction sidecar for '{}': {}",
2719 txn_uuid, e
2720 ),
2721 })
2722 })?;
2723 Ok(())
2724 }
2725
2726 fn table_full_uri(&self, table_name: &str) -> String {
2727 format!("{}/{}.lance", self.root, table_name)
2728 }
2729
2730 fn table_path(&self, table_name: &str) -> Path {
2732 self.base_path
2733 .clone()
2734 .join(format!("{}.lance", table_name).as_str())
2735 }
2736
2737 fn table_reserved_file_path(&self, table_name: &str) -> Path {
2739 self.base_path
2740 .clone()
2741 .join(format!("{}.lance", table_name).as_str())
2742 .join(".lance-reserved")
2743 }
2744
2745 fn table_deregistered_file_path(&self, table_name: &str) -> Path {
2747 self.base_path
2748 .clone()
2749 .join(format!("{}.lance", table_name).as_str())
2750 .join(".lance-deregistered")
2751 }
2752
2753 pub(crate) async fn check_table_status(&self, table_name: &str) -> Result<TableStatus> {
2759 let table_path = self.table_path(table_name);
2760 match self.object_store.read_dir(table_path).await {
2761 Ok(entries) => {
2762 let exists = !entries.is_empty();
2763 let is_deregistered = entries.iter().any(|e| e.ends_with(".lance-deregistered"));
2764 let has_reserved_file = entries.iter().any(|e| e.ends_with(".lance-reserved"));
2765 Ok(TableStatus {
2766 exists,
2767 is_deregistered,
2768 has_reserved_file,
2769 })
2770 }
2771 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(TableStatus {
2774 exists: false,
2775 is_deregistered: false,
2776 has_reserved_file: false,
2777 }),
2778 Err(e) => Err(Self::classify_storage_error(e)),
2782 }
2783 }
2784
2785 fn classify_storage_error(err: Error) -> Error {
2790 if matches!(&err, Error::Namespace { .. }) {
2791 return err;
2792 }
2793 let detail = err.to_string();
2794 if let Error::IO { source, .. } = &err
2795 && let Some(os_err) = source.downcast_ref::<ObjectStoreError>()
2796 {
2797 if is_throttle_error(os_err) {
2798 return NamespaceError::Throttling {
2799 message: format!(
2800 "Storage request was throttled while resolving table: {detail}"
2801 ),
2802 }
2803 .into();
2804 }
2805 if Self::is_service_unavailable_error(os_err) {
2806 return NamespaceError::ServiceUnavailable {
2807 message: format!("Storage service unavailable while resolving table: {detail}"),
2808 }
2809 .into();
2810 }
2811 }
2812 NamespaceError::Internal {
2813 message: format!("Storage error while resolving table: {detail}"),
2814 }
2815 .into()
2816 }
2817
2818 fn is_service_unavailable_error(err: &ObjectStoreError) -> bool {
2822 if let ObjectStoreError::Generic { source, .. } = err {
2823 let message = source.to_string().to_ascii_lowercase();
2824 message.contains("503 service unavailable")
2825 || message.contains("502 bad gateway")
2826 || message.contains("504 gateway timeout")
2827 } else {
2828 false
2829 }
2830 }
2831
2832 fn is_manifest_table_absent_error(err: &Error) -> bool {
2836 if manifest::ManifestNamespace::is_not_found_load_error(err) {
2837 return true;
2838 }
2839 if let Error::Namespace { source, .. } = err
2840 && let Some(ns_err) = source.downcast_ref::<NamespaceError>()
2841 {
2842 return matches!(ns_err, NamespaceError::TableNotFound { .. });
2843 }
2844 false
2845 }
2846
2847 fn map_open_error(err: Error, not_found: NamespaceError) -> Error {
2851 if matches!(&err, Error::IO { .. })
2852 && !manifest::ManifestNamespace::is_not_found_load_error(&err)
2853 {
2854 return Self::classify_storage_error(err);
2855 }
2856 not_found.into()
2857 }
2858
2859 async fn get_storage_options_for_table(
2879 &self,
2880 table_uri: &str,
2881 vend_credentials: bool,
2882 identity: Option<&Identity>,
2883 ) -> Result<Option<HashMap<String, String>>> {
2884 if vend_credentials && let Some(ref vendor) = self.credential_vendor {
2885 let vended = vendor.vend_credentials(table_uri, identity).await?;
2886 return Ok(Some(vended.storage_options));
2887 }
2888 if self.vend_input_storage_options {
2891 let mut options = self.storage_options.clone().unwrap_or_default();
2892 if let Some(refresh_interval_millis) =
2894 self.vend_input_storage_options_refresh_interval_millis
2895 {
2896 let now_millis = std::time::SystemTime::now()
2897 .duration_since(std::time::UNIX_EPOCH)
2898 .unwrap()
2899 .as_millis() as u64;
2900 let expires_at_millis = now_millis + refresh_interval_millis;
2901 options.insert(
2902 "expires_at_millis".to_string(),
2903 expires_at_millis.to_string(),
2904 );
2905 }
2906 return Ok(Some(options));
2907 }
2908 Ok(None)
2911 }
2912
2913 pub async fn migrate(&self) -> Result<usize> {
2967 let Some(manifest_ns) = self.manifest_ns_for_write().await? else {
2969 return Ok(0); };
2971
2972 let manifest_locations = manifest_ns.list_manifest_table_locations().await?;
2974
2975 let dir_tables = self
2978 .filter_declared_tables(self.list_directory_tables().await?, false)
2979 .await?;
2980
2981 let mut migrated_count = 0;
2986 for table_name in dir_tables {
2987 let dir_name = format!("{}.lance", table_name);
2989 if !manifest_locations.contains(&dir_name) {
2990 manifest_ns.register_table(&table_name, dir_name).await?;
2991 migrated_count += 1;
2992 }
2993 }
2994
2995 Ok(migrated_count)
2996 }
2997
2998 async fn delete_physical_version_files(
3006 &self,
3007 table_entries: &[TableDeleteEntry],
3008 branch: Option<&str>,
3009 ) -> Result<i64> {
3010 let mut deleted_count = 0i64;
3011 for te in table_entries {
3012 let table_uri = self.resolve_table_location(&te.table_id).await?;
3013 let table_uri = match branch {
3014 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3015 None => table_uri,
3016 };
3017 let table_path = self.object_store_path_from_uri(&table_uri)?;
3018 let versions_dir_path = table_path.clone().join(VERSIONS_DIR);
3019
3020 let manifest_metas: Vec<_> = self
3022 .object_store
3023 .read_dir_all(&versions_dir_path, None)
3024 .try_collect()
3025 .await
3026 .map_err(|e| {
3027 lance_core::Error::from(NamespaceError::Internal {
3028 message: format!(
3029 "Failed to list manifest files for table at '{}': {}",
3030 table_uri, e
3031 ),
3032 })
3033 })?;
3034 let location_by_version: HashMap<u64, Path> = manifest_metas
3035 .into_iter()
3036 .filter_map(|meta| {
3037 let version = Self::manifest_version_from_filename(meta.location.filename()?)?;
3038 Some((version, meta.location))
3039 })
3040 .collect();
3041
3042 for (&v, version_path) in &location_by_version {
3043 let vi = v as i64;
3044 if !te.ranges.iter().any(|&(s, e)| vi >= s && (e < 0 || vi < e)) {
3045 continue;
3046 }
3047 match self.object_store.inner.delete(version_path).await {
3048 Ok(_) => {
3049 deleted_count += 1;
3050 }
3051 Err(object_store::Error::NotFound { .. }) => {}
3052 Err(e) => {
3053 return Err(NamespaceError::Internal {
3054 message: format!(
3055 "Failed to delete version {} for table at '{}': {}",
3056 v, table_uri, e
3057 ),
3058 }
3059 .into());
3060 }
3061 }
3062 }
3063 }
3064 Ok(deleted_count)
3065 }
3066
3067 #[allow(clippy::too_many_arguments)]
3072 fn apply_query_params_to_scanner(
3073 scanner: &mut Scanner,
3074 filter: Option<&str>,
3075 columns: Option<&QueryTableRequestColumns>,
3076 vector_column: Option<&str>,
3077 vector: &QueryTableRequestVector,
3078 k: i32,
3079 offset: Option<i32>,
3080 prefilter: Option<bool>,
3081 bypass_vector_index: Option<bool>,
3082 nprobes: Option<i32>,
3083 ef: Option<i32>,
3084 refine_factor: Option<i32>,
3085 distance_type: Option<&str>,
3086 fast_search_flag: Option<bool>,
3087 with_row_id: Option<bool>,
3088 lower_bound: Option<f32>,
3089 upper_bound: Option<f32>,
3090 operation: &str,
3091 ) -> Result<()> {
3092 if let Some(pf) = prefilter {
3094 scanner.prefilter(pf);
3095 }
3096
3097 if let Some(filter) = filter {
3098 scanner.filter(filter).map_err(|e| {
3099 Error::invalid_input_source(
3100 format!("Invalid filter expression for {}: {}", operation, e).into(),
3101 )
3102 })?;
3103 }
3104
3105 if let Some(cols) = columns {
3106 if let Some(ref names) = cols.column_names {
3107 scanner.project(names.as_slice()).map_err(|e| {
3108 Error::invalid_input_source(
3109 format!("Invalid column projection for {}: {}", operation, e).into(),
3110 )
3111 })?;
3112 } else if let Some(ref aliases) = cols.column_aliases {
3113 let pairs: Vec<(&str, &str)> = aliases
3115 .iter()
3116 .map(|(alias, src)| (alias.as_str(), src.as_str()))
3117 .collect();
3118 scanner.project_with_transform(&pairs).map_err(|e| {
3119 Error::invalid_input_source(
3120 format!("Invalid column aliases for {}: {}", operation, e).into(),
3121 )
3122 })?;
3123 }
3124 }
3125
3126 let query_vec: Option<Vec<f32>> = vector
3128 .single_vector
3129 .as_ref()
3130 .filter(|v| !v.is_empty())
3131 .cloned()
3132 .or_else(|| {
3133 vector
3134 .multi_vector
3135 .as_ref()
3136 .and_then(|mv| mv.first())
3137 .filter(|v| !v.is_empty())
3138 .cloned()
3139 });
3140
3141 if let Some(q_vec) = query_vec {
3142 let col = vector_column.unwrap_or("vector");
3143 let q = Arc::new(Float32Array::from(q_vec));
3144 scanner
3145 .nearest(col, q.as_ref(), k.max(1) as usize)
3146 .map_err(|e| {
3147 Error::invalid_input_source(
3148 format!("Invalid vector query for {}: {}", operation, e).into(),
3149 )
3150 })?;
3151
3152 if let Some(n) = nprobes {
3154 scanner.nprobes(n.max(1) as usize);
3155 }
3156 if let Some(e) = ef {
3157 scanner.ef(e.max(1) as usize);
3158 }
3159 if let Some(rf) = refine_factor {
3160 scanner.refine(rf.max(0) as u32);
3161 }
3162 if let Some(true) = bypass_vector_index {
3164 scanner.use_index(false);
3165 }
3166 if let Some(true) = fast_search_flag {
3167 scanner.fast_search();
3168 }
3169 if lower_bound.is_some() || upper_bound.is_some() {
3170 scanner.distance_range(lower_bound, upper_bound);
3171 }
3172 if let Some(dt) = distance_type {
3173 let metric = Self::parse_metric_type(Some(dt))?;
3174 scanner.distance_metric(metric);
3175 }
3176 if let Some(off) = offset.filter(|&o| o > 0) {
3178 scanner.limit(None, Some(off as i64)).map_err(|e| {
3179 Error::invalid_input_source(
3180 format!("Invalid offset for {}: {}", operation, e).into(),
3181 )
3182 })?;
3183 }
3184 } else {
3185 let limit = if k > 0 { Some(k as i64) } else { None };
3187 scanner
3188 .limit(limit, offset.map(|o| o as i64))
3189 .map_err(|e| {
3190 Error::invalid_input_source(
3191 format!("Invalid limit/offset for {}: {}", operation, e).into(),
3192 )
3193 })?;
3194 }
3195
3196 if let Some(true) = with_row_id {
3197 scanner.with_row_id();
3198 }
3199
3200 Ok(())
3201 }
3202
3203 pub fn retrieve_ops_metrics(&self) -> HashMap<String, u64> {
3210 self.ops_metrics
3211 .as_ref()
3212 .map(|m| m.retrieve())
3213 .unwrap_or_default()
3214 }
3215
3216 pub fn reset_ops_metrics(&self) {
3220 if let Some(ref metrics) = self.ops_metrics {
3221 metrics.reset();
3222 }
3223 }
3224
3225 fn record_op(&self, operation: &str) {
3227 if let Some(ref metrics) = self.ops_metrics {
3228 metrics.increment(operation);
3229 }
3230 }
3231}
3232
3233#[async_trait]
3234impl LanceNamespace for DirectoryNamespace {
3235 async fn list_namespaces(
3236 &self,
3237 request: ListNamespacesRequest,
3238 ) -> Result<ListNamespacesResponse> {
3239 self.record_op("list_namespaces");
3240 self.ensure_read_manifest().await?;
3241 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3242 return manifest_ns.list_namespaces(request).await;
3243 }
3244
3245 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3246 return Err(self.child_namespace_requires_manifest_error());
3247 }
3248 Self::validate_root_namespace_id(&request.id)?;
3249 Ok(ListNamespacesResponse::new(vec![]))
3250 }
3251
3252 async fn describe_namespace(
3253 &self,
3254 request: DescribeNamespaceRequest,
3255 ) -> Result<DescribeNamespaceResponse> {
3256 self.record_op("describe_namespace");
3257 self.ensure_read_manifest().await?;
3258 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3259 return manifest_ns.describe_namespace(request).await;
3260 }
3261
3262 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3263 return Err(self.child_namespace_requires_manifest_error());
3264 }
3265 Self::validate_root_namespace_id(&request.id)?;
3266 #[allow(clippy::needless_update)]
3267 Ok(DescribeNamespaceResponse {
3268 properties: Some(HashMap::new()),
3269 ..Default::default()
3270 })
3271 }
3272
3273 async fn create_namespace(
3274 &self,
3275 request: CreateNamespaceRequest,
3276 ) -> Result<CreateNamespaceResponse> {
3277 self.record_op("create_namespace");
3278 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3279 return manifest_ns.create_namespace(request).await;
3280 }
3281
3282 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3283 return Err(NamespaceError::NamespaceAlreadyExists {
3284 message: "root namespace".to_string(),
3285 }
3286 .into());
3287 }
3288
3289 Err(NamespaceError::Unsupported {
3290 message: "Child namespaces are only supported when manifest mode is enabled"
3291 .to_string(),
3292 }
3293 .into())
3294 }
3295
3296 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
3297 self.record_op("drop_namespace");
3298 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3299 return manifest_ns.drop_namespace(request).await;
3300 }
3301
3302 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3303 return Err(NamespaceError::InvalidInput {
3304 message: "Root namespace cannot be dropped".to_string(),
3305 }
3306 .into());
3307 }
3308
3309 Err(NamespaceError::Unsupported {
3310 message: "Child namespaces are only supported when manifest mode is enabled"
3311 .to_string(),
3312 }
3313 .into())
3314 }
3315
3316 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
3317 self.record_op("namespace_exists");
3318 self.ensure_read_manifest().await?;
3319 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3320 return manifest_ns.namespace_exists(request).await;
3321 }
3322
3323 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3324 return Ok(());
3325 }
3326
3327 Err(self.child_namespace_requires_manifest_error())
3328 }
3329
3330 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
3331 self.record_op("list_tables");
3332 let namespace_id = request.id.as_ref().ok_or_else(|| {
3334 lance_core::Error::from(NamespaceError::InvalidInput {
3335 message: "Namespace ID is required".to_string(),
3336 })
3337 })?;
3338
3339 if !namespace_id.is_empty()
3345 || !self.dir_listing_enabled
3346 || self.dir_listing_to_manifest_migration_enabled
3347 {
3348 self.ensure_read_manifest().await?;
3349 }
3350
3351 if !namespace_id.is_empty() {
3353 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3354 return manifest_ns.list_tables(request).await;
3355 }
3356 return Err(self.child_namespace_requires_manifest_error());
3357 }
3358
3359 if let Some(manifest_ns) = self.manifest_ns_for_read()
3361 && !self.dir_listing_enabled
3362 {
3363 return manifest_ns.list_tables(request).await;
3364 }
3365 if !self.dir_listing_enabled {
3366 return Ok(ListTablesResponse::new(vec![]));
3367 }
3368
3369 let mut tables = if self.manifest_ns_for_read().is_some()
3372 && self.dir_listing_enabled
3373 && self.dir_listing_to_manifest_migration_enabled
3374 {
3375 let manifest_locations = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3377 manifest_ns.list_manifest_table_locations().await?
3378 } else {
3379 std::collections::HashSet::new()
3380 };
3381
3382 let mut manifest_request = request.clone();
3384 manifest_request.limit = None;
3385 manifest_request.page_token = None;
3386 let manifest_tables = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3387 let manifest_response = manifest_ns.list_tables(manifest_request).await?;
3388 manifest_response.tables
3389 } else {
3390 vec![]
3391 };
3392
3393 let mut all_tables: Vec<String> = manifest_tables;
3396 let dir_tables = self.list_directory_tables().await?;
3397 for table_name in dir_tables {
3398 let full_location = format!("{}/{}.lance", self.root, table_name);
3401 let relative_location = format!("{}.lance", table_name);
3402 if !manifest_locations.contains(&full_location)
3403 && !manifest_locations.contains(&relative_location)
3404 {
3405 all_tables.push(table_name);
3406 }
3407 }
3408
3409 all_tables
3410 } else {
3411 self.list_directory_tables().await?
3412 };
3413
3414 tables = self
3415 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
3416 .await?;
3417
3418 let next_page_token =
3420 Self::apply_pagination(&mut tables, request.page_token, request.limit);
3421 let mut response = ListTablesResponse::new(tables);
3422 response.page_token = next_page_token;
3423 Ok(response)
3424 }
3425
3426 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
3427 self.record_op("describe_table");
3428 self.describe_table_impl(request).await
3429 }
3430
3431 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
3432 self.record_op("table_exists");
3433 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
3434 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
3435 let skip_manifest_for_root = self.dir_listing_enabled
3436 && is_root_level
3437 && !self.dir_listing_to_manifest_migration_enabled;
3438 if is_child_table
3441 || !self.dir_listing_enabled
3442 || self.dir_listing_to_manifest_migration_enabled
3443 {
3444 self.ensure_read_manifest().await?;
3445 }
3446 if let Some(manifest_ns) = self.manifest_ns_for_read()
3447 && !skip_manifest_for_root
3448 {
3449 match manifest_ns.table_exists(request.clone()).await {
3450 Ok(()) => return Ok(()),
3451 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
3452 return Err(e);
3455 }
3456 Err(e) if self.dir_listing_enabled && is_root_level => {
3457 if !Self::is_manifest_table_absent_error(&e) {
3461 return Err(Self::classify_storage_error(e));
3462 }
3463 }
3464 Err(e) => return Err(e),
3465 }
3466 }
3467 if is_child_table {
3468 return Err(self.child_namespace_requires_manifest_error());
3469 }
3470
3471 let table_name = Self::table_name_from_id(&request.id)?;
3472 let table_id = Self::format_table_id_from_request(&request.id);
3473 if !self.dir_listing_enabled {
3474 return Err(NamespaceError::TableNotFound { message: table_id }.into());
3475 }
3476
3477 let status = self.check_table_status(&table_name).await?;
3479
3480 if !status.exists {
3481 return Err(NamespaceError::TableNotFound {
3482 message: table_id.clone(),
3483 }
3484 .into());
3485 }
3486
3487 if status.is_deregistered {
3488 return Err(NamespaceError::TableNotFound {
3489 message: format!("Table is deregistered: {}", table_id),
3490 }
3491 .into());
3492 }
3493
3494 Ok(())
3495 }
3496
3497 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3498 self.record_op("drop_table");
3499 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3500 return manifest_ns.drop_table(request).await;
3501 }
3502
3503 let table_name = Self::table_name_from_id(&request.id)?;
3504 let table_uri = self.table_full_uri(&table_name);
3505 let table_path = self.table_path(&table_name);
3506
3507 self.object_store
3508 .remove_dir_all(table_path)
3509 .await
3510 .map_err(|e| {
3511 lance_core::Error::from(NamespaceError::Internal {
3512 message: format!("Failed to drop table {}: {:?}", table_name, e),
3513 })
3514 })?;
3515
3516 Ok(DropTableResponse {
3517 id: request.id,
3518 location: Some(table_uri),
3519 ..Default::default()
3520 })
3521 }
3522
3523 async fn create_table(
3524 &self,
3525 request: CreateTableRequest,
3526 request_data: Bytes,
3527 ) -> Result<CreateTableResponse> {
3528 self.record_op("create_table");
3529 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3530 return manifest_ns.create_table(request, request_data).await;
3531 }
3532
3533 Self::validate_dir_only_properties(request.properties.as_ref(), "create_table")?;
3534
3535 let table_name = Self::table_name_from_id(&request.id)?;
3536 let table_uri = self.table_full_uri(&table_name);
3537 let status = self.check_table_status(&table_name).await?;
3538 let (reader, _num_rows) =
3539 Self::ipc_reader_from_request_data(&request_data, "create_table")?;
3540
3541 if status.exists && self.table_has_actual_manifests(&table_name).await? {
3542 return Err(NamespaceError::TableAlreadyExists {
3543 message: table_name,
3544 }
3545 .into());
3546 }
3547
3548 let write_result = self
3549 .write_reader_to_table(
3550 &table_uri,
3551 reader,
3552 WriteMode::Create,
3553 request.storage_options.clone(),
3554 )
3555 .await;
3556 if let Err(err) = write_result {
3557 if self.table_uri_has_actual_manifests(&table_uri).await? {
3558 return Err(NamespaceError::TableAlreadyExists {
3559 message: table_name,
3560 }
3561 .into());
3562 }
3563 return Err(err);
3564 }
3565 Ok(CreateTableResponse {
3566 version: Some(1),
3567 location: Some(table_uri),
3568 storage_options: self.storage_options.clone(),
3569 properties: request.properties,
3570 ..Default::default()
3571 })
3572 }
3573
3574 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3575 self.record_op("declare_table");
3576 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3577 let mut response = manifest_ns.declare_table(request.clone()).await?;
3578 if let Some(ref location) = response.location {
3579 let vend = request.vend_credentials.unwrap_or(true);
3581 let identity = request.identity.as_deref();
3582 response.storage_options = self
3583 .get_storage_options_for_table(location, vend, identity)
3584 .await?;
3585 }
3586 if self.table_version_tracking_enabled {
3588 response.managed_versioning = Some(true);
3589 }
3590 return Ok(response);
3591 }
3592
3593 Self::validate_dir_only_properties(request.properties.as_ref(), "declare_table")?;
3594
3595 let table_name = Self::table_name_from_id(&request.id)?;
3596 let table_uri = self.table_full_uri(&table_name);
3597
3598 if let Some(location) = &request.location {
3600 let location = location.trim_end_matches('/');
3601 if location != table_uri {
3602 return Err(NamespaceError::InvalidInput {
3603 message: format!(
3604 "Cannot declare table {} at location {}, must be at location {}",
3605 table_name, location, table_uri
3606 ),
3607 }
3608 .into());
3609 }
3610 }
3611
3612 let status = self.check_table_status(&table_name).await?;
3616 if status.exists && !status.has_reserved_file {
3617 return Err(NamespaceError::TableAlreadyExists {
3619 message: table_name.to_string(),
3620 }
3621 .into());
3622 }
3623
3624 let reserved_file_path = self.table_reserved_file_path(&table_name);
3628
3629 put_marker_file_atomic(
3630 &self.object_store,
3631 &reserved_file_path,
3632 &format!("table {}", table_name),
3633 )
3634 .await
3635 .map_err(|e| match e {
3636 MarkerFileError::AlreadyExists { .. } => {
3637 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3638 message: table_name.to_string(),
3639 })
3640 }
3641 MarkerFileError::Other { message } => {
3642 lance_core::Error::from(NamespaceError::Internal { message })
3643 }
3644 })?;
3645
3646 let vend_credentials = request.vend_credentials.unwrap_or(true);
3648 let identity = request.identity.as_deref();
3649 let storage_options = self
3650 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
3651 .await?;
3652
3653 Ok(DeclareTableResponse {
3654 location: Some(table_uri),
3655 storage_options,
3656 properties: request.properties,
3657 managed_versioning: if self.table_version_tracking_enabled {
3658 Some(true)
3659 } else {
3660 None
3661 },
3662 ..Default::default()
3663 })
3664 }
3665
3666 async fn register_table(
3667 &self,
3668 request: lance_namespace::models::RegisterTableRequest,
3669 ) -> Result<lance_namespace::models::RegisterTableResponse> {
3670 self.record_op("register_table");
3671 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3673 return LanceNamespace::register_table(manifest_ns.as_ref(), request).await;
3674 }
3675
3676 Err(NamespaceError::Unsupported {
3678 message: "register_table is only supported when manifest mode is enabled".to_string(),
3679 }
3680 .into())
3681 }
3682
3683 async fn deregister_table(
3684 &self,
3685 request: lance_namespace::models::DeregisterTableRequest,
3686 ) -> Result<lance_namespace::models::DeregisterTableResponse> {
3687 self.record_op("deregister_table");
3688 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3690 return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await;
3691 }
3692
3693 let table_name = Self::table_name_from_id(&request.id)?;
3695 let table_uri = self.table_full_uri(&table_name);
3696
3697 let status = self.check_table_status(&table_name).await?;
3700
3701 if !status.exists {
3702 return Err(NamespaceError::TableNotFound {
3703 message: table_name.to_string(),
3704 }
3705 .into());
3706 }
3707
3708 if status.is_deregistered {
3709 return Err(NamespaceError::TableNotFound {
3710 message: format!("Table is already deregistered: {}", table_name),
3711 }
3712 .into());
3713 }
3714
3715 let deregistered_path = self.table_deregistered_file_path(&table_name);
3721 put_marker_file_atomic(
3722 &self.object_store,
3723 &deregistered_path,
3724 &format!("deregistration marker for table {}", table_name),
3725 )
3726 .await
3727 .map_err(|e| match e {
3728 MarkerFileError::AlreadyExists { .. } => {
3729 lance_core::Error::from(NamespaceError::InvalidTableState {
3730 message: format!("Table is already deregistered: {}", table_name),
3731 })
3732 }
3733 MarkerFileError::Other { message } => {
3734 lance_core::Error::from(NamespaceError::Internal { message })
3735 }
3736 })?;
3737
3738 Ok(lance_namespace::models::DeregisterTableResponse {
3739 id: request.id,
3740 location: Some(table_uri),
3741 ..Default::default()
3742 })
3743 }
3744
3745 async fn alter_table_add_columns(
3746 &self,
3747 request: AlterTableAddColumnsRequest,
3748 ) -> Result<AlterTableAddColumnsResponse> {
3749 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3750 return manifest_ns.alter_table_add_columns(request).await;
3751 }
3752
3753 let table_name = Self::table_name_from_id(&request.id)?;
3755 let table_uri = self.table_full_uri(&table_name);
3756
3757 let status = self.check_table_status(&table_name).await?;
3759 if !status.exists {
3760 return Err(NamespaceError::TableNotFound {
3761 message: table_name,
3762 }
3763 .into());
3764 }
3765 if status.is_deregistered {
3766 return Err(NamespaceError::TableNotFound {
3767 message: format!("Table is deregistered: {}", table_name),
3768 }
3769 .into());
3770 }
3771
3772 let mut dataset = self
3773 .configured_builder(&table_uri)
3774 .load()
3775 .await
3776 .map_err(|e| {
3777 Error::io_source(box_error(std::io::Error::other(format!(
3778 "Failed to open dataset: {}",
3779 e
3780 ))))
3781 })?;
3782
3783 let sql_expressions = build_sql_expressions(&request.new_columns)?;
3784
3785 dataset
3786 .add_columns(
3787 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3788 None,
3789 None,
3790 )
3791 .await
3792 .map_err(|e| {
3793 Error::io_source(box_error(std::io::Error::other(format!(
3794 "Failed to add columns: {}",
3795 e
3796 ))))
3797 })?;
3798
3799 let version = dataset.version().version as i64;
3800 Ok(AlterTableAddColumnsResponse::new(version))
3801 }
3802
3803 async fn alter_table_alter_columns(
3804 &self,
3805 request: AlterTableAlterColumnsRequest,
3806 ) -> Result<AlterTableAlterColumnsResponse> {
3807 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3808 return manifest_ns.alter_table_alter_columns(request).await;
3809 }
3810
3811 let table_name = Self::table_name_from_id(&request.id)?;
3812 let table_uri = self.table_full_uri(&table_name);
3813
3814 let status = self.check_table_status(&table_name).await?;
3816 if !status.exists {
3817 return Err(NamespaceError::TableNotFound {
3818 message: table_name,
3819 }
3820 .into());
3821 }
3822 if status.is_deregistered {
3823 return Err(NamespaceError::TableNotFound {
3824 message: format!("Table is deregistered: {}", table_name),
3825 }
3826 .into());
3827 }
3828
3829 let mut dataset = self
3830 .configured_builder(&table_uri)
3831 .load()
3832 .await
3833 .map_err(|e| {
3834 Error::io_source(box_error(std::io::Error::other(format!(
3835 "Failed to open dataset: {}",
3836 e
3837 ))))
3838 })?;
3839
3840 let alterations = build_column_alterations(&request.alterations)?;
3841
3842 dataset.alter_columns(&alterations).await.map_err(|e| {
3843 Error::io_source(box_error(std::io::Error::other(format!(
3844 "Failed to alter columns: {}",
3845 e
3846 ))))
3847 })?;
3848
3849 let version = dataset.version().version as i64;
3850 Ok(AlterTableAlterColumnsResponse::new(version))
3851 }
3852
3853 async fn alter_table_drop_columns(
3854 &self,
3855 request: AlterTableDropColumnsRequest,
3856 ) -> Result<AlterTableDropColumnsResponse> {
3857 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3858 return manifest_ns.alter_table_drop_columns(request).await;
3859 }
3860
3861 let table_name = Self::table_name_from_id(&request.id)?;
3862 let table_uri = self.table_full_uri(&table_name);
3863
3864 let status = self.check_table_status(&table_name).await?;
3866 if !status.exists {
3867 return Err(NamespaceError::TableNotFound {
3868 message: table_name,
3869 }
3870 .into());
3871 }
3872 if status.is_deregistered {
3873 return Err(NamespaceError::TableNotFound {
3874 message: format!("Table is deregistered: {}", table_name),
3875 }
3876 .into());
3877 }
3878
3879 let mut dataset = self
3880 .configured_builder(&table_uri)
3881 .load()
3882 .await
3883 .map_err(|e| {
3884 Error::io_source(box_error(std::io::Error::other(format!(
3885 "Failed to open dataset: {}",
3886 e
3887 ))))
3888 })?;
3889
3890 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
3891 dataset.drop_columns(&columns).await.map_err(|e| {
3892 Error::io_source(box_error(std::io::Error::other(format!(
3893 "Failed to drop columns: {}",
3894 e
3895 ))))
3896 })?;
3897
3898 let version = dataset.version().version as i64;
3899 Ok(AlterTableDropColumnsResponse::new(version))
3900 }
3901
3902 async fn list_table_versions(
3903 &self,
3904 request: ListTableVersionsRequest,
3905 ) -> Result<ListTableVersionsResponse> {
3906 self.record_op("list_table_versions");
3907 let branch = Self::normalized_branch(request.branch.as_deref())?;
3908 let table_uri = self.resolve_table_location(&request.id).await?;
3909 let table_uri = match branch {
3910 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3911 None => table_uri,
3912 };
3913 let want_descending = request.descending == Some(true);
3914 let table_versions = self
3915 .list_table_versions_from_storage(&table_uri, want_descending, request.limit)
3916 .await?;
3917
3918 Ok(ListTableVersionsResponse {
3919 versions: table_versions,
3920 page_token: None,
3921 ..Default::default()
3922 })
3923 }
3924
3925 async fn create_table_version(
3926 &self,
3927 request: CreateTableVersionRequest,
3928 ) -> Result<CreateTableVersionResponse> {
3929 self.record_op("create_table_version");
3930 let branch = Self::normalized_branch(request.branch.as_deref())?;
3931 let table_uri = self.resolve_table_location(&request.id).await?;
3932 let (table_uri, table_path, branch_parent_version) = match branch {
3933 Some(b) => self.resolve_branch_for_commit(&table_uri, b).await?,
3934 None => {
3935 let table_path = self.object_store_path_from_uri(&table_uri)?;
3936 (table_uri, table_path, None)
3937 }
3938 };
3939
3940 let staging_manifest_path = &request.manifest_path;
3941 let version = request.version as u64;
3942
3943 let naming_scheme = match request.naming_scheme.as_deref() {
3945 Some("V1") => ManifestNamingScheme::V1,
3946 _ => ManifestNamingScheme::V2,
3947 };
3948
3949 let final_path = naming_scheme.manifest_path(&table_path, version);
3951
3952 let staging_path = Path::parse(staging_manifest_path).map_err(|e| {
3953 lance_core::Error::from(NamespaceError::InvalidInput {
3954 message: format!(
3955 "Invalid staging manifest path '{}': {}",
3956 staging_manifest_path, e
3957 ),
3958 })
3959 })?;
3960
3961 match self.object_store.inner.head(&final_path).await {
3963 Ok(existing_meta) => {
3964 return self
3965 .resolve_existing_table_version(ExistingTableVersionResolve {
3966 staging_path: &staging_path,
3967 final_path: &final_path,
3968 version,
3969 table_uri: &table_uri,
3970 final_meta: &existing_meta,
3971 request_manifest_size: request.manifest_size,
3972 })
3973 .await;
3974 }
3975 Err(ObjectStoreError::NotFound { .. }) => {}
3976 Err(e) => {
3977 return Err(lance_core::Error::from(NamespaceError::Internal {
3978 message: format!(
3979 "Failed to stat version {} for table at '{}': {}",
3980 version, table_uri, e
3981 ),
3982 }));
3983 }
3984 }
3985
3986 let is_branch = branch.is_some();
3989 self.enforce_create_table_version_cas(
3990 &table_path,
3991 version,
3992 &table_uri,
3993 is_branch,
3994 branch_parent_version,
3995 )
3996 .await?;
3997
3998 let copy_result = self
4000 .materialize_version_manifest_create(&staging_path, &final_path, staging_manifest_path)
4001 .await;
4002
4003 match copy_result {
4004 Ok(()) => {}
4005 Err(ObjectStoreError::AlreadyExists { .. })
4006 | Err(ObjectStoreError::Precondition { .. }) => {
4007 let existing_meta = self.object_store.inner.head(&final_path).await.map_err(|e| {
4009 lance_core::Error::from(NamespaceError::Internal {
4010 message: format!(
4011 "Version {} conflict for table at '{}' but failed to stat winner: {}",
4012 version, table_uri, e
4013 ),
4014 })
4015 })?;
4016 return self
4017 .resolve_existing_table_version(ExistingTableVersionResolve {
4018 staging_path: &staging_path,
4019 final_path: &final_path,
4020 version,
4021 table_uri: &table_uri,
4022 final_meta: &existing_meta,
4023 request_manifest_size: request.manifest_size,
4024 })
4025 .await;
4026 }
4027 Err(ObjectStoreError::NotFound { .. }) => {
4028 return Err(lance_core::Error::from(NamespaceError::InvalidInput {
4029 message: format!(
4030 "Staging manifest not found at '{}' for version {} of table at '{}'",
4031 staging_manifest_path, version, table_uri
4032 ),
4033 }));
4034 }
4035 Err(e) => {
4036 return Err(lance_core::Error::from(NamespaceError::Internal {
4037 message: format!(
4038 "Failed to create version {} for table at '{}': {}",
4039 version, table_uri, e
4040 ),
4041 }));
4042 }
4043 }
4044
4045 let final_meta = self
4046 .object_store
4047 .inner
4048 .head(&final_path)
4049 .await
4050 .map_err(|e| {
4051 lance_core::Error::from(NamespaceError::Internal {
4052 message: format!(
4053 "Failed to stat created version {} for table at '{}': {}",
4054 version, table_uri, e
4055 ),
4056 })
4057 })?;
4058
4059 if let Err(e) = self.object_store.inner.delete(&staging_path).await {
4061 log::warn!(
4062 "Failed to delete staging manifest at '{}': {:?}",
4063 staging_path,
4064 e
4065 );
4066 }
4067
4068 Ok(Self::create_table_version_response(
4069 version,
4070 &final_path,
4071 &final_meta,
4072 ))
4073 }
4074
4075 async fn describe_table_version(
4076 &self,
4077 request: DescribeTableVersionRequest,
4078 ) -> Result<DescribeTableVersionResponse> {
4079 self.record_op("describe_table_version");
4080 let branch = Self::normalized_branch(request.branch.as_deref())?;
4081 let table_uri = self.resolve_table_location(&request.id).await?;
4082 let table_uri = match branch {
4083 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
4084 None => table_uri,
4085 };
4086 let versions = self
4087 .list_table_versions_from_storage(&table_uri, true, None)
4088 .await?;
4089 let table_version = if let Some(requested_version) = request.version {
4090 versions
4091 .into_iter()
4092 .find(|version| version.version == requested_version)
4093 .ok_or_else(|| {
4094 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4095 message: format!(
4096 "version {} for table {}",
4097 requested_version,
4098 Self::format_table_id_from_request(&request.id)
4099 ),
4100 })
4101 })?
4102 } else {
4103 versions.into_iter().next().ok_or_else(|| {
4104 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4105 message: format!(
4106 "latest version for table {}",
4107 Self::format_table_id_from_request(&request.id)
4108 ),
4109 })
4110 })?
4111 };
4112
4113 Ok(DescribeTableVersionResponse {
4114 version: Box::new(table_version),
4115 ..Default::default()
4116 })
4117 }
4118
4119 async fn batch_delete_table_versions(
4120 &self,
4121 request: BatchDeleteTableVersionsRequest,
4122 ) -> Result<BatchDeleteTableVersionsResponse> {
4123 self.record_op("batch_delete_table_versions");
4124 let branch = Self::normalized_branch(request.branch.as_deref())?;
4125 let ranges: Vec<(i64, i64)> = request
4128 .ranges
4129 .iter()
4130 .map(|r| (r.start_version, r.end_version))
4131 .collect();
4132
4133 const MAX_VERSIONS_PER_REQUEST: i128 = 1_000_000;
4137 let requested: i128 = ranges
4138 .iter()
4139 .map(|(s, e)| {
4140 if *e < 0 {
4141 0
4142 } else {
4143 (*e as i128 - *s as i128).max(0)
4144 }
4145 })
4146 .sum();
4147 if requested > MAX_VERSIONS_PER_REQUEST {
4148 return Err(NamespaceError::InvalidInput {
4149 message: format!(
4150 "batch_delete requested {} versions; limit is {}",
4151 requested, MAX_VERSIONS_PER_REQUEST
4152 ),
4153 }
4154 .into());
4155 }
4156
4157 let table_entries = vec![TableDeleteEntry {
4158 table_id: request.id.clone(),
4159 ranges,
4160 }];
4161
4162 let total_deleted_count = self
4163 .delete_physical_version_files(&table_entries, branch)
4164 .await?;
4165
4166 Ok(BatchDeleteTableVersionsResponse {
4167 deleted_count: Some(total_deleted_count),
4168 transaction_id: None,
4169 ..Default::default()
4170 })
4171 }
4172
4173 async fn create_table_index(
4174 &self,
4175 request: CreateTableIndexRequest,
4176 ) -> Result<CreateTableIndexResponse> {
4177 self.record_op("create_table_index");
4178 let table_uri = self.resolve_table_location(&request.id).await?;
4179 let mut dataset = self
4180 .load_dataset(&table_uri, None, "create_table_index")
4181 .await?;
4182 let index_request = Self::build_index_params(&request)?;
4183
4184 dataset
4185 .create_index(
4186 &[request.column.as_str()],
4187 index_request.index_type(),
4188 request.name.clone(),
4189 index_request.params(),
4190 false,
4191 )
4192 .await
4193 .map_err(|e| {
4194 let err_msg = format!("{}", e);
4195 let ns_err = if err_msg.contains("already exists") {
4196 NamespaceError::TableIndexAlreadyExists {
4197 message: format!(
4198 "Index '{}' already exists on table '{}': {:?}",
4199 request.name.as_deref().unwrap_or("<auto-generated>"),
4200 table_uri,
4201 e
4202 ),
4203 }
4204 } else if err_msg.contains("not found") || err_msg.contains("does not exist") {
4205 NamespaceError::TableColumnNotFound {
4206 message: format!(
4207 "Column '{}' not found for table '{}': {:?}",
4208 request.column, table_uri, e
4209 ),
4210 }
4211 } else {
4212 NamespaceError::Internal {
4213 message: format!(
4214 "Failed to create {} index '{}' on column '{}' for table '{}': {:?}",
4215 request.index_type,
4216 request.name.as_deref().unwrap_or("<auto-generated>"),
4217 request.column,
4218 table_uri,
4219 e
4220 ),
4221 }
4222 };
4223 lance_core::Error::from(ns_err)
4224 })?;
4225
4226 let transaction_id = dataset
4227 .read_transaction()
4228 .await
4229 .map_err(|e| {
4230 lance_core::Error::from(NamespaceError::Internal {
4231 message: format!(
4232 "Failed to read committed transaction after creating index on '{}': {}",
4233 table_uri, e
4234 ),
4235 })
4236 })?
4237 .map(|transaction| transaction.uuid);
4238
4239 Ok(CreateTableIndexResponse {
4240 transaction_id,
4241 ..Default::default()
4242 })
4243 }
4244
4245 async fn list_table_indices(
4246 &self,
4247 request: ListTableIndicesRequest,
4248 ) -> Result<ListTableIndicesResponse> {
4249 self.record_op("list_table_indices");
4250 let table_uri = self.resolve_table_location(&request.id).await?;
4251 let dataset = self
4252 .load_dataset(&table_uri, request.version, "list_table_indices")
4253 .await?;
4254 let total_rows = dataset.count_rows(None).await.map_err(|e| {
4255 lance_core::Error::from(NamespaceError::Internal {
4256 message: format!("Failed to count rows for table '{}': {:?}", table_uri, e),
4257 })
4258 })? as u64;
4259 let mut indices = dataset
4260 .describe_indices(None)
4261 .await
4262 .map_err(|e| {
4263 lance_core::Error::from(NamespaceError::Internal {
4264 message: format!("Failed to describe table indices for '{}': {:?}", table_uri, e),
4265 })
4266 })?
4267 .into_iter()
4268 .filter(|description| {
4269 description
4270 .metadata()
4271 .first()
4272 .map(|metadata| !is_system_index(metadata))
4273 .unwrap_or(false)
4274 })
4275 .map(|description| {
4276 let columns = description
4277 .field_ids()
4278 .iter()
4279 .map(|field_id| {
4280 dataset
4281 .schema()
4282 .field_path_minimal(i32::try_from(*field_id).map_err(|e| {
4283 lance_core::Error::from(NamespaceError::Internal {
4284 message: format!(
4285 "Field id {} does not fit in i32 for table '{}': {}",
4286 field_id, table_uri, e
4287 ),
4288 })
4289 })?)
4290 .map_err(|e| {
4291 lance_core::Error::from(NamespaceError::Internal {
4292 message: format!(
4293 "Failed to resolve field path for field_id {} in table '{}': {}",
4294 field_id, table_uri, e
4295 ),
4296 })
4297 })
4298 })
4299 .collect::<Result<Vec<_>>>()?;
4300
4301 let segments = description.segments();
4302 let created_at = segments
4303 .iter()
4304 .filter_map(|segment| segment.created_at)
4305 .min()
4306 .map(|ts| ts.to_rfc3339());
4307
4308 #[allow(clippy::needless_update)]
4311 let content = IndexContent {
4312 index_name: description.name().to_string(),
4313 index_uuid: description.metadata()[0].uuid.to_string(),
4314 columns,
4315 status: "SUCCEEDED".to_string(),
4316 index_type: Some(description.index_type().to_string()),
4317 type_url: Some(description.type_url().to_string()),
4318 num_indexed_rows: Some(description.rows_indexed() as i64),
4319 num_unindexed_rows: Some(
4320 total_rows.saturating_sub(description.rows_indexed()) as i64,
4321 ),
4322 size_bytes: description.total_size_bytes().map(|size| size as i64),
4323 num_segments: Some(segments.len() as i32),
4324 created_at,
4325 index_version: segments.first().map(|segment| segment.index_version),
4326 index_details: description.details().ok(),
4327 ..Default::default()
4328 };
4329 Ok(content)
4330 })
4331 .collect::<Result<Vec<_>>>()?;
4332
4333 let page_token = Self::paginate_indices(&mut indices, request.page_token, request.limit);
4334 Ok(ListTableIndicesResponse {
4335 indexes: indices,
4336 page_token,
4337 ..Default::default()
4338 })
4339 }
4340
4341 async fn describe_table_index_stats(
4342 &self,
4343 request: DescribeTableIndexStatsRequest,
4344 ) -> Result<DescribeTableIndexStatsResponse> {
4345 self.record_op("describe_table_index_stats");
4346 let table_uri = self.resolve_table_location(&request.id).await?;
4347 let dataset = self
4348 .load_dataset(&table_uri, request.version, "describe_table_index_stats")
4349 .await?;
4350 let index_name = request.index_name.as_deref().ok_or_else(|| {
4351 lance_core::Error::from(NamespaceError::InvalidInput {
4352 message: "Index name is required for describe_table_index_stats".to_string(),
4353 })
4354 })?;
4355 let metadatas = dataset
4356 .load_indices_by_name(index_name)
4357 .await
4358 .map_err(|e| {
4359 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4360 message: format!(
4361 "Failed to load index '{}' metadata for table '{}': {}",
4362 index_name, table_uri, e
4363 ),
4364 })
4365 })?;
4366 if metadatas.first().is_some_and(is_system_index) {
4367 return Err(NamespaceError::Unsupported {
4368 message: format!("System index '{}' is not exposed by this API", index_name),
4369 }
4370 .into());
4371 }
4372
4373 let stats = <Dataset as DatasetIndexExt>::index_statistics(&dataset, index_name)
4374 .await
4375 .map_err(|e| {
4376 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4377 message: format!(
4378 "Failed to describe index statistics for '{}' on table '{}': {}",
4379 index_name, table_uri, e
4380 ),
4381 })
4382 })?;
4383 let stats: serde_json::Value = serde_json::from_str(&stats).map_err(|e| {
4384 lance_core::Error::from(NamespaceError::Internal {
4385 message: format!(
4386 "Failed to parse index statistics for '{}' on table '{}': {}",
4387 index_name, table_uri, e
4388 ),
4389 })
4390 })?;
4391
4392 Ok(Self::describe_table_index_stats_response(&stats))
4393 }
4394
4395 async fn describe_transaction(
4396 &self,
4397 request: DescribeTransactionRequest,
4398 ) -> Result<DescribeTransactionResponse> {
4399 self.record_op("describe_transaction");
4400 let mut request_id = request.id.ok_or_else(|| {
4401 lance_core::Error::from(NamespaceError::InvalidInput {
4402 message: "Transaction id must include table id and transaction identifier"
4403 .to_string(),
4404 })
4405 })?;
4406 if request_id.len() < 2 {
4407 return Err(NamespaceError::InvalidInput {
4408 message: format!(
4409 "Transaction request id must include table id and transaction identifier, got {:?}",
4410 request_id
4411 ),
4412 }
4413 .into());
4414 }
4415
4416 let id = request_id.pop().expect("request_id len checked above");
4417 let table_id = Some(request_id);
4418 let table_uri = self.resolve_table_location(&table_id).await?;
4419 let dataset = self
4420 .load_dataset(&table_uri, None, "describe_transaction")
4421 .await?;
4422 let (version, transaction) = self.find_transaction(&dataset, &id).await?;
4423
4424 let sidecar = self
4427 .load_transaction_alteration(&table_uri, &transaction.uuid)
4428 .await?;
4429
4430 Ok(Self::transaction_response(version, &transaction, sidecar))
4431 }
4432
4433 async fn alter_transaction(
4434 &self,
4435 request: AlterTransactionRequest,
4436 ) -> Result<AlterTransactionResponse> {
4437 self.record_op("alter_transaction");
4438
4439 let mut request_id = request.id.ok_or_else(|| {
4441 lance_core::Error::from(NamespaceError::InvalidInput {
4442 message: "Transaction id must include table id and transaction identifier"
4443 .to_string(),
4444 })
4445 })?;
4446 if request_id.len() < 2 {
4447 return Err(NamespaceError::InvalidInput {
4448 message: format!(
4449 "Transaction request id must include table id and transaction identifier, got {:?}",
4450 request_id
4451 ),
4452 }
4453 .into());
4454 }
4455
4456 let txn_id = request_id.pop().expect("request_id len checked above");
4457 let table_id = Some(request_id);
4458 let table_uri = self.resolve_table_location(&table_id).await?;
4459 let dataset = self
4460 .load_dataset(&table_uri, None, "alter_transaction")
4461 .await?;
4462 let (version, transaction) = self.find_transaction(&dataset, &txn_id).await?;
4463
4464 const RESERVED_KEYS: &[&str] = &["uuid", "version", "read_version", "operation", "tag"];
4468 let is_reserved = |key: &str| RESERVED_KEYS.contains(&key);
4469
4470 let mut sidecar = self
4473 .load_transaction_alteration(&table_uri, &transaction.uuid)
4474 .await?
4475 .unwrap_or_default();
4476
4477 for action in &request.actions {
4478 if let Some(ref set_status) = action.set_status_action
4479 && let Some(ref status) = set_status.status
4480 {
4481 let normalized = status.to_lowercase().replace('_', "");
4483 match normalized.as_str() {
4484 "queued" | "running" | "succeeded" | "failed" | "canceled" => {
4485 sidecar.status = Some(status.clone());
4486 }
4487 _ => {
4488 return Err(NamespaceError::InvalidInput {
4489 message: format!(
4490 "Invalid transaction status '{}'. Valid values are: Queued, Running, Succeeded, Failed, Canceled",
4491 status
4492 ),
4493 }
4494 .into());
4495 }
4496 }
4497 }
4498
4499 if let Some(ref set_property) = action.set_property_action
4500 && let (Some(key), Some(value)) = (&set_property.key, &set_property.value)
4501 {
4502 if is_reserved(key) {
4503 return Err(NamespaceError::InvalidInput {
4504 message: format!("Property '{}' is reserved and cannot be modified", key),
4505 }
4506 .into());
4507 }
4508 let mode = set_property
4509 .mode
4510 .as_deref()
4511 .unwrap_or("Overwrite")
4512 .to_lowercase();
4513 match mode.as_str() {
4514 "overwrite" => {
4515 sidecar.properties.insert(key.clone(), value.clone());
4516 }
4517 "fail" => {
4518 let exists = sidecar.properties.contains_key(key)
4521 || transaction
4522 .transaction_properties
4523 .as_ref()
4524 .is_some_and(|props| props.contains_key(key));
4525 if exists {
4526 return Err(NamespaceError::ConcurrentModification {
4527 message: format!(
4528 "Property '{}' already exists and mode is 'Fail'",
4529 key
4530 ),
4531 }
4532 .into());
4533 }
4534 sidecar.properties.insert(key.clone(), value.clone());
4535 }
4536 "skip" => {
4537 let exists = sidecar.properties.contains_key(key)
4538 || transaction
4539 .transaction_properties
4540 .as_ref()
4541 .is_some_and(|props| props.contains_key(key));
4542 if !exists {
4543 sidecar.properties.insert(key.clone(), value.clone());
4544 }
4545 }
4546 _ => {
4547 return Err(NamespaceError::InvalidInput {
4548 message: format!(
4549 "Invalid set_property mode '{}'. Valid values are: Overwrite, Fail, Skip",
4550 mode
4551 ),
4552 }
4553 .into());
4554 }
4555 }
4556 }
4557
4558 if let Some(ref unset_property) = action.unset_property_action
4559 && let Some(ref key) = unset_property.key
4560 {
4561 if is_reserved(key) {
4562 return Err(NamespaceError::InvalidInput {
4563 message: format!("Property '{}' is reserved and cannot be modified", key),
4564 }
4565 .into());
4566 }
4567 let mode = unset_property
4568 .mode
4569 .as_deref()
4570 .unwrap_or("Skip")
4571 .to_lowercase();
4572 let exists_in_transaction = transaction
4573 .transaction_properties
4574 .as_ref()
4575 .is_some_and(|props| props.contains_key(key));
4576 match mode.as_str() {
4577 "skip" => {
4578 sidecar.properties.remove(key);
4579 if exists_in_transaction {
4580 sidecar.removed_properties.insert(key.clone());
4583 }
4584 }
4585 "fail" => {
4586 if !sidecar.properties.contains_key(key) && !exists_in_transaction {
4587 return Err(NamespaceError::InvalidInput {
4588 message: format!(
4589 "Property '{}' does not exist and mode is 'Fail'",
4590 key
4591 ),
4592 }
4593 .into());
4594 }
4595 sidecar.properties.remove(key);
4596 if exists_in_transaction {
4597 sidecar.removed_properties.insert(key.clone());
4598 }
4599 }
4600 _ => {
4601 return Err(NamespaceError::InvalidInput {
4602 message: format!(
4603 "Invalid unset_property mode '{}'. Valid values are: Skip, Fail",
4604 mode
4605 ),
4606 }
4607 .into());
4608 }
4609 }
4610 }
4611 }
4612
4613 self.save_transaction_alteration(&table_uri, &transaction.uuid, &sidecar)
4617 .await?;
4618
4619 let final_status = sidecar
4622 .status
4623 .clone()
4624 .unwrap_or_else(|| "SUCCEEDED".to_string());
4625 let response = Self::transaction_response(version, &transaction, Some(sidecar));
4626 Ok(AlterTransactionResponse {
4627 status: final_status,
4628 properties: response.properties,
4629 ..Default::default()
4630 })
4631 }
4632
4633 async fn create_table_scalar_index(
4634 &self,
4635 request: CreateTableIndexRequest,
4636 ) -> Result<CreateTableScalarIndexResponse> {
4637 self.record_op("create_table_scalar_index");
4638 let index_type = Self::parse_index_type(&request.index_type)?;
4639 if !index_type.is_scalar() {
4640 return Err(NamespaceError::InvalidInput {
4641 message: format!(
4642 "create_table_scalar_index only supports scalar index types, got {}",
4643 request.index_type
4644 ),
4645 }
4646 .into());
4647 }
4648
4649 let response = self.create_table_index(request).await?;
4650 Ok(CreateTableScalarIndexResponse {
4651 transaction_id: response.transaction_id,
4652 ..Default::default()
4653 })
4654 }
4655
4656 async fn drop_table_index(
4657 &self,
4658 request: DropTableIndexRequest,
4659 ) -> Result<DropTableIndexResponse> {
4660 self.record_op("drop_table_index");
4661 let table_uri = self.resolve_table_location(&request.id).await?;
4662 let index_name = request.index_name.as_deref().ok_or_else(|| {
4663 lance_core::Error::from(NamespaceError::InvalidInput {
4664 message: "Index name is required for drop_table_index".to_string(),
4665 })
4666 })?;
4667 let mut dataset = self
4668 .load_dataset(&table_uri, None, "drop_table_index")
4669 .await?;
4670 let metadatas = dataset
4671 .load_indices_by_name(index_name)
4672 .await
4673 .map_err(|e| {
4674 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4675 message: format!(
4676 "Failed to load index '{}' before dropping it from table '{}': {}",
4677 index_name, table_uri, e
4678 ),
4679 })
4680 })?;
4681 if metadatas.first().is_some_and(is_system_index) {
4682 return Err(NamespaceError::Unsupported {
4683 message: format!(
4684 "System index '{}' cannot be dropped via this API",
4685 index_name
4686 ),
4687 }
4688 .into());
4689 }
4690
4691 dataset.drop_index(index_name).await.map_err(|e| {
4692 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4693 message: format!(
4694 "Failed to drop index '{}' from table '{}': {}",
4695 index_name, table_uri, e
4696 ),
4697 })
4698 })?;
4699
4700 let transaction_id = dataset
4701 .read_transaction()
4702 .await
4703 .map_err(|e| {
4704 lance_core::Error::from(NamespaceError::Internal {
4705 message: format!(
4706 "Failed to read committed transaction after dropping index '{}' from '{}': {}",
4707 index_name, table_uri, e
4708 ),
4709 })
4710 })?
4711 .map(|transaction| transaction.uuid);
4712
4713 Ok(DropTableIndexResponse {
4714 transaction_id,
4715 ..Default::default()
4716 })
4717 }
4718
4719 async fn list_all_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
4720 let mut tables = self.list_directory_tables().await?;
4723 tables = self
4724 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
4725 .await?;
4726 Self::apply_pagination(&mut tables, request.page_token, request.limit);
4727 Ok(ListTablesResponse::new(tables))
4728 }
4729
4730 async fn restore_table(&self, request: RestoreTableRequest) -> Result<RestoreTableResponse> {
4731 let version = request.version;
4732 if version < 0 {
4733 return Err(Error::invalid_input_source(
4734 format!(
4735 "Table version for restore_table must be non-negative, got {}",
4736 version
4737 )
4738 .into(),
4739 ));
4740 }
4741
4742 let branch = Self::normalized_branch(request.branch.as_deref())?;
4743 let table_uri = self.resolve_table_location(&request.id).await?;
4744 let mut dataset = match branch {
4745 Some(branch) => self.open_validated_branch(&table_uri, branch).await?,
4746 None => self.load_dataset(&table_uri, None, "restore_table").await?,
4747 };
4748
4749 dataset = dataset
4750 .checkout_version(version as u64)
4751 .await
4752 .map_err(|e| {
4753 Error::namespace_source(
4754 format!(
4755 "Failed to checkout version {} for restore at '{}': {}",
4756 version, table_uri, e
4757 )
4758 .into(),
4759 )
4760 })?;
4761
4762 dataset.restore().await.map_err(|e| {
4763 Error::namespace_source(
4764 format!(
4765 "Failed to restore table at '{}' to version {}: {}",
4766 table_uri, version, e
4767 )
4768 .into(),
4769 )
4770 })?;
4771
4772 let transaction_id = dataset
4773 .read_transaction()
4774 .await
4775 .map_err(|e| {
4776 Error::namespace_source(
4777 format!(
4778 "Failed to read transaction after restoring '{}': {}",
4779 table_uri, e
4780 )
4781 .into(),
4782 )
4783 })?
4784 .map(|t| t.uuid);
4785
4786 Ok(RestoreTableResponse {
4787 transaction_id,
4788 ..Default::default()
4789 })
4790 }
4791
4792 async fn update_table_schema_metadata(
4793 &self,
4794 request: UpdateTableSchemaMetadataRequest,
4795 ) -> Result<UpdateTableSchemaMetadataResponse> {
4796 let table_uri = self.resolve_table_location(&request.id).await?;
4797 let mut dataset = self
4798 .load_dataset(&table_uri, None, "update_table_schema_metadata")
4799 .await?;
4800
4801 let new_metadata = request.metadata.unwrap_or_default();
4802 let updated_metadata = dataset
4803 .update_schema_metadata(new_metadata.iter().map(|(k, v)| (k.as_str(), v.as_str())))
4804 .await
4805 .map_err(|e| {
4806 Error::namespace_source(
4807 format!(
4808 "Failed to update schema metadata for table at '{}': {}",
4809 table_uri, e
4810 )
4811 .into(),
4812 )
4813 })?;
4814
4815 let transaction_id = dataset
4816 .read_transaction()
4817 .await
4818 .map_err(|e| {
4819 Error::namespace_source(
4820 format!(
4821 "Failed to read transaction after updating metadata for '{}': {}",
4822 table_uri, e
4823 )
4824 .into(),
4825 )
4826 })?
4827 .map(|t| t.uuid);
4828
4829 Ok(UpdateTableSchemaMetadataResponse {
4830 metadata: Some(updated_metadata),
4831 transaction_id,
4832 ..Default::default()
4833 })
4834 }
4835
4836 async fn get_table_stats(
4837 &self,
4838 request: GetTableStatsRequest,
4839 ) -> Result<GetTableStatsResponse> {
4840 let table_uri = self.resolve_table_location(&request.id).await?;
4841 let dataset = Arc::new(
4842 self.load_dataset(&table_uri, None, "get_table_stats")
4843 .await?,
4844 );
4845
4846 let data_stats = dataset.calculate_data_stats().await.map_err(|e| {
4848 Error::namespace_source(
4849 format!(
4850 "Failed to calculate data statistics for table at '{}': {}",
4851 table_uri, e
4852 )
4853 .into(),
4854 )
4855 })?;
4856 let total_bytes: i64 = data_stats
4857 .fields
4858 .iter()
4859 .map(|f| f.bytes_on_disk as i64)
4860 .sum();
4861
4862 let fragment_row_futures: Vec<_> = dataset
4864 .get_fragments()
4865 .into_iter()
4866 .map(|f| async move { f.physical_rows().await })
4867 .collect();
4868 let fragment_row_results = futures::future::join_all(fragment_row_futures).await;
4869 let mut fragment_row_counts: Vec<i64> = fragment_row_results
4870 .into_iter()
4871 .filter_map(|r| r.ok())
4872 .map(|r| r as i64)
4873 .collect();
4874
4875 let num_fragments = fragment_row_counts.len() as i64;
4876 let num_rows: i64 = fragment_row_counts.iter().sum();
4877
4878 const SMALL_FRAGMENT_THRESHOLD: i64 = 1024 * 1024;
4881 let num_small_fragments = fragment_row_counts
4882 .iter()
4883 .filter(|&&r| r < SMALL_FRAGMENT_THRESHOLD)
4884 .count() as i64;
4885
4886 fragment_row_counts.sort_unstable();
4888 let lengths = if fragment_row_counts.is_empty() {
4889 FragmentSummary::new(0, 0, 0, 0, 0, 0, 0)
4890 } else {
4891 let len = fragment_row_counts.len();
4892 let min = fragment_row_counts[0];
4893 let max = fragment_row_counts[len - 1];
4894 let mean = num_rows / num_fragments;
4895 let pct = |p: f64| fragment_row_counts[((len - 1) as f64 * p) as usize];
4896 FragmentSummary::new(min, max, mean, pct(0.25), pct(0.50), pct(0.75), pct(0.99))
4897 };
4898
4899 let indices = dataset.load_indices().await.map_err(|e| {
4901 Error::namespace_source(
4902 format!("Failed to load indices for table at '{}': {}", table_uri, e).into(),
4903 )
4904 })?;
4905 let num_indices = indices.iter().filter(|m| !is_system_index(m)).count() as i64;
4906
4907 let fragment_stats = FragmentStats::new(num_fragments, num_small_fragments, lengths);
4908 Ok(GetTableStatsResponse::new(
4909 total_bytes,
4910 num_rows,
4911 num_indices,
4912 fragment_stats,
4913 ))
4914 }
4915
4916 async fn explain_table_query_plan(
4917 &self,
4918 request: ExplainTableQueryPlanRequest,
4919 ) -> Result<String> {
4920 let table_uri = self.resolve_table_location(&request.id).await?;
4921 let dataset = self
4922 .load_dataset(
4923 &table_uri,
4924 request.query.version,
4925 "explain_table_query_plan",
4926 )
4927 .await?;
4928 let verbose = request.verbose.unwrap_or(false);
4929
4930 let mut scanner = dataset.scan();
4931 Self::apply_query_params_to_scanner(
4932 &mut scanner,
4933 request.query.filter.as_deref(),
4934 request.query.columns.as_deref(),
4935 request.query.vector_column.as_deref(),
4936 &request.query.vector,
4937 request.query.k,
4938 request.query.offset,
4939 request.query.prefilter,
4940 request.query.bypass_vector_index,
4941 request.query.nprobes,
4942 request.query.ef,
4943 request.query.refine_factor,
4944 request.query.distance_type.as_deref(),
4945 request.query.fast_search,
4946 request.query.with_row_id,
4947 request.query.lower_bound,
4948 request.query.upper_bound,
4949 "explain_table_query_plan",
4950 )?;
4951
4952 scanner.explain_plan(verbose).await.map_err(|e| {
4953 Error::namespace_source(
4954 format!(
4955 "Failed to explain query plan for table at '{}': {}",
4956 table_uri, e
4957 )
4958 .into(),
4959 )
4960 })
4961 }
4962
4963 async fn analyze_table_query_plan(
4964 &self,
4965 request: AnalyzeTableQueryPlanRequest,
4966 ) -> Result<String> {
4967 let table_uri = self.resolve_table_location(&request.id).await?;
4968 let dataset = self
4969 .load_dataset(&table_uri, request.version, "analyze_table_query_plan")
4970 .await?;
4971
4972 let mut scanner = dataset.scan();
4973 Self::apply_query_params_to_scanner(
4974 &mut scanner,
4975 request.filter.as_deref(),
4976 request.columns.as_deref(),
4977 request.vector_column.as_deref(),
4978 &request.vector,
4979 request.k,
4980 request.offset,
4981 request.prefilter,
4982 request.bypass_vector_index,
4983 request.nprobes,
4984 request.ef,
4985 request.refine_factor,
4986 request.distance_type.as_deref(),
4987 request.fast_search,
4988 request.with_row_id,
4989 request.lower_bound,
4990 request.upper_bound,
4991 "analyze_table_query_plan",
4992 )?;
4993
4994 scanner.analyze_plan().await.map_err(|e| {
4995 Error::namespace_source(
4996 format!(
4997 "Failed to analyze query plan for table at '{}': {}",
4998 table_uri, e
4999 )
5000 .into(),
5001 )
5002 })
5003 }
5004
5005 async fn count_table_rows(&self, request: CountTableRowsRequest) -> Result<i64> {
5006 self.record_op("count_table_rows");
5007 let table_uri = self.resolve_table_location(&request.id).await?;
5008 let dataset = self
5009 .load_dataset(&table_uri, request.version, "count_table_rows")
5010 .await?;
5011
5012 let count =
5013 dataset
5014 .count_rows(request.predicate)
5015 .await
5016 .map_err(|e| NamespaceError::Internal {
5017 message: format!("Failed to count rows for table at '{}': {:?}", table_uri, e),
5018 })?;
5019
5020 Ok(count as i64)
5021 }
5022
5023 async fn insert_into_table(
5024 &self,
5025 request: InsertIntoTableRequest,
5026 request_data: Bytes,
5027 ) -> Result<InsertIntoTableResponse> {
5028 self.record_op("insert_into_table");
5029 let table_uri = self.resolve_table_location(&request.id).await?;
5030 let (reader, _num_rows) =
5031 Self::ipc_reader_from_request_data(&request_data, "insert_into_table")?;
5032
5033 let mode = match request.mode.as_deref() {
5034 Some(m) if m.eq_ignore_ascii_case("overwrite") => WriteMode::Overwrite,
5035 Some(m) if m.eq_ignore_ascii_case("append") => WriteMode::Append,
5036 None => WriteMode::Append,
5037 Some(m) => {
5038 return Err(lance_namespace::error::NamespaceError::InvalidInput {
5039 message: format!(
5040 "Unsupported write mode '{}'. Supported modes are: 'append', 'overwrite'",
5041 m
5042 ),
5043 }
5044 .into());
5045 }
5046 };
5047
5048 if !self.table_uri_has_actual_manifests(&table_uri).await? {
5049 self.write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
5050 .await?;
5051 } else {
5052 self.write_reader_to_table(&table_uri, reader, mode, None)
5053 .await?;
5054 }
5055
5056 Ok(InsertIntoTableResponse {
5057 transaction_id: None,
5058 ..Default::default()
5059 })
5060 }
5061
5062 async fn merge_insert_into_table(
5063 &self,
5064 request: MergeInsertIntoTableRequest,
5065 request_data: Bytes,
5066 ) -> Result<MergeInsertIntoTableResponse> {
5067 self.record_op("merge_insert_into_table");
5068 let table_uri = self.resolve_table_location(&request.id).await?;
5069 let on = request.on.as_ref().ok_or_else(|| {
5070 lance_core::Error::from(NamespaceError::InvalidInput {
5071 message: "'on' field is required for merge_insert_into_table".to_string(),
5072 })
5073 })?;
5074
5075 let table_has_manifests = self.table_uri_has_actual_manifests(&table_uri).await?;
5076 let (reader, num_rows) =
5077 Self::ipc_reader_from_request_data(&request_data, "merge_insert_into_table")?;
5078
5079 if !table_has_manifests {
5080 let dataset = self
5081 .write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
5082 .await?;
5083 let version = dataset.version().version as i64;
5084 return Ok(MergeInsertIntoTableResponse {
5085 transaction_id: None,
5086 num_updated_rows: Some(0),
5087 num_inserted_rows: Some(num_rows as i64),
5088 num_deleted_rows: Some(0),
5089 version: Some(version),
5090 ..Default::default()
5091 });
5092 }
5093
5094 let dataset = Arc::new(
5095 self.load_dataset(&table_uri, None, "merge_insert_into_table")
5096 .await?,
5097 );
5098
5099 let mut merge_builder = MergeInsertBuilder::try_new(dataset.clone(), vec![on.clone()])
5100 .map_err(|e| {
5101 lance_core::Error::from(NamespaceError::InvalidInput {
5102 message: format!("Failed to create merge_insert_into_table builder: {}", e),
5103 })
5104 })?;
5105
5106 if let Some(filter) = request.when_matched_update_all_filt.as_deref() {
5107 let behavior = WhenMatched::update_if(dataset.as_ref(), filter).map_err(|e| {
5108 lance_core::Error::from(NamespaceError::InvalidInput {
5109 message: format!(
5110 "Invalid when_matched_update_all_filt for merge_insert_into_table: {}",
5111 e
5112 ),
5113 })
5114 })?;
5115 merge_builder.when_matched(behavior);
5116 } else if request.when_matched_update_all.unwrap_or(false) {
5117 merge_builder.when_matched(WhenMatched::UpdateAll);
5118 }
5119
5120 if matches!(request.when_not_matched_insert_all, Some(false)) {
5121 merge_builder.when_not_matched(WhenNotMatched::DoNothing);
5122 } else {
5123 merge_builder.when_not_matched(WhenNotMatched::InsertAll);
5124 }
5125
5126 if let Some(filter) = request.when_not_matched_by_source_delete_filt.as_deref() {
5127 let behavior = WhenNotMatchedBySource::delete_if(dataset.as_ref(), filter).map_err(|e| {
5128 lance_core::Error::from(NamespaceError::InvalidInput {
5129 message: format!(
5130 "Invalid when_not_matched_by_source_delete_filt for merge_insert_into_table: {}",
5131 e
5132 ),
5133 })
5134 })?;
5135 merge_builder.when_not_matched_by_source(behavior);
5136 } else if request.when_not_matched_by_source_delete.unwrap_or(false) {
5137 merge_builder.when_not_matched_by_source(WhenNotMatchedBySource::Delete);
5138 }
5139
5140 if let Some(use_index) = request.use_index {
5141 merge_builder.use_index(use_index);
5142 }
5143
5144 let (dataset, stats) = merge_builder
5145 .try_build()
5146 .map_err(|e| {
5147 lance_core::Error::from(NamespaceError::InvalidInput {
5148 message: format!("Failed to build merge_insert_into_table job: {}", e),
5149 })
5150 })?
5151 .execute_reader(reader)
5152 .await
5153 .map_err(|e| Self::map_mutation_error(e, "merge_insert_into_table", &table_uri))?;
5154
5155 Ok(MergeInsertIntoTableResponse {
5156 transaction_id: None,
5157 num_updated_rows: Some(stats.num_updated_rows as i64),
5158 num_inserted_rows: Some(stats.num_inserted_rows as i64),
5159 num_deleted_rows: Some(stats.num_deleted_rows as i64),
5160 version: Some(dataset.version().version as i64),
5161 ..Default::default()
5162 })
5163 }
5164
5165 async fn update_table(&self, request: UpdateTableRequest) -> Result<UpdateTableResponse> {
5166 self.record_op("update_table");
5167
5168 if request.updates.is_empty() {
5169 return Err(NamespaceError::InvalidInput {
5170 message: "update_table requires at least one [column, expression] pair".to_string(),
5171 }
5172 .into());
5173 }
5174
5175 let mut seen_columns: HashMap<String, ()> = HashMap::with_capacity(request.updates.len());
5178 for (idx, pair) in request.updates.iter().enumerate() {
5179 if pair.len() != 2 {
5180 return Err(NamespaceError::InvalidInput {
5181 message: format!(
5182 "update_table updates[{}] must be a [column, expression] pair, got {} elements",
5183 idx,
5184 pair.len()
5185 ),
5186 }
5187 .into());
5188 }
5189 let column = &pair[0];
5190 if column.trim().is_empty() {
5191 return Err(NamespaceError::InvalidInput {
5192 message: format!("update_table updates[{}] has an empty column name", idx),
5193 }
5194 .into());
5195 }
5196 if seen_columns.insert(column.clone(), ()).is_some() {
5197 return Err(NamespaceError::InvalidInput {
5198 message: format!(
5199 "update_table cannot update column '{}' more than once",
5200 column
5201 ),
5202 }
5203 .into());
5204 }
5205 }
5206
5207 let table_uri = self.resolve_table_location(&request.id).await?;
5208 let dataset = Arc::new(self.load_dataset(&table_uri, None, "update_table").await?);
5209
5210 let mut builder = UpdateBuilder::new(dataset);
5211 for pair in &request.updates {
5212 builder = builder.set(&pair[0], &pair[1]).map_err(|e| {
5214 lance_core::Error::from(NamespaceError::InvalidInput {
5215 message: format!("Invalid update expression for column '{}': {}", pair[0], e),
5216 })
5217 })?;
5218 }
5219 if let Some(predicate) = request.predicate.as_deref()
5220 && !predicate.trim().is_empty()
5221 {
5222 builder = builder.update_where(predicate).map_err(|e| {
5223 lance_core::Error::from(NamespaceError::InvalidInput {
5224 message: format!("Invalid update_table predicate '{}': {}", predicate, e),
5225 })
5226 })?;
5227 }
5228
5229 let job = builder.build().map_err(|e| {
5230 lance_core::Error::from(NamespaceError::InvalidInput {
5231 message: format!("Failed to build update_table job: {}", e),
5232 })
5233 })?;
5234
5235 let result = job
5236 .execute()
5237 .await
5238 .map_err(|e| Self::map_mutation_error(e, "update_table", &table_uri))?;
5239
5240 let version = result.new_dataset.version().version as i64;
5241 Ok(UpdateTableResponse {
5242 transaction_id: None,
5243 updated_rows: result.rows_updated as i64,
5244 version,
5245 properties: None,
5246 ..Default::default()
5247 })
5248 }
5249
5250 async fn delete_from_table(
5251 &self,
5252 request: DeleteFromTableRequest,
5253 ) -> Result<DeleteFromTableResponse> {
5254 self.record_op("delete_from_table");
5255
5256 if request.predicate.trim().is_empty() {
5257 return Err(NamespaceError::InvalidInput {
5258 message: "delete_from_table requires a non-empty predicate".to_string(),
5259 }
5260 .into());
5261 }
5262
5263 let table_uri = self.resolve_table_location(&request.id).await?;
5264 let mut dataset = self
5265 .load_dataset(&table_uri, None, "delete_from_table")
5266 .await?;
5267
5268 let result = dataset
5269 .delete(&request.predicate)
5270 .await
5271 .map_err(|e| Self::map_mutation_error(e, "delete_from_table", &table_uri))?;
5272
5273 Ok(DeleteFromTableResponse {
5274 transaction_id: None,
5275 version: Some(result.new_dataset.version().version as i64),
5276 ..Default::default()
5277 })
5278 }
5279
5280 async fn query_table(&self, request: QueryTableRequest) -> Result<Bytes> {
5281 use arrow::ipc::writer::FileWriter;
5282
5283 self.record_op("query_table");
5284 let table_uri = self.resolve_table_location(&request.id).await?;
5285 let dataset = self
5286 .load_dataset(&table_uri, request.version, "query_table")
5287 .await?;
5288
5289 let mut scanner = dataset.scan();
5291
5292 let has_vector_query = request
5295 .vector
5296 .single_vector
5297 .as_ref()
5298 .map(|sv| !sv.is_empty())
5299 .unwrap_or(false)
5300 || request
5301 .vector
5302 .multi_vector
5303 .as_ref()
5304 .map(|mv| !mv.is_empty())
5305 .unwrap_or(false);
5306
5307 if let Some(prefilter) = request.prefilter {
5309 scanner.prefilter(prefilter);
5310 }
5311
5312 if has_vector_query {
5314 let vector_column = request.vector_column.as_deref().unwrap_or("vector");
5315
5316 let query_vector: Vec<f32> = request
5318 .vector
5319 .single_vector
5320 .clone()
5321 .or_else(|| {
5322 request
5323 .vector
5324 .multi_vector
5325 .as_ref()
5326 .and_then(|mv| mv.first().cloned())
5327 })
5328 .unwrap_or_default();
5329
5330 if !query_vector.is_empty() {
5331 let k = if request.k > 0 {
5332 request.k as usize
5333 } else {
5334 10
5335 };
5336 let query_array = Float32Array::from(query_vector);
5337 scanner
5338 .nearest(vector_column, &query_array, k)
5339 .map_err(|e| NamespaceError::InvalidInput {
5340 message: format!("Invalid vector search: {:?}", e),
5341 })?;
5342
5343 if let Some(ref distance_type) = request.distance_type {
5345 let metric = match distance_type.to_lowercase().as_str() {
5346 "l2" | "euclidean" => MetricType::L2,
5347 "cosine" => MetricType::Cosine,
5348 "dot" | "inner_product" => MetricType::Dot,
5349 "hamming" => MetricType::Hamming,
5350 _ => {
5351 return Err(NamespaceError::InvalidInput {
5352 message: format!("Unknown distance type: {}", distance_type),
5353 }
5354 .into());
5355 }
5356 };
5357 scanner.distance_metric(metric);
5358 }
5359
5360 if let Some(nprobes) = request.nprobes {
5362 scanner.minimum_nprobes(nprobes as usize);
5363 }
5364
5365 if let Some(ef) = request.ef {
5367 scanner.ef(ef as usize);
5368 }
5369
5370 if let Some(refine_factor) = request.refine_factor {
5372 scanner.refine(refine_factor as u32);
5373 }
5374
5375 if request.lower_bound.is_some() || request.upper_bound.is_some() {
5377 scanner.distance_range(request.lower_bound, request.upper_bound);
5378 }
5379
5380 if let Some(bypass) = request.bypass_vector_index {
5382 scanner.use_index(!bypass);
5383 }
5384
5385 if request.fast_search == Some(true) {
5387 scanner.fast_search();
5388 }
5389 }
5390 }
5391
5392 if let Some(ref fts_query) = request.full_text_query {
5394 if let Some(ref string_query) = fts_query.string_query {
5396 let mut fts = FullTextSearchQuery::new(string_query.query.clone());
5397
5398 if let Some(ref columns) = string_query.columns
5400 && !columns.is_empty()
5401 {
5402 fts = fts
5403 .with_columns(columns)
5404 .map_err(|e| NamespaceError::InvalidInput {
5405 message: format!("Invalid FTS columns: {:?}", e),
5406 })?;
5407 }
5408
5409 scanner
5410 .full_text_search(fts)
5411 .map_err(|e| NamespaceError::InvalidInput {
5412 message: format!("Invalid full text search: {:?}", e),
5413 })?;
5414 } else if let Some(ref structured_query) = fts_query.structured_query {
5415 let engine_query = build_engine_fts_query(&structured_query.query)?;
5417 let fts = FullTextSearchQuery::new_query(engine_query);
5418 scanner
5419 .full_text_search(fts)
5420 .map_err(|e| NamespaceError::InvalidInput {
5421 message: format!("Invalid full text search: {:?}", e),
5422 })?;
5423 }
5424 }
5425
5426 if let Some(ref columns) = request.columns {
5428 if let Some(ref column_names) = columns.column_names
5429 && !column_names.is_empty()
5430 {
5431 scanner
5432 .project(column_names)
5433 .map_err(|e| NamespaceError::InvalidInput {
5434 message: format!("Invalid column projection: {:?}", e),
5435 })?;
5436 } else if let Some(ref column_aliases) = columns.column_aliases
5437 && !column_aliases.is_empty()
5438 {
5439 let transform_pairs: Vec<(String, String)> = column_aliases
5441 .iter()
5442 .map(|(alias, sql)| (alias.clone(), sql.clone()))
5443 .collect();
5444 scanner
5445 .project_with_transform(
5446 &transform_pairs
5447 .iter()
5448 .map(|(a, s)| (a.as_str(), s.as_str()))
5449 .collect::<Vec<_>>(),
5450 )
5451 .map_err(|e| NamespaceError::InvalidInput {
5452 message: format!("Invalid column alias expression: {:?}", e),
5453 })?;
5454 }
5455 }
5456
5457 if let Some(ref filter) = request.filter
5459 && !filter.is_empty()
5460 {
5461 scanner
5462 .filter(filter)
5463 .map_err(|e| NamespaceError::InvalidInput {
5464 message: format!("Invalid filter expression: {:?}", e),
5465 })?;
5466 }
5467
5468 if request.with_row_id == Some(true) {
5470 scanner.with_row_id();
5471 }
5472
5473 if !has_vector_query && request.k > 0 {
5477 let offset = request.offset.map(|o| o as i64);
5478 scanner.limit(Some(request.k as i64), offset).map_err(|e| {
5479 NamespaceError::InvalidInput {
5480 message: format!("Invalid limit/offset: {:?}", e),
5481 }
5482 })?;
5483 } else if has_vector_query && request.offset.is_some() {
5484 let offset = request.offset.map(|o| o as i64);
5486 scanner
5487 .limit(None, offset)
5488 .map_err(|e| NamespaceError::InvalidInput {
5489 message: format!("Invalid offset: {:?}", e),
5490 })?;
5491 }
5492
5493 let batch = scanner
5495 .try_into_batch()
5496 .await
5497 .map_err(|e| NamespaceError::Internal {
5498 message: format!("Failed to execute query: {:?}", e),
5499 })?;
5500
5501 let schema = batch.schema();
5503 let mut buffer = Vec::new();
5504 {
5505 let mut writer = FileWriter::try_new(&mut buffer, &schema).map_err(|e| {
5506 NamespaceError::Internal {
5507 message: format!("Failed to create IPC writer: {:?}", e),
5508 }
5509 })?;
5510 writer.write(&batch).map_err(|e| NamespaceError::Internal {
5511 message: format!("Failed to write batch to IPC: {:?}", e),
5512 })?;
5513 writer.finish().map_err(|e| NamespaceError::Internal {
5514 message: format!("Failed to finish IPC writer: {:?}", e),
5515 })?;
5516 }
5517
5518 Ok(Bytes::from(buffer))
5519 }
5520
5521 async fn list_table_tags(
5522 &self,
5523 request: ListTableTagsRequest,
5524 ) -> Result<ListTableTagsResponse> {
5525 self.record_op("list_table_tags");
5526 let table_uri = self.resolve_table_location(&request.id).await?;
5527 let dataset = self
5528 .load_dataset(&table_uri, None, "list_table_tags")
5529 .await?;
5530
5531 let raw_tags = dataset.tags().list().await.map_err(|e| {
5532 lance_core::Error::from(NamespaceError::Internal {
5533 message: format!("Failed to list tags for table at '{}': {}", table_uri, e),
5534 })
5535 })?;
5536
5537 let tags = raw_tags
5538 .into_iter()
5539 .map(|(name, contents)| {
5540 let mut tag_model =
5541 ModelTagContents::new(contents.version as i64, contents.manifest_size as i64);
5542 tag_model.branch = contents.branch;
5543 (name, tag_model)
5544 })
5545 .collect();
5546
5547 Ok(ListTableTagsResponse {
5548 tags,
5549 page_token: None,
5550 ..Default::default()
5551 })
5552 }
5553
5554 async fn get_table_tag_version(
5555 &self,
5556 request: GetTableTagVersionRequest,
5557 ) -> Result<GetTableTagVersionResponse> {
5558 self.record_op("get_table_tag_version");
5559 if request.tag.is_empty() {
5560 return Err(NamespaceError::InvalidInput {
5561 message: "tag name must not be empty for get_table_tag_version".to_string(),
5562 }
5563 .into());
5564 }
5565
5566 let table_uri = self.resolve_table_location(&request.id).await?;
5567 let dataset = self
5568 .load_dataset(&table_uri, None, "get_table_tag_version")
5569 .await?;
5570
5571 let contents = dataset
5572 .tags()
5573 .get(&request.tag)
5574 .await
5575 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5576
5577 Ok(GetTableTagVersionResponse {
5578 version: contents.version as i64,
5579 branch: contents.branch,
5580 ..Default::default()
5581 })
5582 }
5583
5584 async fn create_table_tag(
5585 &self,
5586 request: CreateTableTagRequest,
5587 ) -> Result<CreateTableTagResponse> {
5588 self.record_op("create_table_tag");
5589 if request.tag.is_empty() {
5590 return Err(NamespaceError::InvalidInput {
5591 message: "tag name must not be empty for create_table_tag".to_string(),
5592 }
5593 .into());
5594 }
5595 if request.version <= 0 {
5596 return Err(NamespaceError::InvalidInput {
5597 message: format!(
5598 "tag version must be a positive integer, got {} for create_table_tag",
5599 request.version
5600 ),
5601 }
5602 .into());
5603 }
5604
5605 let table_uri = self.resolve_table_location(&request.id).await?;
5606 let dataset = self
5607 .load_dataset(&table_uri, None, "create_table_tag")
5608 .await?;
5609
5610 dataset
5611 .tags()
5612 .create(&request.tag, request.version as u64)
5613 .await
5614 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5615
5616 Ok(CreateTableTagResponse {
5617 transaction_id: None,
5618 ..Default::default()
5619 })
5620 }
5621
5622 async fn delete_table_tag(
5623 &self,
5624 request: DeleteTableTagRequest,
5625 ) -> Result<DeleteTableTagResponse> {
5626 self.record_op("delete_table_tag");
5627 if request.tag.is_empty() {
5628 return Err(NamespaceError::InvalidInput {
5629 message: "tag name must not be empty for delete_table_tag".to_string(),
5630 }
5631 .into());
5632 }
5633
5634 let table_uri = self.resolve_table_location(&request.id).await?;
5635 let dataset = self
5636 .load_dataset(&table_uri, None, "delete_table_tag")
5637 .await?;
5638
5639 dataset
5640 .tags()
5641 .delete(&request.tag)
5642 .await
5643 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5644
5645 Ok(DeleteTableTagResponse {
5646 transaction_id: None,
5647 ..Default::default()
5648 })
5649 }
5650
5651 async fn update_table_tag(
5652 &self,
5653 request: UpdateTableTagRequest,
5654 ) -> Result<UpdateTableTagResponse> {
5655 self.record_op("update_table_tag");
5656 if request.tag.is_empty() {
5657 return Err(NamespaceError::InvalidInput {
5658 message: "tag name must not be empty for update_table_tag".to_string(),
5659 }
5660 .into());
5661 }
5662 if request.version <= 0 {
5663 return Err(NamespaceError::InvalidInput {
5664 message: format!(
5665 "tag version must be a positive integer, got {} for update_table_tag",
5666 request.version
5667 ),
5668 }
5669 .into());
5670 }
5671
5672 let table_uri = self.resolve_table_location(&request.id).await?;
5673 let dataset = self
5674 .load_dataset(&table_uri, None, "update_table_tag")
5675 .await?;
5676
5677 dataset
5678 .tags()
5679 .update(&request.tag, request.version as u64)
5680 .await
5681 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5682
5683 Ok(UpdateTableTagResponse {
5684 transaction_id: None,
5685 ..Default::default()
5686 })
5687 }
5688
5689 async fn create_table_branch(
5690 &self,
5691 request: CreateTableBranchRequest,
5692 ) -> Result<CreateTableBranchResponse> {
5693 self.record_op("create_table_branch");
5694 if request.name.is_empty() {
5695 return Err(NamespaceError::InvalidInput {
5696 message: "branch name must not be empty for create_table_branch".to_string(),
5697 }
5698 .into());
5699 }
5700 let from_version = match request.from_version {
5701 Some(v) if v <= 0 => {
5702 return Err(NamespaceError::InvalidInput {
5703 message: format!(
5704 "from_version must be a positive integer, got {} for create_table_branch",
5705 v
5706 ),
5707 }
5708 .into());
5709 }
5710 Some(v) => Some(v as u64),
5711 None => None,
5712 };
5713
5714 let table_uri = self.resolve_table_location(&request.id).await?;
5715 let mut dataset = self
5716 .load_dataset(&table_uri, None, "create_table_branch")
5717 .await?;
5718
5719 if dataset.branches().get(&request.name).await.is_ok() {
5724 return Err(NamespaceError::TableBranchAlreadyExists {
5725 message: format!("branch '{}' for table at '{}'", request.name, table_uri),
5726 }
5727 .into());
5728 }
5729
5730 dataset
5731 .create_branch(
5732 &request.name,
5733 (request.from_branch.as_deref(), from_version),
5734 None,
5735 )
5736 .await
5737 .map_err(|e| {
5738 if matches!(e, lance_core::Error::DatasetNotFound { .. }) {
5741 NamespaceError::InvalidInput {
5742 message: format!(
5743 "from_branch/from_version for branch '{}' refers to a source that does not exist: {}",
5744 request.name, e
5745 ),
5746 }
5747 .into()
5748 } else {
5749 Self::map_branch_error(e, &request.name, &table_uri)
5750 }
5751 })?;
5752
5753 Ok(CreateTableBranchResponse {
5754 transaction_id: None,
5755 ..Default::default()
5756 })
5757 }
5758
5759 async fn list_table_branches(
5760 &self,
5761 request: ListTableBranchesRequest,
5762 ) -> Result<ListTableBranchesResponse> {
5763 self.record_op("list_table_branches");
5764 let table_uri = self.resolve_table_location(&request.id).await?;
5765 let dataset = self
5766 .load_dataset(&table_uri, None, "list_table_branches")
5767 .await?;
5768
5769 let raw_branches = dataset.list_branches().await.map_err(|e| {
5770 lance_core::Error::from(NamespaceError::Internal {
5771 message: format!(
5772 "Failed to list branches for table at '{}': {}",
5773 table_uri, e
5774 ),
5775 })
5776 })?;
5777
5778 let branches = raw_branches
5779 .into_iter()
5780 .map(|(name, contents)| {
5781 let mut branch_model = ModelBranchContents::new(
5784 contents.parent_version as i64,
5785 contents.create_at as i64,
5786 contents.manifest_size as i64,
5787 );
5788 branch_model.parent_branch = contents.parent_branch;
5789 branch_model.metadata = if contents.metadata.is_empty() {
5790 None
5791 } else {
5792 Some(contents.metadata)
5793 };
5794 (name, branch_model)
5795 })
5796 .collect();
5797
5798 Ok(ListTableBranchesResponse {
5799 branches,
5800 page_token: None,
5801 ..Default::default()
5802 })
5803 }
5804
5805 async fn delete_table_branch(
5806 &self,
5807 request: DeleteTableBranchRequest,
5808 ) -> Result<DeleteTableBranchResponse> {
5809 self.record_op("delete_table_branch");
5810 if request.name.is_empty() {
5811 return Err(NamespaceError::InvalidInput {
5812 message: "branch name must not be empty for delete_table_branch".to_string(),
5813 }
5814 .into());
5815 }
5816
5817 let table_uri = self.resolve_table_location(&request.id).await?;
5818 let mut dataset = self
5819 .load_dataset(&table_uri, None, "delete_table_branch")
5820 .await?;
5821
5822 dataset
5823 .delete_branch(&request.name)
5824 .await
5825 .map_err(|e| match e {
5826 lance_core::Error::RefConflict { message } => NamespaceError::InvalidInput {
5827 message: format!(
5828 "branch '{}' for table at '{}': {}",
5829 request.name, table_uri, message
5830 ),
5831 }
5832 .into(),
5833 other => Self::map_branch_error(other, &request.name, &table_uri),
5834 })?;
5835
5836 Ok(DeleteTableBranchResponse {
5837 transaction_id: None,
5838 ..Default::default()
5839 })
5840 }
5841
5842 fn namespace_id(&self) -> String {
5843 format!("DirectoryNamespace {{ root: {:?} }}", self.root)
5844 }
5845}
5846
5847#[derive(Debug)]
5849pub(crate) enum MarkerFileError {
5850 AlreadyExists { description: String },
5852 Other { message: String },
5854}
5855
5856impl std::fmt::Display for MarkerFileError {
5857 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5858 match self {
5859 Self::AlreadyExists { description } => {
5860 write!(f, "{} already exists", description)
5861 }
5862 Self::Other { message } => write!(f, "{}", message),
5863 }
5864 }
5865}
5866
5867pub(crate) async fn put_marker_file_atomic(
5888 object_store: &ObjectStore,
5889 path: &Path,
5890 file_description: &str,
5891) -> std::result::Result<(), MarkerFileError> {
5892 let staging_name = format!("lance-marker.staging.{}", uuid::Uuid::new_v4().simple());
5893 let path_str = path.as_ref();
5894 let staging_path = match path_str.rfind('/') {
5895 Some(idx) => Path::from(format!("{}/{}", &path_str[..idx], staging_name)),
5896 None => Path::from(staging_name.as_str()),
5897 };
5898
5899 object_store
5900 .inner
5901 .put(&staging_path, bytes::Bytes::from_static(b"reserved").into())
5902 .await
5903 .map_err(|e| MarkerFileError::Other {
5904 message: format!("Failed to stage {}: {:?}", file_description, e),
5905 })?;
5906
5907 let mut staging_consumed = false;
5910 let publish_result = match object_store
5911 .inner
5912 .rename_if_not_exists(&staging_path, path)
5913 .await
5914 {
5915 Ok(()) => {
5916 staging_consumed = true;
5917 Ok(())
5918 }
5919 Err(ObjectStoreError::NotImplemented { .. })
5920 | Err(ObjectStoreError::NotSupported { .. }) => {
5921 match object_store
5922 .inner
5923 .copy_if_not_exists(&staging_path, path)
5924 .await
5925 {
5926 Ok(()) => Ok(()),
5927 Err(ObjectStoreError::NotImplemented { .. })
5928 | Err(ObjectStoreError::NotSupported { .. }) => object_store
5929 .inner
5930 .put_opts(
5931 path,
5932 bytes::Bytes::from_static(b"reserved").into(),
5933 PutOptions {
5934 mode: PutMode::Create,
5935 ..Default::default()
5936 },
5937 )
5938 .await
5939 .map(|_| ()),
5940 Err(e) => Err(e),
5941 }
5942 }
5943 Err(e) => Err(e),
5944 };
5945
5946 if !staging_consumed {
5947 delete_staging_marker_best_effort(object_store, &staging_path).await;
5948 }
5949
5950 match publish_result {
5951 Ok(()) => Ok(()),
5952 Err(ObjectStoreError::AlreadyExists { .. })
5953 | Err(ObjectStoreError::Precondition { .. }) => Err(MarkerFileError::AlreadyExists {
5954 description: file_description.to_string(),
5955 }),
5956 Err(e) => Err(MarkerFileError::Other {
5957 message: format!("Failed to create {}: {:?}", file_description, e),
5958 }),
5959 }
5960}
5961
5962async fn delete_staging_marker_best_effort(object_store: &ObjectStore, staging_path: &Path) {
5966 const MAX_ATTEMPTS: u32 = 3;
5967 const BACKOFF_MS: [u64; 2] = [20, 50];
5968
5969 let mut last_err: Option<ObjectStoreError> = None;
5970 for attempt in 0..MAX_ATTEMPTS {
5971 match object_store.inner.delete(staging_path).await {
5972 Ok(()) => return,
5973 Err(ObjectStoreError::NotFound { .. }) => return,
5974 Err(e) => {
5975 last_err = Some(e);
5976 if let Some(&delay_ms) = BACKOFF_MS.get(attempt as usize) {
5977 tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
5978 }
5979 }
5980 }
5981 }
5982 if let Some(del_err) = last_err {
5983 log::warn!(
5984 "Failed to delete staging marker at '{}' after {} attempts: {:?}",
5985 staging_path,
5986 MAX_ATTEMPTS,
5987 del_err
5988 );
5989 }
5990}
5991
5992fn build_engine_fts_query(
5996 query: &lance_namespace::models::FtsQuery,
5997) -> std::result::Result<FtsQuery, NamespaceError> {
5998 if let Some(ref m) = query.r#match {
5999 Ok(FtsQuery::Match(build_engine_match_query(m)?))
6000 } else if let Some(ref p) = query.phrase {
6001 let mut phrase = PhraseQuery::new(p.terms.clone());
6002 if let Some(ref column) = p.column {
6003 phrase = phrase.with_column(Some(column.clone()));
6004 }
6005 if let Some(slop) = p.slop {
6006 phrase = phrase.with_slop(slop as u32);
6007 }
6008 Ok(FtsQuery::Phrase(phrase))
6009 } else if let Some(ref mm) = query.multi_match {
6010 let match_queries = mm
6011 .match_queries
6012 .iter()
6013 .map(build_engine_match_query)
6014 .collect::<std::result::Result<Vec<_>, _>>()?;
6015 Ok(FtsQuery::MultiMatch(MultiMatchQuery { match_queries }))
6016 } else if let Some(ref b) = query.boolean {
6017 let mut clauses: Vec<(Occur, FtsQuery)> = Vec::new();
6018 for clause in &b.must {
6019 clauses.push((Occur::Must, build_engine_fts_query(clause)?));
6020 }
6021 for clause in &b.should {
6022 clauses.push((Occur::Should, build_engine_fts_query(clause)?));
6023 }
6024 for clause in &b.must_not {
6025 clauses.push((Occur::MustNot, build_engine_fts_query(clause)?));
6026 }
6027 Ok(FtsQuery::Boolean(BooleanQuery::new(clauses)))
6028 } else if let Some(ref boost) = query.boost {
6029 let positive = build_engine_fts_query(&boost.positive)?;
6030 let negative = build_engine_fts_query(&boost.negative)?;
6031 Ok(FtsQuery::Boost(BoostQuery::new(
6032 positive,
6033 negative,
6034 boost.negative_boost,
6035 )))
6036 } else {
6037 Err(NamespaceError::InvalidInput {
6038 message: "structured_query.query must set exactly one of match, phrase, multi_match, \
6039 boolean, or boost"
6040 .to_string(),
6041 })
6042 }
6043}
6044
6045fn build_engine_match_query(
6046 m: &lance_namespace::models::MatchQuery,
6047) -> std::result::Result<MatchQuery, NamespaceError> {
6048 let mut match_query = MatchQuery::new(m.terms.clone());
6049 if let Some(ref column) = m.column {
6050 match_query = match_query.with_column(Some(column.clone()));
6051 }
6052 if let Some(boost) = m.boost {
6053 match_query = match_query.with_boost(boost);
6054 }
6055 if let Some(fuzziness) = m.fuzziness {
6056 match_query = match_query.with_fuzziness(Some(fuzziness as u32));
6057 }
6058 if let Some(max_expansions) = m.max_expansions {
6059 match_query = match_query.with_max_expansions(max_expansions as usize);
6060 }
6061 if let Some(ref operator) = m.operator {
6062 let op =
6063 Operator::try_from(operator.as_str()).map_err(|e| NamespaceError::InvalidInput {
6064 message: format!("Invalid FTS operator: {:?}", e),
6065 })?;
6066 match_query = match_query.with_operator(op);
6067 }
6068 if let Some(prefix_length) = m.prefix_length {
6069 match_query = match_query.with_prefix_length(prefix_length as u32);
6070 }
6071 Ok(match_query)
6072}
6073
6074#[cfg(test)]
6075mod tests {
6076 use super::*;
6077 use arrow_ipc::reader::{FileReader, StreamReader};
6078
6079 #[test]
6080 fn test_build_engine_fts_query_match() {
6081 let mut ns_match = lance_namespace::models::MatchQuery::new("hello world".to_string());
6082 ns_match.column = Some("body".to_string());
6083 ns_match.operator = Some("AND".to_string());
6084 ns_match.fuzziness = Some(1);
6085 ns_match.max_expansions = Some(30);
6086 ns_match.boost = Some(2.0);
6087 ns_match.prefix_length = Some(2);
6088
6089 let mut ns_query = lance_namespace::models::FtsQuery::new();
6090 ns_query.r#match = Some(Box::new(ns_match));
6091
6092 match build_engine_fts_query(&ns_query).unwrap() {
6093 FtsQuery::Match(m) => {
6094 assert_eq!(m.terms, "hello world");
6095 assert_eq!(m.column, Some("body".to_string()));
6096 assert_eq!(m.operator, Operator::And);
6097 assert_eq!(m.fuzziness, Some(1));
6098 assert_eq!(m.max_expansions, 30);
6099 assert_eq!(m.boost, 2.0);
6100 assert_eq!(m.prefix_length, 2);
6101 }
6102 other => panic!("expected Match, got {:?}", other),
6103 }
6104 }
6105
6106 fn ns_match_query(terms: &str, column: &str) -> lance_namespace::models::FtsQuery {
6109 let mut m = lance_namespace::models::MatchQuery::new(terms.to_string());
6110 m.column = Some(column.to_string());
6111 let mut q = lance_namespace::models::FtsQuery::new();
6112 q.r#match = Some(Box::new(m));
6113 q
6114 }
6115
6116 #[test]
6117 fn test_build_engine_fts_query_phrase() {
6118 let mut ns_phrase = lance_namespace::models::PhraseQuery::new("hello world".to_string());
6119 ns_phrase.column = Some("body".to_string());
6120 ns_phrase.slop = Some(2);
6121
6122 let mut ns_query = lance_namespace::models::FtsQuery::new();
6123 ns_query.phrase = Some(Box::new(ns_phrase));
6124
6125 match build_engine_fts_query(&ns_query).unwrap() {
6126 FtsQuery::Phrase(p) => {
6127 assert_eq!(p.terms, "hello world");
6128 assert_eq!(p.column, Some("body".to_string()));
6129 assert_eq!(p.slop, 2);
6130 }
6131 other => panic!("expected Phrase, got {:?}", other),
6132 }
6133 }
6134
6135 #[test]
6136 fn test_build_engine_fts_query_multi_match() {
6137 let mut m1 = lance_namespace::models::MatchQuery::new("hello".to_string());
6138 m1.column = Some("title".to_string());
6139 let mut m2 = lance_namespace::models::MatchQuery::new("hello".to_string());
6140 m2.column = Some("body".to_string());
6141 m2.boost = Some(2.0);
6142
6143 let ns_multi = lance_namespace::models::MultiMatchQuery::new(vec![m1, m2]);
6144 let mut ns_query = lance_namespace::models::FtsQuery::new();
6145 ns_query.multi_match = Some(Box::new(ns_multi));
6146
6147 match build_engine_fts_query(&ns_query).unwrap() {
6148 FtsQuery::MultiMatch(mm) => {
6149 assert_eq!(mm.match_queries.len(), 2);
6150 assert_eq!(mm.match_queries[0].terms, "hello");
6151 assert_eq!(mm.match_queries[0].column, Some("title".to_string()));
6152 assert_eq!(mm.match_queries[1].column, Some("body".to_string()));
6153 assert_eq!(mm.match_queries[1].boost, 2.0);
6154 }
6155 other => panic!("expected MultiMatch, got {:?}", other),
6156 }
6157 }
6158
6159 #[test]
6160 fn test_build_engine_fts_query_boolean() {
6161 let ns_boolean = lance_namespace::models::BooleanQuery::new(
6163 vec![ns_match_query("must-term", "body")],
6164 vec![ns_match_query("must-not-term", "body")],
6165 vec![ns_match_query("should-term", "body")],
6166 );
6167 let mut ns_query = lance_namespace::models::FtsQuery::new();
6168 ns_query.boolean = Some(Box::new(ns_boolean));
6169
6170 match build_engine_fts_query(&ns_query).unwrap() {
6171 FtsQuery::Boolean(b) => {
6172 assert!(matches!(&b.must[..], [FtsQuery::Match(m)] if m.terms == "must-term"));
6173 assert!(
6174 matches!(&b.must_not[..], [FtsQuery::Match(m)] if m.terms == "must-not-term")
6175 );
6176 assert!(matches!(&b.should[..], [FtsQuery::Match(m)] if m.terms == "should-term"));
6177 }
6178 other => panic!("expected Boolean, got {:?}", other),
6179 }
6180 }
6181
6182 #[test]
6183 fn test_build_engine_fts_query_boost() {
6184 let mut ns_boost = lance_namespace::models::BoostQuery::new(
6185 ns_match_query("positive-term", "body"),
6186 ns_match_query("negative-term", "body"),
6187 );
6188 ns_boost.negative_boost = Some(0.25);
6189
6190 let mut ns_query = lance_namespace::models::FtsQuery::new();
6191 ns_query.boost = Some(Box::new(ns_boost));
6192
6193 match build_engine_fts_query(&ns_query).unwrap() {
6194 FtsQuery::Boost(b) => {
6195 assert!(
6196 matches!(b.positive.as_ref(), FtsQuery::Match(m) if m.terms == "positive-term")
6197 );
6198 assert!(
6199 matches!(b.negative.as_ref(), FtsQuery::Match(m) if m.terms == "negative-term")
6200 );
6201 assert_eq!(b.negative_boost, 0.25);
6202 }
6203 other => panic!("expected Boost, got {:?}", other),
6204 }
6205 }
6206
6207 #[test]
6208 fn test_build_engine_fts_query_requires_a_variant() {
6209 let empty = lance_namespace::models::FtsQuery::new();
6211 assert!(build_engine_fts_query(&empty).is_err());
6212 }
6213 use lance::dataset::Dataset;
6214 use lance::index::DatasetIndexExt;
6215 use lance_core::utils::tempfile::{TempStdDir, TempStrDir};
6216 use lance_core::utils::testing::CountingObjectStore;
6217 use lance_io::object_store::{providers::local::FileStoreProvider, uri_to_url};
6218 use lance_namespace::error::ErrorCode;
6219 use lance_namespace::models::{
6220 CreateTableRequest, JsonArrowDataType, JsonArrowField, JsonArrowSchema, ListTablesRequest,
6221 QueryTableRequestColumns,
6222 };
6223 use lance_namespace::schema::convert_json_arrow_schema;
6224 use std::io::Cursor;
6225 use std::sync::{
6226 Arc,
6227 atomic::{AtomicUsize, Ordering},
6228 };
6229 use url::Url;
6230
6231 fn assert_plan_contains_all(plan: &str, expected_fragments: &[&str], context: &str) {
6232 for expected_fragment in expected_fragments {
6233 assert!(
6234 plan.contains(expected_fragment),
6235 "{}. Missing fragment: '{}'. Plan:\n{}",
6236 context,
6237 expected_fragment,
6238 plan
6239 );
6240 }
6241 }
6242
6243 fn mutation_error_code(err: lance_core::Error) -> ErrorCode {
6244 match err {
6245 lance_core::Error::Namespace { source, .. } => source
6246 .downcast_ref::<NamespaceError>()
6247 .expect("mutation error should wrap a NamespaceError")
6248 .code(),
6249 other => panic!("expected Namespace error, got: {other:?}"),
6250 }
6251 }
6252
6253 #[test]
6258 fn test_map_mutation_error_commit_conflict_alignment() {
6259 let boxed = || -> Box<dyn std::error::Error + Send + Sync + 'static> {
6260 Box::<dyn std::error::Error + Send + Sync>::from("inner conflict")
6261 };
6262
6263 let throttling_cases = vec![lance_core::Error::commit_conflict_source(1, boxed())];
6264 for err in throttling_cases {
6265 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6266 err,
6267 "update",
6268 "memory://t",
6269 ));
6270 assert_eq!(code, ErrorCode::Throttling);
6271 }
6272
6273 let concurrent_cases = vec![
6274 lance_core::Error::too_much_write_contention("contention"),
6275 lance_core::Error::retryable_commit_conflict_source(1, boxed()),
6276 lance_core::Error::incompatible_transaction_source(boxed()),
6277 lance_core::Error::version_conflict("conflict", 0, 3),
6278 ];
6279 for err in concurrent_cases {
6280 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6281 err,
6282 "update",
6283 "memory://t",
6284 ));
6285 assert_eq!(code, ErrorCode::ConcurrentModification);
6286 }
6287 }
6288
6289 async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) {
6291 let temp_dir = TempStdDir::default();
6292
6293 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
6294 .build()
6295 .await
6296 .unwrap();
6297 (namespace, temp_dir)
6298 }
6299
6300 #[derive(Debug)]
6301 #[allow(dead_code)]
6302 struct CountingFileStoreProvider {
6303 listing_count: Arc<AtomicUsize>,
6304 }
6305
6306 #[async_trait]
6307 impl lance_io::object_store::ObjectStoreProvider for CountingFileStoreProvider {
6308 async fn new_store(
6309 &self,
6310 base_path: Url,
6311 params: &ObjectStoreParams,
6312 ) -> Result<ObjectStore> {
6313 let provider = FileStoreProvider;
6314 let mut store = provider.new_store(base_path, params).await?;
6315 store.inner = Arc::new(CountingObjectStore::new(
6316 store.inner.clone(),
6317 self.listing_count.clone(),
6318 ));
6319 Ok(store)
6320 }
6321
6322 fn extract_path(&self, url: &Url) -> Result<Path> {
6323 let provider = FileStoreProvider;
6324 provider.extract_path(url)
6325 }
6326
6327 fn calculate_object_store_prefix(
6328 &self,
6329 url: &Url,
6330 storage_options: Option<&HashMap<String, String>>,
6331 ) -> Result<String> {
6332 let provider = FileStoreProvider;
6333 provider.calculate_object_store_prefix(url, storage_options)
6334 }
6335 }
6336
6337 #[allow(dead_code)]
6338 fn file_object_store_uri(path: &str) -> String {
6339 let file_url = uri_to_url(path).unwrap();
6340 let mut url = Url::parse("file-object-store:///").unwrap();
6341 url.set_path(file_url.path());
6342 url.to_string()
6343 }
6344
6345 #[allow(dead_code)]
6346 fn build_listing_counting_session(listing_count: Arc<AtomicUsize>) -> Arc<Session> {
6347 let registry = Arc::new(ObjectStoreRegistry::default());
6348 registry.insert(
6349 "file-object-store",
6350 Arc::new(CountingFileStoreProvider { listing_count }),
6351 );
6352 Arc::new(Session::new(0, 0, registry))
6353 }
6354
6355 use futures::stream::BoxStream;
6359 use object_store::{
6360 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
6361 PutMultipartOptions, PutPayload, PutResult, Result as OSResult,
6362 };
6363 use std::ops::Range;
6364
6365 #[derive(Debug, Clone, Copy)]
6366 enum ListBehavior {
6367 Throttle,
6368 ServiceUnavailable,
6369 Internal,
6370 NotFound,
6371 EmptyListing,
6372 }
6373
6374 #[derive(Debug)]
6375 struct FailingListStore {
6376 target: Arc<dyn OSObjectStore>,
6377 behavior: Arc<Mutex<Option<ListBehavior>>>,
6378 }
6379
6380 impl std::fmt::Display for FailingListStore {
6381 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6382 write!(f, "FailingListStore({})", self.target)
6383 }
6384 }
6385
6386 #[async_trait]
6387 impl OSObjectStore for FailingListStore {
6388 async fn put_opts(
6389 &self,
6390 location: &Path,
6391 bytes: PutPayload,
6392 opts: PutOptions,
6393 ) -> OSResult<PutResult> {
6394 self.target.put_opts(location, bytes, opts).await
6395 }
6396
6397 async fn put_multipart_opts(
6398 &self,
6399 location: &Path,
6400 opts: PutMultipartOptions,
6401 ) -> OSResult<Box<dyn MultipartUpload>> {
6402 self.target.put_multipart_opts(location, opts).await
6403 }
6404
6405 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
6406 self.target.get_opts(location, options).await
6407 }
6408
6409 async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> OSResult<Vec<Bytes>> {
6410 self.target.get_ranges(location, ranges).await
6411 }
6412
6413 fn delete_stream(
6414 &self,
6415 locations: BoxStream<'static, OSResult<Path>>,
6416 ) -> BoxStream<'static, OSResult<Path>> {
6417 self.target.delete_stream(locations)
6418 }
6419
6420 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
6421 self.target.list(prefix)
6422 }
6423
6424 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
6425 let behavior = *self.behavior.lock().unwrap();
6426 match behavior {
6427 None => self.target.list_with_delimiter(prefix).await,
6428 Some(ListBehavior::EmptyListing) => Ok(ListResult {
6429 common_prefixes: Vec::new(),
6430 objects: Vec::new(),
6431 }),
6432 Some(ListBehavior::Throttle) => Err(ObjectStoreError::Generic {
6435 store: "test",
6436 source: "Error performing list request: response error, after 3 retries, \
6437 max_retries: 3, retry_timeout: 180s - HTTP status server error \
6438 (503 Service Unavailable): ServerBusy: The server is busy"
6439 .into(),
6440 }),
6441 Some(ListBehavior::ServiceUnavailable) => Err(ObjectStoreError::Generic {
6442 store: "test",
6443 source: "Error performing list request: 503 Service Unavailable".into(),
6444 }),
6445 Some(ListBehavior::Internal) => Err(ObjectStoreError::Generic {
6446 store: "test",
6447 source: "Error performing list request: catastrophic unclassified failure"
6448 .into(),
6449 }),
6450 Some(ListBehavior::NotFound) => Err(ObjectStoreError::NotFound {
6451 path: "test_table.lance".to_string(),
6452 source: "entity not found".into(),
6453 }),
6454 }
6455 }
6456
6457 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
6458 self.target.copy_opts(from, to, opts).await
6459 }
6460 }
6461
6462 #[derive(Debug)]
6463 struct FailingListStoreProvider {
6464 behavior: Arc<Mutex<Option<ListBehavior>>>,
6465 }
6466
6467 #[async_trait]
6468 impl lance_io::object_store::ObjectStoreProvider for FailingListStoreProvider {
6469 async fn new_store(
6470 &self,
6471 base_path: Url,
6472 params: &ObjectStoreParams,
6473 ) -> Result<ObjectStore> {
6474 let mut store = FileStoreProvider.new_store(base_path, params).await?;
6475 store.inner = Arc::new(FailingListStore {
6476 target: store.inner.clone(),
6477 behavior: self.behavior.clone(),
6478 });
6479 Ok(store)
6480 }
6481
6482 fn extract_path(&self, url: &Url) -> Result<Path> {
6483 FileStoreProvider.extract_path(url)
6484 }
6485
6486 fn calculate_object_store_prefix(
6487 &self,
6488 url: &Url,
6489 storage_options: Option<&HashMap<String, String>>,
6490 ) -> Result<String> {
6491 FileStoreProvider.calculate_object_store_prefix(url, storage_options)
6492 }
6493 }
6494
6495 fn build_failing_list_session(behavior: Arc<Mutex<Option<ListBehavior>>>) -> Arc<Session> {
6496 let registry = Arc::new(ObjectStoreRegistry::default());
6497 registry.insert(
6498 "file-object-store",
6499 Arc::new(FailingListStoreProvider { behavior }),
6500 );
6501 Arc::new(Session::new(0, 0, registry))
6502 }
6503
6504 async fn failing_list_namespace() -> (
6508 DirectoryNamespace,
6509 TempStdDir,
6510 Arc<Mutex<Option<ListBehavior>>>,
6511 ) {
6512 let temp_dir = TempStdDir::default();
6513 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
6514 let behavior = Arc::new(Mutex::new(None));
6515 let session = build_failing_list_session(behavior.clone());
6516 let namespace = DirectoryNamespaceBuilder::new(root_uri)
6517 .session(session)
6518 .manifest_enabled(false)
6519 .dir_listing_enabled(true)
6520 .build()
6521 .await
6522 .unwrap();
6523 (namespace, temp_dir, behavior)
6524 }
6525
6526 async fn create_named_dir_table(namespace: &DirectoryNamespace, name: &str) {
6527 let schema = create_test_schema();
6528 let ipc_data = create_test_ipc_data(&schema);
6529 let mut create_req = CreateTableRequest::new();
6530 create_req.id = Some(vec![name.to_string()]);
6531 namespace
6532 .create_table(create_req, Bytes::from(ipc_data))
6533 .await
6534 .unwrap();
6535 }
6536
6537 #[tokio::test]
6542 async fn test_table_resolution_propagates_storage_errors_not_table_not_found() {
6543 for (behavior, expected_code, evidence) in [
6544 (ListBehavior::Throttle, ErrorCode::Throttling, "serverbusy"),
6545 (
6546 ListBehavior::ServiceUnavailable,
6547 ErrorCode::ServiceUnavailable,
6548 "503 service unavailable",
6549 ),
6550 (ListBehavior::Internal, ErrorCode::Internal, "catastrophic"),
6551 ] {
6552 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6553 create_named_dir_table(&namespace, "checkpoint").await;
6554 *toggle.lock().unwrap() = Some(behavior);
6555
6556 let mut describe_req = DescribeTableRequest::new();
6557 describe_req.id = Some(vec!["checkpoint".to_string()]);
6558 let err = namespace.describe_table(describe_req).await.unwrap_err();
6559 let msg = err.to_string();
6560 assert_eq!(
6561 mutation_error_code(err),
6562 expected_code,
6563 "describe_table under {behavior:?}; msg: {msg}"
6564 );
6565 assert!(
6566 msg.to_ascii_lowercase().contains(evidence),
6567 "describe_table message must carry storage evidence '{evidence}', got: {msg}"
6568 );
6569
6570 let mut exists_req = TableExistsRequest::new();
6571 exists_req.id = Some(vec!["checkpoint".to_string()]);
6572 let err = namespace.table_exists(exists_req).await.unwrap_err();
6573 let msg = err.to_string();
6574 assert_eq!(
6575 mutation_error_code(err),
6576 expected_code,
6577 "table_exists under {behavior:?}; msg: {msg}"
6578 );
6579 assert!(
6580 msg.to_ascii_lowercase().contains(evidence),
6581 "table_exists message must carry storage evidence '{evidence}', got: {msg}"
6582 );
6583 }
6584 }
6585
6586 #[tokio::test]
6589 async fn test_table_resolution_missing_table_yields_table_not_found() {
6590 for behavior in [ListBehavior::NotFound, ListBehavior::EmptyListing] {
6591 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6592 *toggle.lock().unwrap() = Some(behavior);
6593
6594 let mut describe_req = DescribeTableRequest::new();
6595 describe_req.id = Some(vec!["missing".to_string()]);
6596 let err = namespace.describe_table(describe_req).await.unwrap_err();
6597 assert_eq!(
6598 mutation_error_code(err),
6599 ErrorCode::TableNotFound,
6600 "describe_table under {behavior:?} should be TableNotFound"
6601 );
6602
6603 let mut exists_req = TableExistsRequest::new();
6604 exists_req.id = Some(vec!["missing".to_string()]);
6605 let err = namespace.table_exists(exists_req).await.unwrap_err();
6606 assert_eq!(
6607 mutation_error_code(err),
6608 ErrorCode::TableNotFound,
6609 "table_exists under {behavior:?} should be TableNotFound"
6610 );
6611 }
6612 }
6613
6614 #[tokio::test]
6624 async fn test_hybrid_resolution_falls_through_and_does_not_mask_throttle() {
6625 let temp_dir = TempStdDir::default();
6626 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
6627 let behavior = Arc::new(Mutex::new(None));
6628 let session = build_failing_list_session(behavior.clone());
6629
6630 let manifest_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
6632 .session(session.clone())
6633 .manifest_enabled(true)
6634 .dir_listing_enabled(true)
6635 .build()
6636 .await
6637 .unwrap();
6638 create_named_dir_table(&manifest_ns, "seed").await;
6639
6640 let dir_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
6642 .session(session.clone())
6643 .manifest_enabled(false)
6644 .dir_listing_enabled(true)
6645 .build()
6646 .await
6647 .unwrap();
6648 create_named_dir_table(&dir_ns, "checkpoint").await;
6649
6650 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
6653 .session(session)
6654 .manifest_enabled(true)
6655 .dir_listing_enabled(true)
6656 .dir_listing_to_manifest_migration_enabled(true)
6657 .build()
6658 .await
6659 .unwrap();
6660
6661 let mut describe_req = DescribeTableRequest::new();
6664 describe_req.id = Some(vec!["checkpoint".to_string()]);
6665 hybrid_ns
6666 .describe_table(describe_req)
6667 .await
6668 .expect("unregistered on-disk table should resolve via the manifest fall-through");
6669
6670 *behavior.lock().unwrap() = Some(ListBehavior::Throttle);
6674 let mut describe_req = DescribeTableRequest::new();
6675 describe_req.id = Some(vec!["checkpoint".to_string()]);
6676 let err = hybrid_ns.describe_table(describe_req).await.unwrap_err();
6677 let code = mutation_error_code(err);
6678 assert_ne!(
6679 code,
6680 ErrorCode::TableNotFound,
6681 "hybrid resolution masked a throttle as TableNotFound"
6682 );
6683 assert!(
6684 matches!(
6685 code,
6686 ErrorCode::Throttling | ErrorCode::ServiceUnavailable | ErrorCode::Internal
6687 ),
6688 "hybrid resolution should surface a storage error, got {code:?}"
6689 );
6690 }
6691
6692 #[test]
6693 fn test_classify_storage_error_maps_variants_and_preserves_evidence() {
6694 let throttle: Error = ObjectStoreError::Generic {
6695 store: "test",
6696 source: "list request failed, after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
6697 }
6698 .into();
6699 assert!(matches!(throttle, Error::IO { .. }));
6700 let classified = DirectoryNamespace::classify_storage_error(throttle);
6701 let msg = classified.to_string();
6702 assert_eq!(mutation_error_code(classified), ErrorCode::Throttling);
6703 assert!(
6704 msg.to_ascii_lowercase().contains("serverbusy"),
6705 "throttle evidence lost: {msg}"
6706 );
6707
6708 let service: Error = ObjectStoreError::Generic {
6709 store: "test",
6710 source: "504 Gateway Timeout".into(),
6711 }
6712 .into();
6713 assert_eq!(
6714 mutation_error_code(DirectoryNamespace::classify_storage_error(service)),
6715 ErrorCode::ServiceUnavailable
6716 );
6717
6718 let internal: Error = ObjectStoreError::Generic {
6719 store: "test",
6720 source: "disk caught fire".into(),
6721 }
6722 .into();
6723 assert_eq!(
6724 mutation_error_code(DirectoryNamespace::classify_storage_error(internal)),
6725 ErrorCode::Internal
6726 );
6727
6728 let preexisting: Error = NamespaceError::TableAlreadyExists {
6730 message: "t".to_string(),
6731 }
6732 .into();
6733 assert_eq!(
6734 mutation_error_code(DirectoryNamespace::classify_storage_error(preexisting)),
6735 ErrorCode::TableAlreadyExists
6736 );
6737 }
6738
6739 #[test]
6740 fn test_is_manifest_table_absent_error() {
6741 let table_not_found: Error = NamespaceError::TableNotFound {
6742 message: "t".to_string(),
6743 }
6744 .into();
6745 assert!(DirectoryNamespace::is_manifest_table_absent_error(
6746 &table_not_found
6747 ));
6748 let raw_not_found: Error = ObjectStoreError::NotFound {
6749 path: "t".to_string(),
6750 source: "x".into(),
6751 }
6752 .into();
6753 assert!(DirectoryNamespace::is_manifest_table_absent_error(
6754 &raw_not_found
6755 ));
6756
6757 let throttle: Error = ObjectStoreError::Generic {
6758 store: "test",
6759 source: "after 3 retries, max_retries: 3 ServerBusy".into(),
6760 }
6761 .into();
6762 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
6763 &throttle
6764 ));
6765 let internal: Error = NamespaceError::Internal {
6766 message: "boom".to_string(),
6767 }
6768 .into();
6769 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
6770 &internal
6771 ));
6772 }
6773
6774 #[test]
6775 fn test_map_open_error() {
6776 let not_found = || NamespaceError::TableNotFound {
6777 message: "table at 'x' not found: ...".to_string(),
6778 };
6779
6780 let throttle: Error = ObjectStoreError::Generic {
6781 store: "test",
6782 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
6783 }
6784 .into();
6785 assert_eq!(
6786 mutation_error_code(DirectoryNamespace::map_open_error(throttle, not_found())),
6787 ErrorCode::Throttling
6788 );
6789
6790 let generic_io: Error = ObjectStoreError::Generic {
6791 store: "test",
6792 source: "connection reset".into(),
6793 }
6794 .into();
6795 assert_eq!(
6796 mutation_error_code(DirectoryNamespace::map_open_error(generic_io, not_found())),
6797 ErrorCode::Internal
6798 );
6799
6800 let io_not_found: Error = ObjectStoreError::NotFound {
6801 path: "x".to_string(),
6802 source: "missing".into(),
6803 }
6804 .into();
6805 assert_eq!(
6806 mutation_error_code(DirectoryNamespace::map_open_error(
6807 io_not_found,
6808 not_found()
6809 )),
6810 ErrorCode::TableNotFound
6811 );
6812
6813 let dataset_not_found = Error::dataset_not_found("x".to_string(), "missing".into());
6814 assert_eq!(
6815 mutation_error_code(DirectoryNamespace::map_open_error(
6816 dataset_not_found,
6817 not_found()
6818 )),
6819 ErrorCode::TableNotFound
6820 );
6821
6822 let ref_not_found = Error::RefNotFound {
6824 message: "branch 'b' does not exist".to_string(),
6825 };
6826 assert_eq!(
6827 mutation_error_code(DirectoryNamespace::map_open_error(
6828 ref_not_found,
6829 not_found()
6830 )),
6831 ErrorCode::TableNotFound
6832 );
6833
6834 let version_miss = Error::RefNotFound {
6836 message: "version 5 does not exist".to_string(),
6837 };
6838 assert_eq!(
6839 mutation_error_code(DirectoryNamespace::map_open_error(
6840 version_miss,
6841 NamespaceError::TableVersionNotFound {
6842 message: "version 5 not found".to_string(),
6843 },
6844 )),
6845 ErrorCode::TableVersionNotFound
6846 );
6847 let version_throttle: Error = ObjectStoreError::Generic {
6848 store: "test",
6849 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
6850 }
6851 .into();
6852 assert_eq!(
6853 mutation_error_code(DirectoryNamespace::map_open_error(
6854 version_throttle,
6855 NamespaceError::TableVersionNotFound {
6856 message: "version 5 not found".to_string(),
6857 },
6858 )),
6859 ErrorCode::Throttling
6860 );
6861 }
6862
6863 fn create_test_ipc_data(schema: &JsonArrowSchema) -> Vec<u8> {
6865 use arrow::ipc::writer::StreamWriter;
6866
6867 let arrow_schema = convert_json_arrow_schema(schema).unwrap();
6868 let arrow_schema = Arc::new(arrow_schema);
6869 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
6870 let mut buffer = Vec::new();
6871 {
6872 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
6873 writer.write(&batch).unwrap();
6874 writer.finish().unwrap();
6875 }
6876 buffer
6877 }
6878
6879 fn create_ipc_data_from_batches(
6880 schema: Arc<arrow_schema::Schema>,
6881 batches: Vec<arrow::record_batch::RecordBatch>,
6882 ) -> Vec<u8> {
6883 use arrow::ipc::writer::StreamWriter;
6884
6885 let mut buffer = Vec::new();
6886 {
6887 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
6888 for batch in &batches {
6889 writer.write(batch).unwrap();
6890 }
6891 writer.finish().unwrap();
6892 }
6893 buffer
6894 }
6895
6896 fn create_non_empty_test_ipc_data() -> Vec<u8> {
6897 use arrow::array::{Int32Array, StringArray};
6898 use arrow::record_batch::RecordBatch;
6899
6900 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
6901 let batch = RecordBatch::try_new(
6902 schema.clone(),
6903 vec![
6904 Arc::new(Int32Array::from(vec![1, 2])),
6905 Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])),
6906 ],
6907 )
6908 .unwrap();
6909 create_ipc_data_from_batches(schema, vec![batch])
6910 }
6911
6912 fn create_single_row_test_ipc_data() -> Vec<u8> {
6913 use arrow::array::{Int32Array, StringArray};
6914 use arrow::record_batch::RecordBatch;
6915
6916 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
6917 let batch = RecordBatch::try_new(
6918 schema.clone(),
6919 vec![
6920 Arc::new(Int32Array::from(vec![10])),
6921 Arc::new(StringArray::from(vec![Some("carol")])),
6922 ],
6923 )
6924 .unwrap();
6925 create_ipc_data_from_batches(schema, vec![batch])
6926 }
6927
6928 fn create_test_schema() -> JsonArrowSchema {
6930 let int_type = JsonArrowDataType::new("int32".to_string());
6931 let string_type = JsonArrowDataType::new("utf8".to_string());
6932
6933 let id_field = JsonArrowField {
6934 name: "id".to_string(),
6935 r#type: Box::new(int_type),
6936 nullable: false,
6937 metadata: None,
6938 };
6939
6940 let name_field = JsonArrowField {
6941 name: "name".to_string(),
6942 r#type: Box::new(string_type),
6943 nullable: true,
6944 metadata: None,
6945 };
6946
6947 JsonArrowSchema {
6948 fields: vec![id_field, name_field],
6949 metadata: None,
6950 }
6951 }
6952
6953 fn create_scalar_table_ipc_data() -> Vec<u8> {
6954 use arrow::array::{Int32Array, StringArray};
6955 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6956
6957 let schema = Arc::new(ArrowSchema::new(vec![
6958 Field::new("id", DataType::Int32, false),
6959 Field::new("name", DataType::Utf8, true),
6960 ]));
6961 let batch = arrow::record_batch::RecordBatch::try_new(
6962 schema.clone(),
6963 vec![
6964 Arc::new(Int32Array::from(vec![1, 2, 3])),
6965 Arc::new(StringArray::from(vec!["alice", "bob", "cory"])),
6966 ],
6967 )
6968 .unwrap();
6969 create_ipc_data_from_batches(schema, vec![batch])
6970 }
6971
6972 async fn create_legacy_manifest_without_primary_key_metadata(root: &str) {
6973 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6974 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
6975
6976 let schema = Arc::new(ArrowSchema::new(vec![
6977 Field::new("object_id", DataType::Utf8, false),
6978 Field::new("object_type", DataType::Utf8, false),
6979 Field::new("location", DataType::Utf8, true),
6980 Field::new("metadata", DataType::Utf8, true),
6981 Field::new(
6982 "base_objects",
6983 DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))),
6984 true,
6985 ),
6986 ]));
6987 let batch = RecordBatch::new_empty(schema.clone());
6988 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
6989 Dataset::write(Box::new(reader), &format!("{}/__manifest", root), None)
6990 .await
6991 .unwrap();
6992 }
6993
6994 async fn manifest_has_primary_key_metadata(root: &str) -> bool {
6995 let dataset = Dataset::open(&format!("{}/__manifest", root))
6996 .await
6997 .unwrap();
6998 dataset
6999 .schema()
7000 .field("object_id")
7001 .map(|field| {
7002 field
7003 .metadata
7004 .contains_key(lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
7005 })
7006 .unwrap_or(false)
7007 }
7008
7009 fn create_vector_table_ipc_data() -> Vec<u8> {
7010 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
7011 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7012
7013 let schema = Arc::new(ArrowSchema::new(vec![
7014 Field::new("id", DataType::Int32, false),
7015 Field::new(
7016 "vector",
7017 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
7018 true,
7019 ),
7020 ]));
7021 let vector_field = Arc::new(Field::new("item", DataType::Float32, true));
7022 let vectors = FixedSizeListArray::try_new(
7023 vector_field,
7024 2,
7025 Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6])),
7026 None,
7027 )
7028 .unwrap();
7029 let batch = arrow::record_batch::RecordBatch::try_new(
7030 schema.clone(),
7031 vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(vectors)],
7032 )
7033 .unwrap();
7034 create_ipc_data_from_batches(schema, vec![batch])
7035 }
7036
7037 async fn create_scalar_table(namespace: &DirectoryNamespace, table_name: &str) {
7038 let mut create_table_request = CreateTableRequest::new();
7039 create_table_request.id = Some(vec![table_name.to_string()]);
7040 namespace
7041 .create_table(
7042 create_table_request,
7043 Bytes::from(create_scalar_table_ipc_data()),
7044 )
7045 .await
7046 .unwrap();
7047 }
7048
7049 async fn create_vector_table(namespace: &DirectoryNamespace, table_name: &str) {
7050 let mut create_table_request = CreateTableRequest::new();
7051 create_table_request.id = Some(vec![table_name.to_string()]);
7052 namespace
7053 .create_table(
7054 create_table_request,
7055 Bytes::from(create_vector_table_ipc_data()),
7056 )
7057 .await
7058 .unwrap();
7059 }
7060
7061 async fn open_dataset(namespace: &DirectoryNamespace, table_name: &str) -> Dataset {
7062 let mut describe_request = DescribeTableRequest::new();
7063 describe_request.id = Some(vec![table_name.to_string()]);
7064 let table_uri = namespace
7065 .describe_table(describe_request)
7066 .await
7067 .unwrap()
7068 .location
7069 .expect("table location should exist");
7070 Dataset::open(&table_uri).await.unwrap()
7071 }
7072
7073 async fn create_scalar_index(
7074 namespace: &DirectoryNamespace,
7075 table_name: &str,
7076 index_name: &str,
7077 ) -> Option<String> {
7078 use lance_namespace::models::CreateTableIndexRequest;
7079
7080 let mut create_index_request =
7081 CreateTableIndexRequest::new("id".to_string(), "BTREE".to_string());
7082 create_index_request.id = Some(vec![table_name.to_string()]);
7083 create_index_request.name = Some(index_name.to_string());
7084 namespace
7085 .create_table_scalar_index(create_index_request)
7086 .await
7087 .unwrap()
7088 .transaction_id
7089 }
7090
7091 async fn create_branch_with_commits(
7096 namespace: &DirectoryNamespace,
7097 table_name: &str,
7098 branch_name: &str,
7099 extra_versions: usize,
7100 ) -> String {
7101 let mut main = open_dataset(namespace, table_name).await;
7102 let fork_version = main.version().version;
7103 let branch = main
7104 .create_branch(branch_name, fork_version, None)
7105 .await
7106 .unwrap();
7107 let branch_uri = branch.uri().to_string();
7108 for i in 0..extra_versions {
7109 append_scalar_version(&branch_uri, (i as i32 + 1) * 100).await;
7110 }
7111 branch_uri
7112 }
7113
7114 async fn append_scalar_version(uri: &str, seed: i32) {
7117 use arrow::array::{Int32Array, StringArray};
7118 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7119 let schema = Arc::new(ArrowSchema::new(vec![
7120 Field::new("id", DataType::Int32, false),
7121 Field::new("name", DataType::Utf8, true),
7122 ]));
7123 let batch = arrow::record_batch::RecordBatch::try_new(
7124 schema.clone(),
7125 vec![
7126 Arc::new(Int32Array::from(vec![seed, seed + 1])),
7127 Arc::new(StringArray::from(vec![Some("x"), Some("y")])),
7128 ],
7129 )
7130 .unwrap();
7131 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
7132 Dataset::write(
7133 reader,
7134 uri,
7135 Some(WriteParams {
7136 mode: WriteMode::Append,
7137 ..Default::default()
7138 }),
7139 )
7140 .await
7141 .unwrap();
7142 }
7143
7144 async fn list_versions(
7146 namespace: &DirectoryNamespace,
7147 table_name: &str,
7148 branch: Option<&str>,
7149 ) -> Result<Vec<TableVersion>> {
7150 let req = ListTableVersionsRequest {
7151 id: Some(vec![table_name.to_string()]),
7152 branch: branch.map(|b| b.to_string()),
7153 ..Default::default()
7154 };
7155 namespace.list_table_versions(req).await.map(|r| r.versions)
7156 }
7157
7158 #[tokio::test]
7159 async fn test_list_table_versions_on_branch() {
7160 let (namespace, _temp_dir) = create_test_namespace().await;
7161 create_scalar_table(&namespace, "users").await;
7162 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7163
7164 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7167 .await
7168 .unwrap();
7169 assert!(branch_versions.len() >= 2);
7170 assert!(
7171 branch_versions
7172 .iter()
7173 .all(|v| v.manifest_path.contains("tree/exp")),
7174 "branch versions must resolve to branch manifests: {:?}",
7175 branch_versions
7176 );
7177
7178 let main_versions = list_versions(&namespace, "users", None).await.unwrap();
7180 let main_explicit = list_versions(&namespace, "users", Some("main"))
7181 .await
7182 .unwrap();
7183 assert_eq!(main_versions.len(), main_explicit.len());
7184 assert!(
7185 main_versions
7186 .iter()
7187 .all(|v| !v.manifest_path.contains("tree/"))
7188 );
7189
7190 let missing = list_versions(&namespace, "users", Some("does-not-exist")).await;
7192 assert!(missing.is_err());
7193 assert!(missing.unwrap_err().to_string().contains("not found"));
7194 }
7195
7196 #[tokio::test]
7197 async fn test_describe_table_version_on_branch() {
7198 let (namespace, _temp_dir) = create_test_namespace().await;
7199 create_scalar_table(&namespace, "users").await;
7200 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7201
7202 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7203 .await
7204 .unwrap();
7205 let latest = branch_versions.iter().map(|v| v.version).max().unwrap();
7206
7207 let req = DescribeTableVersionRequest {
7209 id: Some(vec!["users".to_string()]),
7210 branch: Some("exp".to_string()),
7211 ..Default::default()
7212 };
7213 let resp = namespace.describe_table_version(req).await.unwrap();
7214 assert_eq!(resp.version.version, latest);
7215 assert!(resp.version.manifest_path.contains("tree/exp"));
7216
7217 let req = DescribeTableVersionRequest {
7219 id: Some(vec!["users".to_string()]),
7220 version: Some(latest),
7221 branch: Some("exp".to_string()),
7222 ..Default::default()
7223 };
7224 assert!(namespace.describe_table_version(req).await.is_ok());
7225
7226 let req = DescribeTableVersionRequest {
7228 id: Some(vec!["users".to_string()]),
7229 version: Some(999_999),
7230 branch: Some("exp".to_string()),
7231 ..Default::default()
7232 };
7233 assert!(namespace.describe_table_version(req).await.is_err());
7234
7235 let req = DescribeTableVersionRequest {
7237 id: Some(vec!["users".to_string()]),
7238 branch: Some("nope".to_string()),
7239 ..Default::default()
7240 };
7241 let err = namespace.describe_table_version(req).await;
7242 assert!(err.is_err() && err.unwrap_err().to_string().contains("not found"));
7243 }
7244
7245 #[tokio::test]
7246 async fn test_restore_table_on_branch() {
7247 use lance_namespace::models::RestoreTableRequest;
7248
7249 let (namespace, _temp_dir) = create_test_namespace().await;
7250 create_scalar_table(&namespace, "users").await;
7251 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7252
7253 let before = list_versions(&namespace, "users", Some("exp"))
7254 .await
7255 .unwrap();
7256 let branch_latest = before.iter().map(|v| v.version).max().unwrap();
7257 let earliest = before.iter().map(|v| v.version).min().unwrap();
7258 let main_before = list_versions(&namespace, "users", None)
7259 .await
7260 .unwrap()
7261 .len();
7262
7263 let req = RestoreTableRequest {
7266 id: Some(vec!["users".to_string()]),
7267 version: earliest,
7268 branch: Some("exp".to_string()),
7269 ..Default::default()
7270 };
7271 let resp = namespace.restore_table(req).await.unwrap();
7272 assert!(resp.transaction_id.is_some());
7273
7274 let after = list_versions(&namespace, "users", Some("exp"))
7275 .await
7276 .unwrap();
7277 let new_latest = after.iter().map(|v| v.version).max().unwrap();
7278 assert!(
7279 new_latest > branch_latest,
7280 "restore should add a branch version"
7281 );
7282
7283 let main_after = list_versions(&namespace, "users", None)
7284 .await
7285 .unwrap()
7286 .len();
7287 assert_eq!(main_after, main_before, "main must be unaffected");
7288 }
7289
7290 #[tokio::test]
7291 async fn test_batch_delete_table_versions_on_branch() {
7292 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7293
7294 let (namespace, _temp_dir) = create_test_namespace().await;
7295 create_scalar_table(&namespace, "users").await;
7296 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7297
7298 let before = list_versions(&namespace, "users", Some("exp"))
7299 .await
7300 .unwrap();
7301 let main_before = list_versions(&namespace, "users", None).await.unwrap();
7302
7303 let req = BatchDeleteTableVersionsRequest {
7308 id: Some(vec!["users".to_string()]),
7309 branch: Some("exp".to_string()),
7310 ranges: vec![VersionRange::new(0, -1)],
7311 ..Default::default()
7312 };
7313 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7314 assert_eq!(
7315 resp.deleted_count,
7316 Some(before.len() as i64),
7317 "every branch manifest should be physically deleted"
7318 );
7319
7320 assert!(
7322 list_versions(&namespace, "users", Some("exp"))
7323 .await
7324 .is_err()
7325 );
7326 let main_after = list_versions(&namespace, "users", None).await.unwrap();
7327 assert_eq!(
7328 main_after.len(),
7329 main_before.len(),
7330 "main must be untouched"
7331 );
7332 }
7333
7334 #[tokio::test]
7335 async fn test_create_table_version_on_branch() {
7336 use futures::TryStreamExt;
7337 use lance_namespace::models::CreateTableVersionRequest;
7338
7339 let (namespace, _temp_dir) = create_test_namespace().await;
7340 create_scalar_table(&namespace, "users").await;
7341 let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 1).await;
7342
7343 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7345 let versions_dir = branch_ds.versions_dir();
7346 let store = branch_ds.object_store(None).await.unwrap();
7347 let existing = store
7348 .inner
7349 .list(Some(&versions_dir))
7350 .try_collect::<Vec<_>>()
7351 .await
7352 .unwrap()
7353 .into_iter()
7354 .find(|m| {
7355 m.location
7356 .filename()
7357 .map(|f| f.ends_with(".manifest"))
7358 .unwrap_or(false)
7359 })
7360 .expect("a branch manifest");
7361 let bytes = store
7362 .inner
7363 .get(&existing.location)
7364 .await
7365 .unwrap()
7366 .bytes()
7367 .await
7368 .unwrap();
7369 let staging = versions_dir.join("staging_manifest");
7370 store.inner.put(&staging, bytes.into()).await.unwrap();
7371
7372 let main_before = list_versions(&namespace, "users", None)
7373 .await
7374 .unwrap()
7375 .len();
7376 let new_version = list_versions(&namespace, "users", Some("exp"))
7377 .await
7378 .unwrap()
7379 .iter()
7380 .map(|v| v.version)
7381 .max()
7382 .unwrap()
7383 + 1;
7384
7385 let req = CreateTableVersionRequest {
7386 id: Some(vec!["users".to_string()]),
7387 version: new_version,
7388 manifest_path: staging.to_string(),
7389 naming_scheme: Some("V2".to_string()),
7390 branch: Some("exp".to_string()),
7391 ..Default::default()
7392 };
7393 let resp = namespace.create_table_version(req).await.unwrap();
7394 let info = resp.version.expect("version info");
7395 assert!(
7397 info.manifest_path.contains("tree/exp"),
7398 "got {}",
7399 info.manifest_path
7400 );
7401
7402 let after = list_versions(&namespace, "users", Some("exp"))
7404 .await
7405 .unwrap();
7406 assert!(after.iter().any(|v| v.version == new_version));
7407 let main_after = list_versions(&namespace, "users", None)
7408 .await
7409 .unwrap()
7410 .len();
7411 assert_eq!(main_after, main_before, "main must be unaffected");
7412 }
7413
7414 #[tokio::test]
7419 async fn test_external_manifest_store_resolves_branch_from_base_path() {
7420 use futures::TryStreamExt;
7421 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
7422 use lance_table::io::commit::external_manifest::ExternalManifestStore;
7423
7424 let (namespace, _temp_dir) = create_test_namespace().await;
7425 create_scalar_table(&namespace, "users").await; let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 2).await;
7427
7428 let namespace = Arc::new(namespace);
7429 let table_id = vec!["users".to_string()];
7430 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7431 let branch_base = branch_ds.branch_location().path;
7432 let root_base = branch_ds.branch_location().find_main().unwrap().path;
7433 let store = LanceNamespaceExternalManifestStore::new(
7434 namespace.clone(),
7435 table_id.clone(),
7436 root_base.clone(),
7437 );
7438
7439 let (branch_latest, branch_path) = store
7442 .get_latest_version(branch_base.as_ref())
7443 .await
7444 .unwrap()
7445 .expect("branch has versions");
7446 let (_main_latest, main_path) = store
7447 .get_latest_version(root_base.as_ref())
7448 .await
7449 .unwrap()
7450 .expect("main has versions");
7451 assert!(
7452 branch_path.contains("tree/exp"),
7453 "branch latest must resolve to the branch tree: {}",
7454 branch_path
7455 );
7456 assert!(
7457 !main_path.contains("tree/exp"),
7458 "main latest must not resolve to a branch tree: {}",
7459 main_path
7460 );
7461
7462 let described = store
7464 .get(branch_base.as_ref(), branch_latest)
7465 .await
7466 .unwrap();
7467 assert!(
7468 described.contains("tree/exp"),
7469 "describe on the branch must resolve to the branch tree: {}",
7470 described
7471 );
7472
7473 assert!(store.get_latest_version("somewhere/else").await.is_err());
7475
7476 let versions_dir = branch_ds.versions_dir();
7479 let obj = branch_ds.object_store(None).await.unwrap();
7480 let existing = obj
7481 .inner
7482 .list(Some(&versions_dir))
7483 .try_collect::<Vec<_>>()
7484 .await
7485 .unwrap()
7486 .into_iter()
7487 .find(|m| {
7488 m.location
7489 .filename()
7490 .map(|f| f.ends_with(".manifest"))
7491 .unwrap_or(false)
7492 })
7493 .expect("a branch manifest");
7494 let bytes = obj
7495 .inner
7496 .get(&existing.location)
7497 .await
7498 .unwrap()
7499 .bytes()
7500 .await
7501 .unwrap();
7502 let size = bytes.len() as u64;
7503 let staging = versions_dir.clone().join("staging_manifest");
7504 obj.inner.put(&staging, bytes.into()).await.unwrap();
7505
7506 let committed = store
7507 .put(
7508 &branch_base,
7509 branch_latest + 1,
7510 &staging,
7511 size,
7512 None,
7513 obj.inner.as_ref(),
7514 ManifestNamingScheme::V2,
7515 )
7516 .await
7517 .unwrap();
7518 assert!(
7519 committed.path.to_string().contains("tree/exp"),
7520 "a commit through a branch-qualified base must land on the branch tree: {}",
7521 committed.path
7522 );
7523 }
7524
7525 #[tokio::test]
7533 async fn test_write_into_namespace_on_branch_appends_to_branch() {
7534 use lance::dataset::builder::DatasetBuilder;
7535 use lance_namespace::models::CreateTableBranchRequest;
7536
7537 let temp = TempStdDir::default();
7538 let namespace = Arc::new(
7539 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
7540 .manifest_enabled(true)
7541 .table_version_tracking_enabled(true)
7542 .ops_metrics_enabled(true)
7543 .build()
7544 .await
7545 .unwrap(),
7546 );
7547 let ns: Arc<dyn LanceNamespace> = namespace.clone();
7548 let table_id = vec!["t".to_string()];
7549 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
7551 id: Some(table_id.clone()),
7552 name: "exp".to_string(),
7553 ..Default::default()
7554 })
7555 .await
7556 .unwrap();
7557
7558 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
7559 ns.list_table_versions(ListTableVersionsRequest {
7560 id: Some(table_id),
7561 ..Default::default()
7562 })
7563 .await
7564 .unwrap()
7565 .versions
7566 .len()
7567 };
7568 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
7569 let commits_before = namespace
7570 .retrieve_ops_metrics()
7571 .get("create_table_version")
7572 .copied()
7573 .unwrap_or(0);
7574
7575 let branch_ds = Dataset::write_into_namespace_on_branch(
7576 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
7577 ns.clone(),
7578 table_id.clone(),
7579 "exp",
7580 Some(WriteParams {
7581 mode: WriteMode::Append,
7582 ..Default::default()
7583 }),
7584 )
7585 .await
7586 .unwrap();
7587 assert_eq!(branch_ds.manifest.branch.as_deref(), Some("exp"));
7588 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
7589
7590 let commits_after = namespace
7593 .retrieve_ops_metrics()
7594 .get("create_table_version")
7595 .copied()
7596 .unwrap_or(0);
7597 assert_eq!(
7598 commits_after,
7599 commits_before + 1,
7600 "the branch append must register its version via create_table_version"
7601 );
7602 let exp_versions = ns
7603 .list_table_versions(ListTableVersionsRequest {
7604 id: Some(table_id.clone()),
7605 branch: Some("exp".to_string()),
7606 ..Default::default()
7607 })
7608 .await
7609 .unwrap()
7610 .versions;
7611 assert!(
7612 exp_versions
7613 .iter()
7614 .all(|v| v.manifest_path.contains("tree/exp")),
7615 "branch versions must resolve to the branch tree: {:?}",
7616 exp_versions
7617 );
7618 assert_eq!(
7619 main_chain_len(ns.clone(), table_id.clone()).await,
7620 main_before,
7621 "main's catalog must be untouched by the branch append"
7622 );
7623
7624 Dataset::write_into_namespace(
7627 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
7628 ns.clone(),
7629 table_id.clone(),
7630 Some(WriteParams {
7631 mode: WriteMode::Append,
7632 ..Default::default()
7633 }),
7634 )
7635 .await
7636 .unwrap();
7637 assert_eq!(
7638 main_chain_len(ns.clone(), table_id.clone()).await,
7639 main_before + 1,
7640 "a managed main append must register its version in the catalog"
7641 );
7642 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
7643 .await
7644 .unwrap()
7645 .load()
7646 .await
7647 .unwrap();
7648 assert_eq!(
7649 scan_id_column(&fresh).await,
7650 vec![1, 2, 100],
7651 "a fresh managed open must resolve the appended version, not a stale latest"
7652 );
7653 }
7654
7655 #[tokio::test]
7657 async fn test_write_into_namespace_on_branch_rejects_create() {
7658 use arrow::array::{Int32Array, StringArray};
7659 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7660
7661 let (namespace, _temp_dir) = create_test_namespace().await;
7662 let namespace = Arc::new(namespace);
7663
7664 let schema = Arc::new(ArrowSchema::new(vec![
7665 Field::new("id", DataType::Int32, false),
7666 Field::new("name", DataType::Utf8, true),
7667 ]));
7668 let batch = arrow::record_batch::RecordBatch::try_new(
7669 schema.clone(),
7670 vec![
7671 Arc::new(Int32Array::from(vec![1])),
7672 Arc::new(StringArray::from(vec![Some("a")])),
7673 ],
7674 )
7675 .unwrap();
7676 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
7677
7678 let result = Dataset::write_into_namespace_on_branch(
7679 reader,
7680 namespace.clone(),
7681 vec!["new_table".to_string()],
7682 "exp",
7683 Some(WriteParams {
7684 mode: WriteMode::Create,
7685 ..Default::default()
7686 }),
7687 )
7688 .await;
7689 assert!(result.is_err(), "create on a branch must be rejected");
7690 assert!(
7691 result.unwrap_err().to_string().contains("branch"),
7692 "error should mention the branch restriction"
7693 );
7694 }
7695
7696 #[tokio::test]
7697 async fn test_branch_name_validation_rejects_traversal() {
7698 let (namespace, _temp_dir) = create_test_namespace().await;
7699 create_scalar_table(&namespace, "users").await;
7700
7701 let err = list_versions(&namespace, "users", Some("../evil")).await;
7704 assert!(err.is_err());
7705 assert!(err.unwrap_err().to_string().contains("invalid branch name"));
7706 }
7707
7708 #[tokio::test]
7709 async fn test_branch_ops_reject_zombie_branch() {
7710 use futures::TryStreamExt;
7711 use lance_namespace::models::{
7712 BatchDeleteTableVersionsRequest, CreateTableVersionRequest, RestoreTableRequest,
7713 VersionRange,
7714 };
7715
7716 let (namespace, _temp_dir) = create_test_namespace().await;
7717 create_scalar_table(&namespace, "users").await;
7718
7719 let dataset = open_dataset(&namespace, "users").await;
7720 let store = dataset.object_store(None).await.unwrap();
7721 let manifest = store
7722 .inner
7723 .list(Some(&dataset.versions_dir()))
7724 .try_collect::<Vec<_>>()
7725 .await
7726 .unwrap()
7727 .into_iter()
7728 .find(|m| {
7729 m.location
7730 .filename()
7731 .map(|f| f.ends_with(".manifest"))
7732 .unwrap_or(false)
7733 })
7734 .expect("a manifest");
7735 let bytes = store
7736 .inner
7737 .get(&manifest.location)
7738 .await
7739 .unwrap()
7740 .bytes()
7741 .await
7742 .unwrap();
7743 let zombie = dataset
7744 .branch_location()
7745 .find_branch(Some("ghost"))
7746 .unwrap()
7747 .path
7748 .join(VERSIONS_DIR)
7749 .join(manifest.location.filename().unwrap());
7750 store.inner.put(&zombie, bytes.into()).await.unwrap();
7751
7752 assert!(dataset.branches().get("ghost").await.is_err());
7753
7754 fn rejected<T: std::fmt::Debug>(label: &str, r: Result<T>) {
7755 match r {
7756 Ok(v) => panic!("{label} must reject the zombie branch, got Ok({v:?})"),
7757 Err(e) => assert!(e.to_string().contains("not found"), "{label}: {e}"),
7758 }
7759 }
7760
7761 rejected(
7762 "list",
7763 list_versions(&namespace, "users", Some("ghost")).await,
7764 );
7765 rejected(
7766 "describe",
7767 namespace
7768 .describe_table_version(DescribeTableVersionRequest {
7769 id: Some(vec!["users".to_string()]),
7770 branch: Some("ghost".to_string()),
7771 ..Default::default()
7772 })
7773 .await,
7774 );
7775 rejected(
7776 "create",
7777 namespace
7778 .create_table_version(CreateTableVersionRequest {
7779 id: Some(vec!["users".to_string()]),
7780 version: 2,
7781 manifest_path: zombie.to_string(),
7782 branch: Some("ghost".to_string()),
7783 ..Default::default()
7784 })
7785 .await,
7786 );
7787 rejected(
7788 "restore",
7789 namespace
7790 .restore_table(RestoreTableRequest {
7791 id: Some(vec!["users".to_string()]),
7792 version: 1,
7793 branch: Some("ghost".to_string()),
7794 ..Default::default()
7795 })
7796 .await,
7797 );
7798 rejected(
7799 "batch_delete",
7800 namespace
7801 .batch_delete_table_versions(BatchDeleteTableVersionsRequest {
7802 id: Some(vec!["users".to_string()]),
7803 branch: Some("ghost".to_string()),
7804 ranges: vec![VersionRange::new(1, 1)],
7805 ..Default::default()
7806 })
7807 .await,
7808 );
7809 }
7810
7811 #[tokio::test]
7816 async fn test_batch_delete_table_versions_main_v2() {
7817 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7818
7819 let (namespace, _temp_dir) = create_test_namespace().await;
7820 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
7822 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
7826 assert!(before.len() >= 3);
7827 assert!(
7830 before
7831 .iter()
7832 .all(|v| v.manifest_path.rsplit('/').next().unwrap().len() == 29),
7833 "expected V2-named manifests: {:?}",
7834 before
7835 );
7836 let min_v = before.iter().map(|v| v.version).min().unwrap();
7837 let max_v = before.iter().map(|v| v.version).max().unwrap();
7838
7839 let req = BatchDeleteTableVersionsRequest {
7842 id: Some(vec!["users".to_string()]),
7843 ranges: vec![VersionRange::new(min_v, max_v)],
7844 ..Default::default()
7845 };
7846 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7847 assert_eq!(
7848 resp.deleted_count,
7849 Some((before.len() - 1) as i64),
7850 "V2 manifests must actually be deleted (was 0 before the fix)"
7851 );
7852
7853 let after = list_versions(&namespace, "users", None).await.unwrap();
7854 assert_eq!(after.len(), 1);
7855 assert_eq!(after[0].version, max_v);
7856 }
7857
7858 #[tokio::test]
7860 async fn test_batch_delete_end_is_exclusive() {
7861 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7862
7863 let (namespace, _temp_dir) = create_test_namespace().await;
7864 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
7866 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
7870 let min_v = before.iter().map(|v| v.version).min().unwrap();
7871
7872 let req = BatchDeleteTableVersionsRequest {
7873 id: Some(vec!["users".to_string()]),
7874 ranges: vec![VersionRange::new(min_v, min_v + 1)],
7875 ..Default::default()
7876 };
7877 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7878 assert_eq!(
7879 resp.deleted_count,
7880 Some(1),
7881 "only min_v is in [min_v, min_v+1)"
7882 );
7883
7884 let after = list_versions(&namespace, "users", None).await.unwrap();
7885 assert!(
7886 !after.iter().any(|v| v.version == min_v),
7887 "min_v must be deleted"
7888 );
7889 assert_eq!(after.len(), before.len() - 1, "exactly one version removed");
7890 }
7891
7892 #[tokio::test]
7893 async fn test_batch_delete_rejects_unbounded_range() {
7894 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7895
7896 let (namespace, _temp_dir) = create_test_namespace().await;
7897 create_scalar_table(&namespace, "users").await;
7898
7899 let req = BatchDeleteTableVersionsRequest {
7902 id: Some(vec!["users".to_string()]),
7903 ranges: vec![VersionRange::new(0, i64::MAX)],
7904 ..Default::default()
7905 };
7906 let err = namespace.batch_delete_table_versions(req).await;
7907 assert!(err.is_err());
7908 assert!(
7909 err.unwrap_err().to_string().contains("limit"),
7910 "expected a range-too-large error"
7911 );
7912 }
7913
7914 async fn create_managed_namespace(path: &str) -> Arc<dyn LanceNamespace> {
7916 Arc::new(
7917 DirectoryNamespaceBuilder::new(path)
7918 .manifest_enabled(true)
7919 .table_version_tracking_enabled(true)
7920 .build()
7921 .await
7922 .unwrap(),
7923 )
7924 }
7925
7926 fn single_int_schema() -> Arc<arrow::datatypes::Schema> {
7927 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7928 Arc::new(ArrowSchema::new(vec![Field::new(
7929 "id",
7930 DataType::Int32,
7931 false,
7932 )]))
7933 }
7934
7935 fn single_int_batch(seed: i32) -> arrow::record_batch::RecordBatch {
7936 use arrow::array::Int32Array;
7937 arrow::record_batch::RecordBatch::try_new(
7938 single_int_schema(),
7939 vec![Arc::new(Int32Array::from(vec![seed]))],
7940 )
7941 .unwrap()
7942 }
7943
7944 async fn create_managed_table(ns: &Arc<dyn LanceNamespace>, table_id: &[String]) -> Dataset {
7947 let mut ds = Dataset::write_into_namespace(
7948 RecordBatchIterator::new(vec![Ok(single_int_batch(1))], single_int_schema()),
7949 ns.clone(),
7950 table_id.to_vec(),
7951 Some(WriteParams {
7952 mode: WriteMode::Create,
7953 ..Default::default()
7954 }),
7955 )
7956 .await
7957 .unwrap();
7958 ds.append(
7959 RecordBatchIterator::new(vec![Ok(single_int_batch(2))], single_int_schema()),
7960 None,
7961 )
7962 .await
7963 .unwrap();
7964 ds
7965 }
7966
7967 async fn scan_id_column(ds: &Dataset) -> Vec<i32> {
7969 use arrow::array::Int32Array;
7970 use futures::TryStreamExt;
7971 let batches: Vec<arrow::record_batch::RecordBatch> = ds
7972 .scan()
7973 .try_into_stream()
7974 .await
7975 .unwrap()
7976 .try_collect()
7977 .await
7978 .unwrap();
7979 let mut ids: Vec<i32> = batches
7980 .iter()
7981 .flat_map(|b| {
7982 b.column(0)
7983 .as_any()
7984 .downcast_ref::<Int32Array>()
7985 .unwrap()
7986 .values()
7987 .to_vec()
7988 })
7989 .collect();
7990 ids.sort();
7991 ids
7992 }
7993
7994 #[tokio::test]
7999 async fn test_managed_branch_open_and_commit() {
8000 use futures::TryStreamExt;
8001 use lance::dataset::builder::DatasetBuilder;
8002 use lance_namespace::models::CreateTableBranchRequest;
8003
8004 let temp = TempStdDir::default();
8005 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8006 let table_id = vec!["t".to_string()];
8007 create_managed_table(&ns, &table_id).await;
8008 let main_before = ns
8009 .list_table_versions(ListTableVersionsRequest {
8010 id: Some(table_id.clone()),
8011 ..Default::default()
8012 })
8013 .await
8014 .unwrap()
8015 .versions
8016 .len();
8017
8018 ns.create_table_branch(CreateTableBranchRequest {
8021 id: Some(table_id.clone()),
8022 name: "exp".to_string(),
8023 ..Default::default()
8024 })
8025 .await
8026 .unwrap();
8027
8028 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8031 .await
8032 .unwrap()
8033 .with_branch("exp", None)
8034 .load()
8035 .await
8036 .unwrap();
8037 assert_eq!(
8038 branch_ds.manifest.branch.as_deref(),
8039 Some("exp"),
8040 "with_branch on a managed table must open the branch chain"
8041 );
8042 let branch_base = branch_ds.branch_location().path;
8043 assert!(
8044 branch_base.as_ref().ends_with("tree/exp"),
8045 "the branch dataset must be rooted at the branch chain: {}",
8046 branch_base
8047 );
8048 let branch_v_before = branch_ds.version().version;
8049
8050 branch_ds
8052 .append(
8053 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8054 None,
8055 )
8056 .await
8057 .unwrap();
8058 assert_eq!(
8059 branch_ds.manifest.branch.as_deref(),
8060 Some("exp"),
8061 "the commit must stay on the branch"
8062 );
8063 assert!(
8064 branch_ds.version().version > branch_v_before,
8065 "the branch version must advance after the commit"
8066 );
8067 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
8068
8069 let store = branch_ds.object_store(None).await.unwrap();
8073 let branch_data = branch_base.clone().join("data");
8074 let branch_files = store
8075 .inner
8076 .list(Some(&branch_data))
8077 .try_collect::<Vec<_>>()
8078 .await
8079 .unwrap();
8080 assert!(
8081 !branch_files.is_empty(),
8082 "the branch commit must place data files under the branch chain"
8083 );
8084
8085 let table_uri = ns
8087 .describe_table(DescribeTableRequest {
8088 id: Some(table_id.clone()),
8089 ..Default::default()
8090 })
8091 .await
8092 .unwrap()
8093 .location
8094 .unwrap();
8095 let fs_branch_ds = DatasetBuilder::from_uri(&table_uri)
8096 .with_branch("exp", None)
8097 .load()
8098 .await
8099 .unwrap();
8100 assert_eq!(fs_branch_ds.manifest.branch.as_deref(), Some("exp"));
8101 assert_eq!(scan_id_column(&fs_branch_ds).await, vec![1, 2, 3]);
8102
8103 let main_after = ns
8106 .list_table_versions(ListTableVersionsRequest {
8107 id: Some(table_id.clone()),
8108 ..Default::default()
8109 })
8110 .await
8111 .unwrap()
8112 .versions
8113 .len();
8114 assert_eq!(
8115 main_after, main_before,
8116 "committing on the branch must not change main's chain"
8117 );
8118 let main_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8119 .await
8120 .unwrap()
8121 .load()
8122 .await
8123 .unwrap();
8124 assert_eq!(main_ds.manifest.branch, None);
8125 assert_eq!(scan_id_column(&main_ds).await, vec![1, 2]);
8126 }
8127
8128 #[tokio::test]
8133 async fn test_managed_branch_tags() {
8134 use lance::dataset::builder::DatasetBuilder;
8135 use lance::dataset::refs::Ref;
8136 use lance_namespace::models::CreateTableBranchRequest;
8137
8138 let temp = TempStdDir::default();
8139 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8140 let table_id = vec!["t".to_string()];
8141 let main_ds = create_managed_table(&ns, &table_id).await;
8142 ns.create_table_branch(CreateTableBranchRequest {
8143 id: Some(table_id.clone()),
8144 name: "exp".to_string(),
8145 ..Default::default()
8146 })
8147 .await
8148 .unwrap();
8149 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8150 .await
8151 .unwrap()
8152 .with_branch("exp", None)
8153 .load()
8154 .await
8155 .unwrap();
8156 branch_ds
8157 .append(
8158 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8159 None,
8160 )
8161 .await
8162 .unwrap();
8163 let branch_version = branch_ds.version().version;
8164
8165 main_ds
8168 .tags()
8169 .create("exp-tag", ("exp", Some(branch_version)))
8170 .await
8171 .unwrap();
8172 let tag = main_ds.tags().get("exp-tag").await.unwrap();
8173 assert_eq!(tag.branch.as_deref(), Some("exp"));
8174 assert_eq!(tag.version, branch_version);
8175
8176 branch_ds
8178 .tags()
8179 .create("exp-tag2", branch_version)
8180 .await
8181 .unwrap();
8182 let tag2 = branch_ds.tags().get("exp-tag2").await.unwrap();
8183 assert_eq!(tag2.branch.as_deref(), Some("exp"));
8184
8185 let tag_open = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8188 .await
8189 .unwrap()
8190 .with_tag("exp-tag")
8191 .load()
8192 .await
8193 .unwrap();
8194 assert_eq!(tag_open.manifest.branch.as_deref(), Some("exp"));
8195 assert_eq!(tag_open.version().version, branch_version);
8196 assert_eq!(scan_id_column(&tag_open).await, vec![1, 2, 3]);
8197
8198 let tag_checkout = main_ds
8200 .checkout_version(Ref::Tag("exp-tag".to_string()))
8201 .await
8202 .unwrap();
8203 assert_eq!(tag_checkout.manifest.branch.as_deref(), Some("exp"));
8204 assert_eq!(scan_id_column(&tag_checkout).await, vec![1, 2, 3]);
8205
8206 let err = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8208 .await
8209 .unwrap()
8210 .with_tag("no-such-tag")
8211 .load()
8212 .await;
8213 assert!(err.is_err(), "a missing tag must error");
8214 }
8215
8216 #[tokio::test]
8221 async fn test_managed_cross_branch_checkout() {
8222 use lance::dataset::builder::DatasetBuilder;
8223 use lance::dataset::refs::Ref;
8224 use lance_namespace::models::CreateTableBranchRequest;
8225
8226 let temp = TempStdDir::default();
8227 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8228 let table_id = vec!["t".to_string()];
8229 let mut main_ds = create_managed_table(&ns, &table_id).await;
8230 ns.create_table_branch(CreateTableBranchRequest {
8231 id: Some(table_id.clone()),
8232 name: "exp".to_string(),
8233 ..Default::default()
8234 })
8235 .await
8236 .unwrap();
8237
8238 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8240 .await
8241 .unwrap()
8242 .with_branch("exp", None)
8243 .load()
8244 .await
8245 .unwrap();
8246 branch_ds
8247 .append(
8248 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8249 None,
8250 )
8251 .await
8252 .unwrap();
8253 let overlap_version = branch_ds.version().version;
8254 while main_ds.version().version < overlap_version {
8255 main_ds
8256 .append(
8257 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
8258 None,
8259 )
8260 .await
8261 .unwrap();
8262 }
8263
8264 let on_branch = main_ds
8267 .checkout_version(Ref::Version(Some("exp".to_string()), Some(overlap_version)))
8268 .await
8269 .unwrap();
8270 assert_eq!(on_branch.manifest.branch.as_deref(), Some("exp"));
8271 assert_eq!(scan_id_column(&on_branch).await, vec![1, 2, 3]);
8272
8273 let mut on_branch_latest = main_ds.checkout_branch("exp").await.unwrap();
8275 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8276 assert_eq!(on_branch_latest.version().version, overlap_version);
8277
8278 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
8281 ns.list_table_versions(ListTableVersionsRequest {
8282 id: Some(table_id),
8283 ..Default::default()
8284 })
8285 .await
8286 .unwrap()
8287 .versions
8288 .len()
8289 };
8290 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
8291 on_branch_latest
8292 .append(
8293 RecordBatchIterator::new(vec![Ok(single_int_batch(4))], single_int_schema()),
8294 None,
8295 )
8296 .await
8297 .unwrap();
8298 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8299 assert_eq!(scan_id_column(&on_branch_latest).await, vec![1, 2, 3, 4]);
8300 assert_eq!(
8301 main_chain_len(ns.clone(), table_id.clone()).await,
8302 main_before,
8303 "a commit on the checked-out branch must not advance main's chain"
8304 );
8305
8306 let on_main = branch_ds
8308 .checkout_version(Ref::Version(None, Some(1)))
8309 .await
8310 .unwrap();
8311 assert_eq!(on_main.manifest.branch, None);
8312 assert_eq!(scan_id_column(&on_main).await, vec![1]);
8313
8314 ns.create_table_branch(CreateTableBranchRequest {
8316 id: Some(table_id.clone()),
8317 name: "exp2".to_string(),
8318 ..Default::default()
8319 })
8320 .await
8321 .unwrap();
8322 let on_branch2 = branch_ds.checkout_branch("exp2").await.unwrap();
8323 assert_eq!(on_branch2.manifest.branch.as_deref(), Some("exp2"));
8324
8325 let err = main_ds
8327 .checkout_version(Ref::Version(Some("exp".to_string()), Some(999)))
8328 .await;
8329 assert!(err.is_err(), "a version missing from the branch must error");
8330 }
8331
8332 #[tokio::test]
8338 async fn test_commit_builder_honors_explicit_handler_for_dataset_dest() {
8339 use lance::dataset::write::{CommitBuilder, InsertBuilder};
8340 use lance::dataset::{WriteDestination, builder::DatasetBuilder};
8341 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
8342 use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
8343
8344 let temp = TempStdDir::default();
8345 let namespace = Arc::new(
8346 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
8347 .manifest_enabled(true)
8348 .table_version_tracking_enabled(true)
8349 .ops_metrics_enabled(true)
8350 .build()
8351 .await
8352 .unwrap(),
8353 );
8354 let ns: Arc<dyn LanceNamespace> = namespace.clone();
8355 let table_id = vec!["t".to_string()];
8356 create_managed_table(&ns, &table_id).await; let table_uri = ns
8360 .describe_table(DescribeTableRequest {
8361 id: Some(table_id.clone()),
8362 ..Default::default()
8363 })
8364 .await
8365 .unwrap()
8366 .location
8367 .unwrap();
8368 let plain_ds = Arc::new(Dataset::open(&table_uri).await.unwrap());
8369
8370 let transaction = InsertBuilder::new(WriteDestination::Dataset(plain_ds.clone()))
8371 .with_params(&WriteParams {
8372 mode: WriteMode::Append,
8373 ..Default::default()
8374 })
8375 .execute_uncommitted(vec![single_int_batch(3)])
8376 .await
8377 .unwrap();
8378
8379 let handler = Arc::new(ExternalManifestCommitHandler {
8380 external_manifest_store: Arc::new(
8381 LanceNamespaceExternalManifestStore::for_table_uri(
8382 ns.clone(),
8383 table_id.clone(),
8384 &table_uri,
8385 )
8386 .unwrap(),
8387 ),
8388 });
8389 let commits_before = namespace
8390 .retrieve_ops_metrics()
8391 .get("create_table_version")
8392 .copied()
8393 .unwrap_or(0);
8394 let committed = CommitBuilder::new(WriteDestination::Dataset(plain_ds))
8395 .with_commit_handler(handler)
8396 .execute(transaction)
8397 .await
8398 .unwrap();
8399 assert_eq!(scan_id_column(&committed).await, vec![1, 2, 3]);
8400
8401 let commits_after = namespace
8402 .retrieve_ops_metrics()
8403 .get("create_table_version")
8404 .copied()
8405 .unwrap_or(0);
8406 assert_eq!(
8407 commits_after,
8408 commits_before + 1,
8409 "the explicit handler must route the commit through create_table_version"
8410 );
8411 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8412 .await
8413 .unwrap()
8414 .load()
8415 .await
8416 .unwrap();
8417 assert_eq!(
8418 scan_id_column(&fresh).await,
8419 vec![1, 2, 3],
8420 "a fresh managed open must resolve the committed version"
8421 );
8422 }
8423
8424 #[tokio::test]
8426 async fn test_managed_branch_from_non_latest_fork() {
8427 use lance::dataset::builder::DatasetBuilder;
8428 use lance_namespace::models::CreateTableBranchRequest;
8429
8430 let temp = TempStdDir::default();
8431 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8432 let table_id = vec!["t".to_string()];
8433 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
8436 id: Some(table_id.clone()),
8437 name: "old".to_string(),
8438 from_version: Some(1),
8439 ..Default::default()
8440 })
8441 .await
8442 .unwrap();
8443
8444 let old_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8445 .await
8446 .unwrap()
8447 .with_branch("old", None)
8448 .load()
8449 .await
8450 .unwrap();
8451 assert_eq!(old_ds.manifest.branch.as_deref(), Some("old"));
8452 assert_eq!(
8453 scan_id_column(&old_ds).await,
8454 vec![1],
8455 "the fork must contain only the fork-point data"
8456 );
8457 }
8458
8459 #[test]
8462 fn test_manifest_version_from_filename() {
8463 assert_eq!(
8465 DirectoryNamespace::manifest_version_from_filename("5.manifest"),
8466 Some(5)
8467 );
8468 assert_eq!(
8469 DirectoryNamespace::manifest_version_from_filename("0.manifest"),
8470 Some(0)
8471 );
8472 let v2_five = format!("{:020}.manifest", u64::MAX - 5);
8474 assert_eq!(
8475 DirectoryNamespace::manifest_version_from_filename(&v2_five),
8476 Some(5)
8477 );
8478 let v2_zero = format!("{:020}.manifest", u64::MAX);
8479 assert_eq!(
8480 DirectoryNamespace::manifest_version_from_filename(&v2_zero),
8481 Some(0)
8482 );
8483 assert_eq!(
8485 DirectoryNamespace::manifest_version_from_filename("data.lance"),
8486 None
8487 );
8488 assert_eq!(
8489 DirectoryNamespace::manifest_version_from_filename("d5.manifest"),
8490 None
8491 );
8492 }
8493
8494 #[tokio::test]
8495 async fn test_create_table() {
8496 let (namespace, _temp_dir) = create_test_namespace().await;
8497
8498 let schema = create_test_schema();
8500 let ipc_data = create_test_ipc_data(&schema);
8501
8502 let mut request = CreateTableRequest::new();
8503 request.id = Some(vec!["test_table".to_string()]);
8504
8505 let response = namespace
8506 .create_table(request, bytes::Bytes::from(ipc_data))
8507 .await
8508 .unwrap();
8509
8510 assert!(response.location.is_some());
8511 assert!(response.location.unwrap().ends_with("test_table.lance"));
8512 assert_eq!(response.version, Some(1));
8513 }
8514
8515 #[tokio::test]
8516 async fn test_create_table_without_data() {
8517 let (namespace, _temp_dir) = create_test_namespace().await;
8518
8519 let mut request = CreateTableRequest::new();
8520 request.id = Some(vec!["test_table".to_string()]);
8521
8522 let result = namespace.create_table(request, bytes::Bytes::new()).await;
8523 assert!(result.is_err());
8524 assert!(
8525 result
8526 .unwrap_err()
8527 .to_string()
8528 .contains("Arrow IPC stream) is required")
8529 );
8530 }
8531
8532 #[tokio::test]
8533 async fn test_create_table_with_invalid_id() {
8534 let (namespace, _temp_dir) = create_test_namespace().await;
8535
8536 let schema = create_test_schema();
8538 let ipc_data = create_test_ipc_data(&schema);
8539
8540 let mut request = CreateTableRequest::new();
8542 request.id = Some(vec![]);
8543
8544 let result = namespace
8545 .create_table(request, bytes::Bytes::from(ipc_data.clone()))
8546 .await;
8547 assert!(result.is_err());
8548
8549 let mut create_ns_req = CreateNamespaceRequest::new();
8552 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
8553 namespace.create_namespace(create_ns_req).await.unwrap();
8554
8555 let mut request = CreateTableRequest::new();
8557 request.id = Some(vec!["test_namespace".to_string(), "table".to_string()]);
8558
8559 let result = namespace
8560 .create_table(request, bytes::Bytes::from(ipc_data))
8561 .await;
8562 assert!(
8564 result.is_ok(),
8565 "Multi-level table IDs should work with manifest enabled"
8566 );
8567 }
8568
8569 #[tokio::test]
8570 async fn test_list_tables() {
8571 let (namespace, _temp_dir) = create_test_namespace().await;
8572
8573 let mut request = ListTablesRequest::new();
8575 request.id = Some(vec![]);
8576 let response = namespace.list_tables(request).await.unwrap();
8577 assert_eq!(response.tables.len(), 0);
8578
8579 let schema = create_test_schema();
8581 let ipc_data = create_test_ipc_data(&schema);
8582
8583 let mut create_request = CreateTableRequest::new();
8585 create_request.id = Some(vec!["table1".to_string()]);
8586 namespace
8587 .create_table(create_request, bytes::Bytes::from(ipc_data.clone()))
8588 .await
8589 .unwrap();
8590
8591 let mut create_request = CreateTableRequest::new();
8593 create_request.id = Some(vec!["table2".to_string()]);
8594 namespace
8595 .create_table(create_request, bytes::Bytes::from(ipc_data))
8596 .await
8597 .unwrap();
8598
8599 let mut request = ListTablesRequest::new();
8601 request.id = Some(vec![]);
8602 let response = namespace.list_tables(request).await.unwrap();
8603 let tables = response.tables;
8604 assert_eq!(tables.len(), 2);
8605 assert!(tables.contains(&"table1".to_string()));
8606 assert!(tables.contains(&"table2".to_string()));
8607 }
8608
8609 #[tokio::test]
8610 async fn test_list_tables_pagination() {
8611 let (namespace, _temp_dir) = create_test_namespace().await;
8612
8613 let schema = create_test_schema();
8614 let ipc_data = create_test_ipc_data(&schema);
8615
8616 for name in ["alpha", "bravo", "charlie"] {
8617 let mut req = CreateTableRequest::new();
8618 req.id = Some(vec![name.to_string()]);
8619 namespace
8620 .create_table(req, bytes::Bytes::from(ipc_data.clone()))
8621 .await
8622 .unwrap();
8623 }
8624
8625 let first_page = namespace
8627 .list_tables(ListTablesRequest {
8628 id: Some(vec![]),
8629 limit: Some(2),
8630 ..Default::default()
8631 })
8632 .await
8633 .unwrap();
8634
8635 assert_eq!(first_page.tables, vec!["alpha", "bravo"]);
8636 assert_eq!(first_page.page_token.as_deref(), Some("bravo"));
8637
8638 let second_page = namespace
8640 .list_tables(ListTablesRequest {
8641 id: Some(vec![]),
8642 limit: Some(2),
8643 page_token: first_page.page_token.clone(),
8644 ..Default::default()
8645 })
8646 .await
8647 .unwrap();
8648
8649 assert_eq!(second_page.tables, vec!["charlie"]);
8650 assert!(second_page.page_token.is_none());
8651 }
8652
8653 #[tokio::test]
8654 async fn test_list_tables_pagination_limit_zero() {
8655 let (namespace, _temp_dir) = create_test_namespace().await;
8656
8657 let schema = create_test_schema();
8658 let ipc_data = create_test_ipc_data(&schema);
8659
8660 let mut req = CreateTableRequest::new();
8661 req.id = Some(vec!["alpha".to_string()]);
8662 namespace
8663 .create_table(req, bytes::Bytes::from(ipc_data))
8664 .await
8665 .unwrap();
8666
8667 let response = namespace
8668 .list_tables(ListTablesRequest {
8669 id: Some(vec![]),
8670 limit: Some(0),
8671 ..Default::default()
8672 })
8673 .await
8674 .unwrap();
8675
8676 assert!(response.tables.is_empty());
8677 assert!(response.page_token.is_none());
8678 }
8679
8680 #[tokio::test]
8681 async fn test_list_tables_with_namespace_id() {
8682 let (namespace, _temp_dir) = create_test_namespace().await;
8683
8684 let mut create_ns_req = CreateNamespaceRequest::new();
8686 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
8687 namespace.create_namespace(create_ns_req).await.unwrap();
8688
8689 let mut request = ListTablesRequest::new();
8691 request.id = Some(vec!["test_namespace".to_string()]);
8692
8693 let result = namespace.list_tables(request).await;
8694 assert!(
8696 result.is_ok(),
8697 "list_tables should work with child namespace when manifest is enabled"
8698 );
8699 let response = result.unwrap();
8700 assert_eq!(
8701 response.tables.len(),
8702 0,
8703 "Namespace should have no tables yet"
8704 );
8705 }
8706
8707 #[tokio::test]
8708 async fn test_create_scalar_index() {
8709 let (namespace, _temp_dir) = create_test_namespace().await;
8710 create_scalar_table(&namespace, "users").await;
8711
8712 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8713 let dataset = open_dataset(&namespace, "users").await;
8714 let expected_transaction_id = dataset
8715 .read_transaction()
8716 .await
8717 .unwrap()
8718 .map(|transaction| transaction.uuid);
8719 assert_eq!(transaction_id, expected_transaction_id);
8720 let indices = dataset.load_indices().await.unwrap();
8721 assert!(indices.iter().any(|index| index.name == "users_id_idx"));
8722 }
8723
8724 #[tokio::test]
8725 async fn test_create_vector_index() {
8726 use lance_namespace::models::CreateTableIndexRequest;
8727
8728 let (namespace, _temp_dir) = create_test_namespace().await;
8729 create_vector_table(&namespace, "vectors").await;
8730
8731 let mut create_index_request =
8732 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
8733 create_index_request.id = Some(vec!["vectors".to_string()]);
8734 create_index_request.name = Some("vector_idx".to_string());
8735 create_index_request.distance_type = Some("l2".to_string());
8736 let transaction_id = namespace
8737 .create_table_index(create_index_request)
8738 .await
8739 .unwrap()
8740 .transaction_id;
8741
8742 let dataset = open_dataset(&namespace, "vectors").await;
8743 let expected_transaction_id = dataset
8744 .read_transaction()
8745 .await
8746 .unwrap()
8747 .map(|transaction| transaction.uuid);
8748 assert_eq!(transaction_id, expected_transaction_id);
8749 let indices = dataset.load_indices().await.unwrap();
8750 assert!(indices.iter().any(|index| index.name == "vector_idx"));
8751 }
8752
8753 #[tokio::test]
8754 async fn test_list_table_indices() {
8755 use lance_namespace::models::{CreateTableIndexRequest, ListTableIndicesRequest};
8756
8757 let (namespace, _temp_dir) = create_test_namespace().await;
8758 create_scalar_table(&namespace, "users").await;
8759 create_scalar_index(&namespace, "users", "a_idx").await;
8760 create_scalar_index(&namespace, "users", "b_idx").await;
8761 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8762
8763 let response = namespace
8764 .list_table_indices(ListTableIndicesRequest {
8765 id: Some(vec!["users".to_string()]),
8766 ..Default::default()
8767 })
8768 .await
8769 .unwrap();
8770
8771 assert_eq!(response.indexes.len(), 3);
8772 assert_eq!(response.indexes[0].index_name, "a_idx");
8773 assert_eq!(response.indexes[1].index_name, "b_idx");
8774 assert_eq!(response.indexes[2].index_name, "users_id_idx");
8775 assert!(response.page_token.is_none());
8776 let users_id_idx = response
8777 .indexes
8778 .iter()
8779 .find(|index| index.index_name == "users_id_idx")
8780 .unwrap();
8781 assert_eq!(users_id_idx.columns, vec!["id"]);
8782 assert_eq!(users_id_idx.status, "SUCCEEDED");
8783
8784 assert_eq!(users_id_idx.index_type.as_deref(), Some("BTree"));
8786 assert!(
8787 users_id_idx
8788 .type_url
8789 .as_deref()
8790 .is_some_and(|s| !s.is_empty())
8791 );
8792 assert_eq!(users_id_idx.num_indexed_rows, Some(3));
8793 assert_eq!(users_id_idx.num_unindexed_rows, Some(0));
8794 assert_eq!(users_id_idx.num_segments, Some(1));
8795 assert!(users_id_idx.size_bytes.is_some_and(|size| size > 0));
8796 assert!(users_id_idx.created_at.is_some());
8797 assert!(users_id_idx.index_version.is_some());
8798 assert!(users_id_idx.index_details.is_some());
8799
8800 let dataset = open_dataset(&namespace, "users").await;
8801 let expected_transaction_id = dataset
8802 .read_transaction()
8803 .await
8804 .unwrap()
8805 .map(|transaction| transaction.uuid);
8806 assert_eq!(transaction_id, expected_transaction_id);
8807 let indices = dataset.load_indices().await.unwrap();
8808 assert_eq!(
8809 indices
8810 .iter()
8811 .filter(|index| index.name == "users_id_idx")
8812 .count(),
8813 1
8814 );
8815
8816 let first_page = namespace
8817 .list_table_indices(ListTableIndicesRequest {
8818 id: Some(vec!["users".to_string()]),
8819 limit: Some(2),
8820 ..Default::default()
8821 })
8822 .await
8823 .unwrap();
8824
8825 assert_eq!(first_page.indexes.len(), 2);
8826 assert_eq!(first_page.indexes[0].index_name, "a_idx");
8827 assert_eq!(first_page.indexes[1].index_name, "b_idx");
8828 assert_eq!(first_page.page_token.as_deref(), Some("b_idx"));
8829
8830 let second_page = namespace
8831 .list_table_indices(ListTableIndicesRequest {
8832 id: Some(vec!["users".to_string()]),
8833 page_token: first_page.page_token.clone(),
8834 limit: Some(2),
8835 ..Default::default()
8836 })
8837 .await
8838 .unwrap();
8839
8840 assert_eq!(second_page.indexes.len(), 1);
8841 assert_eq!(second_page.indexes[0].index_name, "users_id_idx");
8842 assert!(second_page.page_token.is_none());
8843
8844 create_vector_table(&namespace, "vectors").await;
8846 let mut create_index_request =
8847 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
8848 create_index_request.id = Some(vec!["vectors".to_string()]);
8849 create_index_request.name = Some("vector_idx".to_string());
8850 create_index_request.distance_type = Some("l2".to_string());
8851 namespace
8852 .create_table_index(create_index_request)
8853 .await
8854 .unwrap();
8855
8856 let vector_response = namespace
8857 .list_table_indices(ListTableIndicesRequest {
8858 id: Some(vec!["vectors".to_string()]),
8859 ..Default::default()
8860 })
8861 .await
8862 .unwrap();
8863
8864 assert_eq!(vector_response.indexes.len(), 1);
8865 let vector_idx = &vector_response.indexes[0];
8866 assert_eq!(vector_idx.index_name, "vector_idx");
8867 assert_eq!(vector_idx.columns, vec!["vector"]);
8868 assert_eq!(vector_idx.index_type.as_deref(), Some("IVF_FLAT"));
8869 assert!(
8870 vector_idx
8871 .type_url
8872 .as_deref()
8873 .is_some_and(|s| !s.is_empty())
8874 );
8875 assert!(vector_idx.num_indexed_rows.is_some());
8876 assert!(vector_idx.num_unindexed_rows.is_some());
8877 assert_eq!(vector_idx.num_segments, Some(1));
8878 assert!(vector_idx.created_at.is_some());
8879 assert!(vector_idx.index_version.is_some());
8880 assert!(vector_idx.index_details.is_some());
8881 }
8882
8883 #[tokio::test]
8884 async fn test_describe_table_index_stats() {
8885 use lance_namespace::models::DescribeTableIndexStatsRequest;
8886
8887 let (namespace, _temp_dir) = create_test_namespace().await;
8888 create_scalar_table(&namespace, "users").await;
8889 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8890
8891 let response = namespace
8892 .describe_table_index_stats(DescribeTableIndexStatsRequest {
8893 id: Some(vec!["users".to_string()]),
8894 index_name: Some("users_id_idx".to_string()),
8895 ..Default::default()
8896 })
8897 .await
8898 .unwrap();
8899 assert_eq!(response.index_type, Some("BTree".to_string()));
8900 assert_eq!(response.num_indices, Some(1));
8901 assert_eq!(response.num_indexed_rows, Some(3));
8902 assert_eq!(response.num_unindexed_rows, Some(0));
8903
8904 let dataset = open_dataset(&namespace, "users").await;
8905 let expected_transaction_id = dataset
8906 .read_transaction()
8907 .await
8908 .unwrap()
8909 .map(|transaction| transaction.uuid);
8910 assert_eq!(transaction_id, expected_transaction_id);
8911 let stats: serde_json::Value =
8912 serde_json::from_str(&dataset.index_statistics("users_id_idx").await.unwrap()).unwrap();
8913 assert_eq!(stats["index_type"], "BTree");
8914 assert_eq!(stats["num_indices"], 1);
8915 assert_eq!(stats["num_indexed_rows"], 3);
8916 assert_eq!(stats["num_unindexed_rows"], 0);
8917 }
8918
8919 #[tokio::test]
8920 async fn test_describe_transaction() {
8921 use lance_namespace::models::DescribeTransactionRequest;
8922
8923 let (namespace, _temp_dir) = create_test_namespace().await;
8924 create_scalar_table(&namespace, "users").await;
8925 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8926 let dataset = open_dataset(&namespace, "users").await;
8927 let latest_transaction = dataset.read_transaction().await.unwrap();
8928 assert_eq!(
8929 transaction_id,
8930 latest_transaction
8931 .as_ref()
8932 .map(|transaction| transaction.uuid.clone())
8933 );
8934
8935 if let Some(transaction_id) = transaction_id {
8936 let response = namespace
8937 .describe_transaction(DescribeTransactionRequest {
8938 id: Some(vec!["users".to_string(), transaction_id.clone()]),
8939 ..Default::default()
8940 })
8941 .await
8942 .unwrap();
8943 assert_eq!(response.status, "SUCCEEDED");
8944 assert_eq!(
8945 response
8946 .properties
8947 .as_ref()
8948 .and_then(|props| props.get("operation")),
8949 Some(&"CreateIndex".to_string())
8950 );
8951 assert_eq!(
8952 response
8953 .properties
8954 .as_ref()
8955 .and_then(|props| props.get("uuid")),
8956 Some(&transaction_id)
8957 );
8958 } else {
8959 assert!(latest_transaction.is_none());
8960 }
8961 }
8962
8963 #[tokio::test]
8964 async fn test_drop_table_index() {
8965 use lance_namespace::models::{DropTableIndexRequest, ListTableIndicesRequest};
8966
8967 let (namespace, _temp_dir) = create_test_namespace().await;
8968 create_scalar_table(&namespace, "users").await;
8969 let create_transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8970
8971 let drop_transaction_id = namespace
8972 .drop_table_index(DropTableIndexRequest {
8973 id: Some(vec!["users".to_string()]),
8974 index_name: Some("users_id_idx".to_string()),
8975 ..Default::default()
8976 })
8977 .await
8978 .unwrap()
8979 .transaction_id;
8980
8981 let dataset = open_dataset(&namespace, "users").await;
8982 let previous_dataset = dataset
8983 .checkout_version(dataset.version().version - 1)
8984 .await
8985 .unwrap();
8986 let previous_transaction_id = previous_dataset
8987 .read_transaction()
8988 .await
8989 .unwrap()
8990 .map(|transaction| transaction.uuid);
8991 assert_eq!(create_transaction_id, previous_transaction_id);
8992 let expected_drop_transaction_id = dataset
8993 .read_transaction()
8994 .await
8995 .unwrap()
8996 .map(|transaction| transaction.uuid);
8997 assert_eq!(drop_transaction_id, expected_drop_transaction_id);
8998 let indices = dataset.load_indices().await.unwrap();
8999 assert!(!indices.iter().any(|index| index.name == "users_id_idx"));
9000
9001 let list_response = namespace
9002 .list_table_indices(ListTableIndicesRequest {
9003 id: Some(vec!["users".to_string()]),
9004 ..Default::default()
9005 })
9006 .await
9007 .unwrap();
9008 assert!(list_response.indexes.is_empty());
9009 }
9010
9011 #[tokio::test]
9012 async fn test_describe_table() {
9013 let (namespace, _temp_dir) = create_test_namespace().await;
9014
9015 let schema = create_test_schema();
9017 let ipc_data = create_test_ipc_data(&schema);
9018
9019 let mut create_request = CreateTableRequest::new();
9020 create_request.id = Some(vec!["test_table".to_string()]);
9021 namespace
9022 .create_table(create_request, bytes::Bytes::from(ipc_data))
9023 .await
9024 .unwrap();
9025
9026 let mut request = DescribeTableRequest::new();
9028 request.id = Some(vec!["test_table".to_string()]);
9029 let response = namespace.describe_table(request).await.unwrap();
9030
9031 assert!(response.location.is_some());
9032 assert!(response.location.unwrap().ends_with("test_table.lance"));
9033 }
9034
9035 #[tokio::test]
9036 async fn test_describe_nonexistent_table() {
9037 let (namespace, _temp_dir) = create_test_namespace().await;
9038
9039 let mut request = DescribeTableRequest::new();
9040 request.id = Some(vec!["nonexistent".to_string()]);
9041
9042 let result = namespace.describe_table(request).await;
9043 assert!(result.is_err());
9044 assert!(result.unwrap_err().to_string().contains("Table not found"));
9045 }
9046
9047 #[tokio::test]
9048 async fn test_table_exists() {
9049 let (namespace, _temp_dir) = create_test_namespace().await;
9050
9051 let schema = create_test_schema();
9053 let ipc_data = create_test_ipc_data(&schema);
9054
9055 let mut create_request = CreateTableRequest::new();
9056 create_request.id = Some(vec!["existing_table".to_string()]);
9057 namespace
9058 .create_table(create_request, bytes::Bytes::from(ipc_data))
9059 .await
9060 .unwrap();
9061
9062 let mut request = TableExistsRequest::new();
9064 request.id = Some(vec!["existing_table".to_string()]);
9065 let result = namespace.table_exists(request).await;
9066 assert!(result.is_ok());
9067
9068 let mut request = TableExistsRequest::new();
9070 request.id = Some(vec!["nonexistent".to_string()]);
9071 let result = namespace.table_exists(request).await;
9072 assert!(result.is_err());
9073 assert!(result.unwrap_err().to_string().contains("Table not found"));
9074 }
9075
9076 #[tokio::test]
9077 async fn test_drop_table() {
9078 let (namespace, _temp_dir) = create_test_namespace().await;
9079
9080 let schema = create_test_schema();
9082 let ipc_data = create_test_ipc_data(&schema);
9083
9084 let mut create_request = CreateTableRequest::new();
9085 create_request.id = Some(vec!["table_to_drop".to_string()]);
9086 namespace
9087 .create_table(create_request, bytes::Bytes::from(ipc_data))
9088 .await
9089 .unwrap();
9090
9091 let mut exists_request = TableExistsRequest::new();
9093 exists_request.id = Some(vec!["table_to_drop".to_string()]);
9094 assert!(namespace.table_exists(exists_request.clone()).await.is_ok());
9095
9096 let mut drop_request = DropTableRequest::new();
9098 drop_request.id = Some(vec!["table_to_drop".to_string()]);
9099 let response = namespace.drop_table(drop_request).await.unwrap();
9100 assert!(response.location.is_some());
9101
9102 assert!(namespace.table_exists(exists_request).await.is_err());
9104 }
9105
9106 #[tokio::test]
9107 async fn test_drop_nonexistent_table() {
9108 let (namespace, _temp_dir) = create_test_namespace().await;
9109
9110 let mut request = DropTableRequest::new();
9111 request.id = Some(vec!["nonexistent".to_string()]);
9112
9113 let result = namespace.drop_table(request).await;
9115 let _ = result;
9118 }
9119
9120 #[tokio::test]
9121 async fn test_root_namespace_operations() {
9122 let (namespace, _temp_dir) = create_test_namespace().await;
9123
9124 let mut request = ListNamespacesRequest::new();
9126 request.id = Some(vec![]);
9127 let result = namespace.list_namespaces(request).await;
9128 assert!(result.is_ok());
9129 assert_eq!(result.unwrap().namespaces.len(), 0);
9130
9131 let mut request = DescribeNamespaceRequest::new();
9133 request.id = Some(vec![]);
9134 let result = namespace.describe_namespace(request).await;
9135 assert!(result.is_ok());
9136
9137 let mut request = NamespaceExistsRequest::new();
9139 request.id = Some(vec![]);
9140 let result = namespace.namespace_exists(request).await;
9141 assert!(result.is_ok());
9142
9143 let mut request = CreateNamespaceRequest::new();
9145 request.id = Some(vec![]);
9146 let result = namespace.create_namespace(request).await;
9147 assert!(result.is_err());
9148 assert!(result.unwrap_err().to_string().contains("already exists"));
9149
9150 let mut request = DropNamespaceRequest::new();
9152 request.id = Some(vec![]);
9153 let result = namespace.drop_namespace(request).await;
9154 assert!(result.is_err());
9155 assert!(
9156 result
9157 .unwrap_err()
9158 .to_string()
9159 .contains("cannot be dropped")
9160 );
9161 }
9162
9163 #[tokio::test]
9164 async fn test_non_root_namespace_operations() {
9165 let (namespace, _temp_dir) = create_test_namespace().await;
9166
9167 let mut request = CreateNamespaceRequest::new();
9170 request.id = Some(vec!["child".to_string()]);
9171 let result = namespace.create_namespace(request).await;
9172 assert!(
9173 result.is_ok(),
9174 "Child namespace creation should succeed with manifest enabled"
9175 );
9176
9177 let mut request = NamespaceExistsRequest::new();
9179 request.id = Some(vec!["child".to_string()]);
9180 let result = namespace.namespace_exists(request).await;
9181 assert!(
9182 result.is_ok(),
9183 "Child namespace should exist after creation"
9184 );
9185
9186 let mut request = DropNamespaceRequest::new();
9188 request.id = Some(vec!["child".to_string()]);
9189 let result = namespace.drop_namespace(request).await;
9190 assert!(
9191 result.is_ok(),
9192 "Child namespace drop should succeed with manifest enabled"
9193 );
9194
9195 let mut request = NamespaceExistsRequest::new();
9197 request.id = Some(vec!["child".to_string()]);
9198 let result = namespace.namespace_exists(request).await;
9199 assert!(
9200 result.is_err(),
9201 "Child namespace should not exist after drop"
9202 );
9203 }
9204
9205 #[tokio::test]
9206 async fn test_config_custom_root() {
9207 let temp_dir = TempStdDir::default();
9208 let custom_path = temp_dir.join("custom");
9209 std::fs::create_dir(&custom_path).unwrap();
9210
9211 let namespace = DirectoryNamespaceBuilder::new(custom_path.to_string_lossy().to_string())
9212 .build()
9213 .await
9214 .unwrap();
9215
9216 let schema = create_test_schema();
9218 let ipc_data = create_test_ipc_data(&schema);
9219
9220 let mut request = CreateTableRequest::new();
9222 request.id = Some(vec!["test_table".to_string()]);
9223
9224 let response = namespace
9225 .create_table(request, bytes::Bytes::from(ipc_data))
9226 .await
9227 .unwrap();
9228
9229 assert!(response.location.unwrap().contains("custom"));
9230 }
9231
9232 #[tokio::test]
9233 async fn test_config_storage_options() {
9234 let temp_dir = TempStdDir::default();
9235
9236 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9237 .storage_option("option1", "value1")
9238 .storage_option("option2", "value2")
9239 .build()
9240 .await
9241 .unwrap();
9242
9243 let schema = create_test_schema();
9245 let ipc_data = create_test_ipc_data(&schema);
9246
9247 let mut request = CreateTableRequest::new();
9249 request.id = Some(vec!["test_table".to_string()]);
9250
9251 let response = namespace
9252 .create_table(request, bytes::Bytes::from(ipc_data))
9253 .await
9254 .unwrap();
9255
9256 let storage_options = response.storage_options.unwrap();
9257 assert_eq!(storage_options.get("option1"), Some(&"value1".to_string()));
9258 assert_eq!(storage_options.get("option2"), Some(&"value2".to_string()));
9259 }
9260
9261 #[tokio::test]
9265 async fn test_no_storage_options_without_vendor() {
9266 use lance_namespace::models::DeclareTableRequest;
9267
9268 let temp_dir = TempStdDir::default();
9269
9270 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9272 .manifest_enabled(false)
9273 .storage_option("aws_access_key_id", "AKID")
9274 .storage_option("aws_secret_access_key", "SECRET")
9275 .storage_option("region", "us-east-1")
9276 .build()
9277 .await
9278 .unwrap();
9279
9280 let schema = create_test_schema();
9281 let ipc_data = create_test_ipc_data(&schema);
9282
9283 let mut create_req = CreateTableRequest::new();
9285 create_req.id = Some(vec!["t1".to_string()]);
9286 namespace
9287 .create_table(create_req, bytes::Bytes::from(ipc_data))
9288 .await
9289 .unwrap();
9290
9291 let mut desc_req = DescribeTableRequest::new();
9293 desc_req.id = Some(vec!["t1".to_string()]);
9294 let resp = namespace.describe_table(desc_req).await.unwrap();
9295 assert!(resp.storage_options.is_none());
9296
9297 let mut decl_req = DeclareTableRequest::new();
9299 decl_req.id = Some(vec!["t2".to_string()]);
9300 let resp = namespace.declare_table(decl_req).await.unwrap();
9301 assert!(resp.storage_options.is_none());
9302 }
9303
9304 #[tokio::test]
9306 async fn test_no_storage_options_without_vendor_manifest() {
9307 let temp_dir = TempStdDir::default();
9308
9309 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9310 .storage_option("aws_access_key_id", "AKID")
9311 .storage_option("aws_secret_access_key", "SECRET")
9312 .storage_option("region", "us-east-1")
9313 .build()
9314 .await
9315 .unwrap();
9316
9317 let schema = create_test_schema();
9318 let ipc_data = create_test_ipc_data(&schema);
9319
9320 let mut create_req = CreateTableRequest::new();
9321 create_req.id = Some(vec!["t1".to_string()]);
9322 namespace
9323 .create_table(create_req, bytes::Bytes::from(ipc_data))
9324 .await
9325 .unwrap();
9326
9327 let mut desc_req = DescribeTableRequest::new();
9329 desc_req.id = Some(vec!["t1".to_string()]);
9330 let resp = namespace.describe_table(desc_req).await.unwrap();
9331 assert!(resp.storage_options.is_none());
9332 }
9333
9334 #[tokio::test]
9335 async fn test_from_properties_manifest_enabled() {
9336 let temp_dir = TempStdDir::default();
9337
9338 let mut properties = HashMap::new();
9339 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9340 properties.insert("manifest_enabled".to_string(), "true".to_string());
9341 properties.insert("dir_listing_enabled".to_string(), "false".to_string());
9342
9343 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9344 assert!(builder.manifest_enabled);
9345 assert!(!builder.dir_listing_enabled);
9346
9347 let namespace = builder.build().await.unwrap();
9348
9349 let schema = create_test_schema();
9351 let ipc_data = create_test_ipc_data(&schema);
9352
9353 let mut request = CreateTableRequest::new();
9355 request.id = Some(vec!["test_table".to_string()]);
9356
9357 let response = namespace
9358 .create_table(request, bytes::Bytes::from(ipc_data))
9359 .await
9360 .unwrap();
9361
9362 assert!(response.location.is_some());
9363 }
9364
9365 #[tokio::test]
9366 async fn test_from_properties_dir_listing_enabled() {
9367 let temp_dir = TempStdDir::default();
9368
9369 let mut properties = HashMap::new();
9370 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9371 properties.insert("manifest_enabled".to_string(), "false".to_string());
9372 properties.insert("dir_listing_enabled".to_string(), "true".to_string());
9373
9374 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9375 assert!(!builder.manifest_enabled);
9376 assert!(builder.dir_listing_enabled);
9377
9378 let namespace = builder.build().await.unwrap();
9379
9380 let schema = create_test_schema();
9382 let ipc_data = create_test_ipc_data(&schema);
9383
9384 let mut request = CreateTableRequest::new();
9386 request.id = Some(vec!["test_table".to_string()]);
9387
9388 let response = namespace
9389 .create_table(request, bytes::Bytes::from(ipc_data))
9390 .await
9391 .unwrap();
9392
9393 assert!(response.location.is_some());
9394 }
9395
9396 #[tokio::test]
9397 async fn test_from_properties_defaults() {
9398 let temp_dir = TempStdDir::default();
9399
9400 let mut properties = HashMap::new();
9401 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9402
9403 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9404 assert!(builder.manifest_enabled);
9406 assert!(builder.dir_listing_enabled);
9407 }
9408
9409 #[tokio::test]
9410 async fn test_from_properties_with_storage_options() {
9411 let temp_dir = TempStdDir::default();
9412
9413 let mut properties = HashMap::new();
9414 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9415 properties.insert("manifest_enabled".to_string(), "true".to_string());
9416 properties.insert("storage.region".to_string(), "us-west-2".to_string());
9417 properties.insert("storage.bucket".to_string(), "my-bucket".to_string());
9418
9419 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9420 assert!(builder.manifest_enabled);
9421 assert!(builder.storage_options.is_some());
9422
9423 let storage_options = builder.storage_options.unwrap();
9424 assert_eq!(
9425 storage_options.get("region"),
9426 Some(&"us-west-2".to_string())
9427 );
9428 assert_eq!(
9429 storage_options.get("bucket"),
9430 Some(&"my-bucket".to_string())
9431 );
9432 }
9433
9434 #[tokio::test]
9435 async fn test_various_arrow_types() {
9436 let (namespace, _temp_dir) = create_test_namespace().await;
9437
9438 let fields = vec![
9440 JsonArrowField {
9441 name: "bool_col".to_string(),
9442 r#type: Box::new(JsonArrowDataType::new("bool".to_string())),
9443 nullable: true,
9444 metadata: None,
9445 },
9446 JsonArrowField {
9447 name: "int8_col".to_string(),
9448 r#type: Box::new(JsonArrowDataType::new("int8".to_string())),
9449 nullable: true,
9450 metadata: None,
9451 },
9452 JsonArrowField {
9453 name: "float64_col".to_string(),
9454 r#type: Box::new(JsonArrowDataType::new("float64".to_string())),
9455 nullable: true,
9456 metadata: None,
9457 },
9458 JsonArrowField {
9459 name: "binary_col".to_string(),
9460 r#type: Box::new(JsonArrowDataType::new("binary".to_string())),
9461 nullable: true,
9462 metadata: None,
9463 },
9464 ];
9465
9466 let schema = JsonArrowSchema {
9467 fields,
9468 metadata: None,
9469 };
9470
9471 let ipc_data = create_test_ipc_data(&schema);
9473
9474 let mut request = CreateTableRequest::new();
9475 request.id = Some(vec!["complex_table".to_string()]);
9476
9477 let response = namespace
9478 .create_table(request, bytes::Bytes::from(ipc_data))
9479 .await
9480 .unwrap();
9481
9482 assert!(response.location.is_some());
9483 }
9484
9485 #[tokio::test]
9486 async fn test_connect_dir() {
9487 let temp_dir = TempStdDir::default();
9488
9489 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9490 .build()
9491 .await
9492 .unwrap();
9493
9494 let mut request = ListTablesRequest::new();
9496 request.id = Some(vec![]);
9497 let response = namespace.list_tables(request).await.unwrap();
9498 assert_eq!(response.tables.len(), 0);
9499 }
9500
9501 #[tokio::test]
9502 async fn test_create_table_with_ipc_data() {
9503 use arrow::array::{Int32Array, StringArray};
9504 use arrow::ipc::writer::StreamWriter;
9505
9506 let (namespace, _temp_dir) = create_test_namespace().await;
9507
9508 let schema = create_test_schema();
9510
9511 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
9513 let arrow_schema = Arc::new(arrow_schema);
9514
9515 let id_array = Int32Array::from(vec![1, 2, 3]);
9517 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
9518 let batch = arrow::record_batch::RecordBatch::try_new(
9519 arrow_schema.clone(),
9520 vec![Arc::new(id_array), Arc::new(name_array)],
9521 )
9522 .unwrap();
9523
9524 let mut buffer = Vec::new();
9526 {
9527 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
9528 writer.write(&batch).unwrap();
9529 writer.finish().unwrap();
9530 }
9531
9532 let mut request = CreateTableRequest::new();
9534 request.id = Some(vec!["test_table_with_data".to_string()]);
9535
9536 let response = namespace
9537 .create_table(request, Bytes::from(buffer))
9538 .await
9539 .unwrap();
9540
9541 assert_eq!(response.version, Some(1));
9542 assert!(
9543 response
9544 .location
9545 .unwrap()
9546 .contains("test_table_with_data.lance")
9547 );
9548
9549 let mut exists_request = TableExistsRequest::new();
9551 exists_request.id = Some(vec!["test_table_with_data".to_string()]);
9552 namespace.table_exists(exists_request).await.unwrap();
9553 }
9554
9555 #[tokio::test]
9556 async fn test_child_namespace_create_and_list() {
9557 let (namespace, _temp_dir) = create_test_namespace().await;
9558
9559 for i in 1..=3 {
9561 let mut create_req = CreateNamespaceRequest::new();
9562 create_req.id = Some(vec![format!("ns{}", i)]);
9563 let result = namespace.create_namespace(create_req).await;
9564 assert!(result.is_ok(), "Failed to create child namespace ns{}", i);
9565 }
9566
9567 let list_req = ListNamespacesRequest {
9569 id: Some(vec![]),
9570 ..Default::default()
9571 };
9572 let result = namespace.list_namespaces(list_req).await;
9573 assert!(result.is_ok());
9574 let namespaces = result.unwrap().namespaces;
9575 assert_eq!(namespaces.len(), 3);
9576 assert!(namespaces.contains(&"ns1".to_string()));
9577 assert!(namespaces.contains(&"ns2".to_string()));
9578 assert!(namespaces.contains(&"ns3".to_string()));
9579 }
9580
9581 #[tokio::test]
9582 async fn test_nested_namespace_hierarchy() {
9583 let (namespace, _temp_dir) = create_test_namespace().await;
9584
9585 let mut create_req = CreateNamespaceRequest::new();
9587 create_req.id = Some(vec!["parent".to_string()]);
9588 namespace.create_namespace(create_req).await.unwrap();
9589
9590 let mut create_req = CreateNamespaceRequest::new();
9592 create_req.id = Some(vec!["parent".to_string(), "child1".to_string()]);
9593 namespace.create_namespace(create_req).await.unwrap();
9594
9595 let mut create_req = CreateNamespaceRequest::new();
9596 create_req.id = Some(vec!["parent".to_string(), "child2".to_string()]);
9597 namespace.create_namespace(create_req).await.unwrap();
9598
9599 let list_req = ListNamespacesRequest {
9601 id: Some(vec!["parent".to_string()]),
9602 ..Default::default()
9603 };
9604 let result = namespace.list_namespaces(list_req).await;
9605 assert!(result.is_ok());
9606 let children = result.unwrap().namespaces;
9607 assert_eq!(children.len(), 2);
9608 assert!(children.contains(&"child1".to_string()));
9609 assert!(children.contains(&"child2".to_string()));
9610
9611 let list_req = ListNamespacesRequest {
9613 id: Some(vec![]),
9614 ..Default::default()
9615 };
9616 let result = namespace.list_namespaces(list_req).await;
9617 assert!(result.is_ok());
9618 let root_namespaces = result.unwrap().namespaces;
9619 assert_eq!(root_namespaces.len(), 1);
9620 assert_eq!(root_namespaces[0], "parent");
9621 }
9622
9623 #[tokio::test]
9624 async fn test_table_in_child_namespace() {
9625 let (namespace, _temp_dir) = create_test_namespace().await;
9626
9627 let mut create_ns_req = CreateNamespaceRequest::new();
9629 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9630 namespace.create_namespace(create_ns_req).await.unwrap();
9631
9632 let schema = create_test_schema();
9634 let ipc_data = create_test_ipc_data(&schema);
9635 let mut create_table_req = CreateTableRequest::new();
9636 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9637 let result = namespace
9638 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9639 .await;
9640 assert!(result.is_ok(), "Failed to create table in child namespace");
9641
9642 let list_req = ListTablesRequest {
9644 id: Some(vec!["test_ns".to_string()]),
9645 ..Default::default()
9646 };
9647 let result = namespace.list_tables(list_req).await;
9648 assert!(result.is_ok());
9649 let tables = result.unwrap().tables;
9650 assert_eq!(tables.len(), 1);
9651 assert_eq!(tables[0], "table1");
9652
9653 let mut exists_req = TableExistsRequest::new();
9655 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9656 let result = namespace.table_exists(exists_req).await;
9657 assert!(result.is_ok());
9658
9659 let mut describe_req = DescribeTableRequest::new();
9661 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9662 let result = namespace.describe_table(describe_req).await;
9663 assert!(result.is_ok());
9664 let response = result.unwrap();
9665 assert!(response.location.is_some());
9666 }
9667
9668 #[tokio::test]
9677 async fn test_read_self_heals_after_manifest_created_by_other_connection() {
9678 let temp_dir = TempStdDir::default();
9679 let root = temp_dir.to_str().unwrap();
9680
9681 let reader = DirectoryNamespaceBuilder::new(root).build().await.unwrap();
9684
9685 let writer = DirectoryNamespaceBuilder::new(root).build().await.unwrap();
9688 let mut create_ns_req = CreateNamespaceRequest::new();
9689 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9690 writer.create_namespace(create_ns_req).await.unwrap();
9691 let ipc_data = create_test_ipc_data(&create_test_schema());
9692 let mut create_table_req = CreateTableRequest::new();
9693 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9694 writer
9695 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9696 .await
9697 .unwrap();
9698
9699 let mut describe_req = DescribeTableRequest::new();
9702 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9703 let resp = reader
9704 .describe_table(describe_req)
9705 .await
9706 .expect("describe_table must resolve a table registered after build");
9707 assert!(resp.location.is_some());
9708
9709 let mut exists_req = TableExistsRequest::new();
9710 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9711 reader
9712 .table_exists(exists_req)
9713 .await
9714 .expect("table_exists must resolve a table registered after build");
9715
9716 let list_req = ListTablesRequest {
9717 id: Some(vec!["test_ns".to_string()]),
9718 ..Default::default()
9719 };
9720 let tables = reader.list_tables(list_req).await.unwrap().tables;
9721 assert_eq!(tables, vec!["table1".to_string()]);
9722 }
9723
9724 #[tokio::test]
9732 async fn test_migration_root_read_self_heals_registered_alias() {
9733 use lance_namespace::models::RegisterTableRequest;
9734
9735 let temp_dir = TempStdDir::default();
9736 let temp_path = temp_dir.to_str().unwrap();
9737
9738 let reader = DirectoryNamespaceBuilder::new(temp_path)
9741 .dir_listing_enabled(true)
9742 .dir_listing_to_manifest_migration_enabled(true)
9743 .build()
9744 .await
9745 .unwrap();
9746
9747 let writer = DirectoryNamespaceBuilder::new(temp_path)
9751 .dir_listing_enabled(true)
9752 .dir_listing_to_manifest_migration_enabled(true)
9753 .build()
9754 .await
9755 .unwrap();
9756 let ipc_data = create_test_ipc_data(&create_test_schema());
9757 let table_uri = format!("{}/external_table.lance", temp_path);
9758 let cursor = Cursor::new(ipc_data);
9759 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
9760 let batches: Vec<_> = stream_reader
9761 .collect::<std::result::Result<Vec<_>, _>>()
9762 .unwrap();
9763 let schema = batches[0].schema();
9764 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
9765 let batch_reader = RecordBatchIterator::new(batch_results, schema);
9766 Dataset::write(Box::new(batch_reader), &table_uri, None)
9767 .await
9768 .unwrap();
9769 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
9770 register_req.id = Some(vec!["registered_table".to_string()]);
9771 writer.register_table(register_req).await.unwrap();
9772
9773 let mut describe_req = DescribeTableRequest::new();
9776 describe_req.id = Some(vec!["registered_table".to_string()]);
9777 reader
9778 .describe_table(describe_req)
9779 .await
9780 .expect("describe_table must resolve a manifest alias registered after build");
9781
9782 let mut exists_req = TableExistsRequest::new();
9783 exists_req.id = Some(vec!["registered_table".to_string()]);
9784 reader
9785 .table_exists(exists_req)
9786 .await
9787 .expect("table_exists must resolve a manifest alias registered after build");
9788
9789 let list_req = ListTablesRequest {
9790 id: Some(vec![]),
9791 ..Default::default()
9792 };
9793 let tables = reader.list_tables(list_req).await.unwrap().tables;
9794 assert!(
9795 tables.contains(&"registered_table".to_string()),
9796 "list_tables must include the manifest alias registered after build, got {:?}",
9797 tables
9798 );
9799 }
9800
9801 #[tokio::test]
9802 async fn test_multiple_tables_in_child_namespace() {
9803 let (namespace, _temp_dir) = create_test_namespace().await;
9804
9805 let mut create_ns_req = CreateNamespaceRequest::new();
9807 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9808 namespace.create_namespace(create_ns_req).await.unwrap();
9809
9810 let schema = create_test_schema();
9812 let ipc_data = create_test_ipc_data(&schema);
9813 for i in 1..=3 {
9814 let mut create_table_req = CreateTableRequest::new();
9815 create_table_req.id = Some(vec!["test_ns".to_string(), format!("table{}", i)]);
9816 namespace
9817 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
9818 .await
9819 .unwrap();
9820 }
9821
9822 let list_req = ListTablesRequest {
9824 id: Some(vec!["test_ns".to_string()]),
9825 ..Default::default()
9826 };
9827 let result = namespace.list_tables(list_req).await;
9828 assert!(result.is_ok());
9829 let tables = result.unwrap().tables;
9830 assert_eq!(tables.len(), 3);
9831 assert!(tables.contains(&"table1".to_string()));
9832 assert!(tables.contains(&"table2".to_string()));
9833 assert!(tables.contains(&"table3".to_string()));
9834 }
9835
9836 #[tokio::test]
9837 async fn test_drop_table_in_child_namespace() {
9838 let (namespace, _temp_dir) = create_test_namespace().await;
9839
9840 let mut create_ns_req = CreateNamespaceRequest::new();
9842 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9843 namespace.create_namespace(create_ns_req).await.unwrap();
9844
9845 let schema = create_test_schema();
9847 let ipc_data = create_test_ipc_data(&schema);
9848 let mut create_table_req = CreateTableRequest::new();
9849 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9850 namespace
9851 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9852 .await
9853 .unwrap();
9854
9855 let mut drop_req = DropTableRequest::new();
9857 drop_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9858 let result = namespace.drop_table(drop_req).await;
9859 assert!(result.is_ok(), "Failed to drop table in child namespace");
9860
9861 let mut exists_req = TableExistsRequest::new();
9863 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9864 let result = namespace.table_exists(exists_req).await;
9865 assert!(result.is_err());
9866 }
9867
9868 #[tokio::test]
9869 async fn test_deeply_nested_namespace() {
9870 let (namespace, _temp_dir) = create_test_namespace().await;
9871
9872 let mut create_req = CreateNamespaceRequest::new();
9874 create_req.id = Some(vec!["level1".to_string()]);
9875 namespace.create_namespace(create_req).await.unwrap();
9876
9877 let mut create_req = CreateNamespaceRequest::new();
9878 create_req.id = Some(vec!["level1".to_string(), "level2".to_string()]);
9879 namespace.create_namespace(create_req).await.unwrap();
9880
9881 let mut create_req = CreateNamespaceRequest::new();
9882 create_req.id = Some(vec![
9883 "level1".to_string(),
9884 "level2".to_string(),
9885 "level3".to_string(),
9886 ]);
9887 namespace.create_namespace(create_req).await.unwrap();
9888
9889 let schema = create_test_schema();
9891 let ipc_data = create_test_ipc_data(&schema);
9892 let mut create_table_req = CreateTableRequest::new();
9893 create_table_req.id = Some(vec![
9894 "level1".to_string(),
9895 "level2".to_string(),
9896 "level3".to_string(),
9897 "table1".to_string(),
9898 ]);
9899 let result = namespace
9900 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9901 .await;
9902 assert!(
9903 result.is_ok(),
9904 "Failed to create table in deeply nested namespace"
9905 );
9906
9907 let mut exists_req = TableExistsRequest::new();
9909 exists_req.id = Some(vec![
9910 "level1".to_string(),
9911 "level2".to_string(),
9912 "level3".to_string(),
9913 "table1".to_string(),
9914 ]);
9915 let result = namespace.table_exists(exists_req).await;
9916 assert!(result.is_ok());
9917 }
9918
9919 #[tokio::test]
9920 async fn test_namespace_with_properties() {
9921 let (namespace, _temp_dir) = create_test_namespace().await;
9922
9923 let mut properties = HashMap::new();
9925 properties.insert("owner".to_string(), "test_user".to_string());
9926 properties.insert("description".to_string(), "Test namespace".to_string());
9927
9928 let mut create_req = CreateNamespaceRequest::new();
9929 create_req.id = Some(vec!["test_ns".to_string()]);
9930 create_req.properties = Some(properties.clone());
9931 namespace.create_namespace(create_req).await.unwrap();
9932
9933 let describe_req = DescribeNamespaceRequest {
9935 id: Some(vec!["test_ns".to_string()]),
9936 ..Default::default()
9937 };
9938 let result = namespace.describe_namespace(describe_req).await;
9939 assert!(result.is_ok());
9940 let response = result.unwrap();
9941 assert!(response.properties.is_some());
9942 let props = response.properties.unwrap();
9943 assert_eq!(props.get("owner"), Some(&"test_user".to_string()));
9944 assert_eq!(
9945 props.get("description"),
9946 Some(&"Test namespace".to_string())
9947 );
9948 }
9949
9950 #[tokio::test]
9951 async fn test_cannot_drop_namespace_with_tables() {
9952 let (namespace, _temp_dir) = create_test_namespace().await;
9953
9954 let mut create_ns_req = CreateNamespaceRequest::new();
9956 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9957 namespace.create_namespace(create_ns_req).await.unwrap();
9958
9959 let schema = create_test_schema();
9961 let ipc_data = create_test_ipc_data(&schema);
9962 let mut create_table_req = CreateTableRequest::new();
9963 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9964 namespace
9965 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9966 .await
9967 .unwrap();
9968
9969 let mut drop_req = DropNamespaceRequest::new();
9971 drop_req.id = Some(vec!["test_ns".to_string()]);
9972 let result = namespace.drop_namespace(drop_req).await;
9973 assert!(
9974 result.is_err(),
9975 "Should not be able to drop namespace with tables"
9976 );
9977 }
9978
9979 #[tokio::test]
9980 async fn test_isolation_between_namespaces() {
9981 let (namespace, _temp_dir) = create_test_namespace().await;
9982
9983 let mut create_req = CreateNamespaceRequest::new();
9985 create_req.id = Some(vec!["ns1".to_string()]);
9986 namespace.create_namespace(create_req).await.unwrap();
9987
9988 let mut create_req = CreateNamespaceRequest::new();
9989 create_req.id = Some(vec!["ns2".to_string()]);
9990 namespace.create_namespace(create_req).await.unwrap();
9991
9992 let schema = create_test_schema();
9994 let ipc_data = create_test_ipc_data(&schema);
9995
9996 let mut create_table_req = CreateTableRequest::new();
9997 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
9998 namespace
9999 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
10000 .await
10001 .unwrap();
10002
10003 let mut create_table_req = CreateTableRequest::new();
10004 create_table_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
10005 namespace
10006 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10007 .await
10008 .unwrap();
10009
10010 let list_req = ListTablesRequest {
10012 id: Some(vec!["ns1".to_string()]),
10013 page_token: None,
10014 limit: None,
10015 ..Default::default()
10016 };
10017 let result = namespace.list_tables(list_req).await.unwrap();
10018 assert_eq!(result.tables.len(), 1);
10019 assert_eq!(result.tables[0], "table1");
10020
10021 let list_req = ListTablesRequest {
10022 id: Some(vec!["ns2".to_string()]),
10023 page_token: None,
10024 limit: None,
10025 ..Default::default()
10026 };
10027 let result = namespace.list_tables(list_req).await.unwrap();
10028 assert_eq!(result.tables.len(), 1);
10029 assert_eq!(result.tables[0], "table1");
10030
10031 let mut drop_req = DropTableRequest::new();
10033 drop_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10034 namespace.drop_table(drop_req).await.unwrap();
10035
10036 let mut exists_req = TableExistsRequest::new();
10038 exists_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10039 assert!(namespace.table_exists(exists_req).await.is_err());
10040
10041 let mut exists_req = TableExistsRequest::new();
10042 exists_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
10043 assert!(namespace.table_exists(exists_req).await.is_ok());
10044 }
10045
10046 #[tokio::test]
10047 async fn test_migrate_directory_tables() {
10048 let temp_dir = TempStdDir::default();
10049 let temp_path = temp_dir.to_str().unwrap();
10050
10051 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
10053 .manifest_enabled(false)
10054 .dir_listing_enabled(true)
10055 .build()
10056 .await
10057 .unwrap();
10058
10059 let schema = create_test_schema();
10061 let ipc_data = create_test_ipc_data(&schema);
10062
10063 for i in 1..=3 {
10064 let mut create_req = CreateTableRequest::new();
10065 create_req.id = Some(vec![format!("table{}", i)]);
10066 dir_only_ns
10067 .create_table(create_req, bytes::Bytes::from(ipc_data.clone()))
10068 .await
10069 .unwrap();
10070 }
10071
10072 drop(dir_only_ns);
10073
10074 let dual_mode_ns = DirectoryNamespaceBuilder::new(temp_path)
10076 .manifest_enabled(true)
10077 .dir_listing_enabled(true)
10078 .build()
10079 .await
10080 .unwrap();
10081
10082 let mut list_req = ListTablesRequest::new();
10084 list_req.id = Some(vec![]);
10085 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
10086 assert_eq!(tables.len(), 3);
10087
10088 let migrated_count = dual_mode_ns.migrate().await.unwrap();
10090 assert_eq!(migrated_count, 3, "Should migrate all 3 tables");
10091
10092 let mut list_req = ListTablesRequest::new();
10094 list_req.id = Some(vec![]);
10095 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
10096 assert_eq!(tables.len(), 3);
10097
10098 let migrated_count = dual_mode_ns.migrate().await.unwrap();
10100 assert_eq!(
10101 migrated_count, 0,
10102 "Should not migrate already-migrated tables"
10103 );
10104
10105 drop(dual_mode_ns);
10106
10107 let manifest_only_ns = DirectoryNamespaceBuilder::new(temp_path)
10109 .manifest_enabled(true)
10110 .dir_listing_enabled(false)
10111 .build()
10112 .await
10113 .unwrap();
10114
10115 let mut list_req = ListTablesRequest::new();
10117 list_req.id = Some(vec![]);
10118 let tables = manifest_only_ns.list_tables(list_req).await.unwrap().tables;
10119 assert_eq!(tables.len(), 3);
10120 assert!(tables.contains(&"table1".to_string()));
10121 assert!(tables.contains(&"table2".to_string()));
10122 assert!(tables.contains(&"table3".to_string()));
10123 }
10124
10125 #[tokio::test]
10126 async fn test_migrate_without_manifest() {
10127 let temp_dir = TempStdDir::default();
10128 let temp_path = temp_dir.to_str().unwrap();
10129
10130 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10132 .manifest_enabled(false)
10133 .dir_listing_enabled(true)
10134 .build()
10135 .await
10136 .unwrap();
10137
10138 let migrated_count = namespace.migrate().await.unwrap();
10140 assert_eq!(migrated_count, 0);
10141 }
10142
10143 #[tokio::test]
10144 async fn test_register_table() {
10145 use lance_namespace::models::{RegisterTableRequest, TableExistsRequest};
10146
10147 let temp_dir = TempStdDir::default();
10148 let temp_path = temp_dir.to_str().unwrap();
10149
10150 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10151 .dir_listing_to_manifest_migration_enabled(true)
10152 .build()
10153 .await
10154 .unwrap();
10155
10156 let schema = create_test_schema();
10158 let ipc_data = create_test_ipc_data(&schema);
10159
10160 let table_uri = format!("{}/external_table.lance", temp_path);
10161 let cursor = Cursor::new(ipc_data);
10162 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
10163 let batches: Vec<_> = stream_reader
10164 .collect::<std::result::Result<Vec<_>, _>>()
10165 .unwrap();
10166 let schema = batches[0].schema();
10167 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
10168 let reader = RecordBatchIterator::new(batch_results, schema);
10169 Dataset::write(Box::new(reader), &table_uri, None)
10170 .await
10171 .unwrap();
10172
10173 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
10175 register_req.id = Some(vec!["registered_table".to_string()]);
10176
10177 let response = namespace.register_table(register_req).await.unwrap();
10178 assert_eq!(response.location, Some("external_table.lance".to_string()));
10179
10180 let mut exists_req = TableExistsRequest::new();
10182 exists_req.id = Some(vec!["registered_table".to_string()]);
10183 assert!(namespace.table_exists(exists_req).await.is_ok());
10184
10185 let mut list_req = ListTablesRequest::new();
10187 list_req.id = Some(vec![]);
10188 let tables = namespace.list_tables(list_req).await.unwrap();
10189 assert!(tables.tables.contains(&"registered_table".to_string()));
10190 }
10191
10192 #[tokio::test]
10193 async fn test_register_table_duplicate_fails() {
10194 use lance_namespace::models::RegisterTableRequest;
10195
10196 let temp_dir = TempStdDir::default();
10197 let temp_path = temp_dir.to_str().unwrap();
10198
10199 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10200 .build()
10201 .await
10202 .unwrap();
10203
10204 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10206 register_req.id = Some(vec!["test_table".to_string()]);
10207
10208 namespace
10209 .register_table(register_req.clone())
10210 .await
10211 .unwrap();
10212
10213 let result = namespace.register_table(register_req).await;
10215 assert!(result.is_err());
10216 assert!(result.unwrap_err().to_string().contains("already exists"));
10217 }
10218
10219 #[tokio::test]
10220 async fn test_deregister_table() {
10221 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
10222
10223 let temp_dir = TempStdDir::default();
10224 let temp_path = temp_dir.to_str().unwrap();
10225
10226 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10229 .manifest_enabled(true)
10230 .dir_listing_enabled(false)
10231 .build()
10232 .await
10233 .unwrap();
10234
10235 let schema = create_test_schema();
10237 let ipc_data = create_test_ipc_data(&schema);
10238
10239 let mut create_req = CreateTableRequest::new();
10240 create_req.id = Some(vec!["test_table".to_string()]);
10241 namespace
10242 .create_table(create_req, bytes::Bytes::from(ipc_data))
10243 .await
10244 .unwrap();
10245
10246 let mut exists_req = TableExistsRequest::new();
10248 exists_req.id = Some(vec!["test_table".to_string()]);
10249 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
10250
10251 let mut deregister_req = DeregisterTableRequest::new();
10253 deregister_req.id = Some(vec!["test_table".to_string()]);
10254 let response = namespace.deregister_table(deregister_req).await.unwrap();
10255
10256 assert!(
10258 response.location.is_some(),
10259 "Deregister should return location"
10260 );
10261 let location = response.location.as_ref().unwrap();
10262 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10265 .expect("Failed to convert temp path to URL");
10266 let expected_prefix = expected_url.to_string();
10267 assert!(
10268 location.starts_with(&expected_prefix),
10269 "Location should start with '{}', got: {}",
10270 expected_prefix,
10271 location
10272 );
10273 assert!(
10274 location.contains("test_table"),
10275 "Location should contain table name: {}",
10276 location
10277 );
10278 assert_eq!(response.id, Some(vec!["test_table".to_string()]));
10279
10280 assert!(namespace.table_exists(exists_req).await.is_err());
10282
10283 let dataset = Dataset::open(location).await;
10285 assert!(
10286 dataset.is_ok(),
10287 "Physical table data should still exist at {}",
10288 location
10289 );
10290 }
10291
10292 #[tokio::test]
10293 async fn test_deregister_table_in_child_namespace() {
10294 use lance_namespace::models::{
10295 CreateNamespaceRequest, DeregisterTableRequest, TableExistsRequest,
10296 };
10297
10298 let temp_dir = TempStdDir::default();
10299 let temp_path = temp_dir.to_str().unwrap();
10300
10301 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10302 .build()
10303 .await
10304 .unwrap();
10305
10306 let mut create_ns_req = CreateNamespaceRequest::new();
10308 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10309 namespace.create_namespace(create_ns_req).await.unwrap();
10310
10311 let schema = create_test_schema();
10313 let ipc_data = create_test_ipc_data(&schema);
10314
10315 let mut create_req = CreateTableRequest::new();
10316 create_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10317 namespace
10318 .create_table(create_req, bytes::Bytes::from(ipc_data))
10319 .await
10320 .unwrap();
10321
10322 let mut deregister_req = DeregisterTableRequest::new();
10324 deregister_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10325 let response = namespace.deregister_table(deregister_req).await.unwrap();
10326
10327 assert!(
10329 response.location.is_some(),
10330 "Deregister should return location"
10331 );
10332 let location = response.location.as_ref().unwrap();
10333 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10336 .expect("Failed to convert temp path to URL");
10337 let expected_prefix = expected_url.to_string();
10338 assert!(
10339 location.starts_with(&expected_prefix),
10340 "Location should start with '{}', got: {}",
10341 expected_prefix,
10342 location
10343 );
10344 assert!(
10345 location.contains("test_ns") && location.contains("test_table"),
10346 "Location should contain namespace and table name: {}",
10347 location
10348 );
10349 assert_eq!(
10350 response.id,
10351 Some(vec!["test_ns".to_string(), "test_table".to_string()])
10352 );
10353
10354 let mut exists_req = TableExistsRequest::new();
10356 exists_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10357 assert!(namespace.table_exists(exists_req).await.is_err());
10358 }
10359
10360 #[tokio::test]
10361 async fn test_register_without_manifest_fails() {
10362 use lance_namespace::models::RegisterTableRequest;
10363
10364 let temp_dir = TempStdDir::default();
10365 let temp_path = temp_dir.to_str().unwrap();
10366
10367 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10369 .manifest_enabled(false)
10370 .build()
10371 .await
10372 .unwrap();
10373
10374 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10376 register_req.id = Some(vec!["test_table".to_string()]);
10377 let result = namespace.register_table(register_req).await;
10378 assert!(result.is_err());
10379 assert!(
10380 result
10381 .unwrap_err()
10382 .to_string()
10383 .contains("manifest mode is enabled")
10384 );
10385
10386 }
10389
10390 #[tokio::test]
10391 async fn test_register_table_rejects_absolute_uri() {
10392 use lance_namespace::models::RegisterTableRequest;
10393
10394 let temp_dir = TempStdDir::default();
10395 let temp_path = temp_dir.to_str().unwrap();
10396
10397 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10398 .build()
10399 .await
10400 .unwrap();
10401
10402 let mut register_req = RegisterTableRequest::new("s3://bucket/table.lance".to_string());
10404 register_req.id = Some(vec!["test_table".to_string()]);
10405 let result = namespace.register_table(register_req).await;
10406 assert!(result.is_err());
10407 let err_msg = result.unwrap_err().to_string();
10408 assert!(err_msg.contains("Absolute URIs are not allowed"));
10409 }
10410
10411 #[tokio::test]
10412 async fn test_register_table_rejects_absolute_path() {
10413 use lance_namespace::models::RegisterTableRequest;
10414
10415 let temp_dir = TempStdDir::default();
10416 let temp_path = temp_dir.to_str().unwrap();
10417
10418 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10419 .build()
10420 .await
10421 .unwrap();
10422
10423 let mut register_req = RegisterTableRequest::new("/tmp/table.lance".to_string());
10425 register_req.id = Some(vec!["test_table".to_string()]);
10426 let result = namespace.register_table(register_req).await;
10427 assert!(result.is_err());
10428 let err_msg = result.unwrap_err().to_string();
10429 assert!(err_msg.contains("Absolute paths are not allowed"));
10430 }
10431
10432 #[tokio::test]
10433 async fn test_register_table_rejects_path_traversal() {
10434 use lance_namespace::models::RegisterTableRequest;
10435
10436 let temp_dir = TempStdDir::default();
10437 let temp_path = temp_dir.to_str().unwrap();
10438
10439 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10440 .build()
10441 .await
10442 .unwrap();
10443
10444 let mut register_req = RegisterTableRequest::new("../outside/table.lance".to_string());
10446 register_req.id = Some(vec!["test_table".to_string()]);
10447 let result = namespace.register_table(register_req).await;
10448 assert!(result.is_err());
10449 let err_msg = result.unwrap_err().to_string();
10450 assert!(err_msg.contains("Path traversal is not allowed"));
10451 }
10452
10453 #[tokio::test]
10454 async fn test_namespace_write() {
10455 use arrow::array::Int32Array;
10456 use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
10457 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
10458 use lance::dataset::{Dataset, WriteMode, WriteParams};
10459 use lance_namespace::LanceNamespace;
10460
10461 let (namespace, _temp_dir) = create_test_namespace().await;
10462 let namespace = Arc::new(namespace) as Arc<dyn LanceNamespace>;
10463
10464 let table_id = vec!["test_ns".to_string(), "test_table".to_string()];
10466 let schema = Arc::new(ArrowSchema::new(vec![
10467 ArrowField::new("a", DataType::Int32, false),
10468 ArrowField::new("b", DataType::Int32, false),
10469 ]));
10470
10471 let data1 = RecordBatch::try_new(
10473 schema.clone(),
10474 vec![
10475 Arc::new(Int32Array::from(vec![1, 2, 3])),
10476 Arc::new(Int32Array::from(vec![10, 20, 30])),
10477 ],
10478 )
10479 .unwrap();
10480
10481 let reader1 = RecordBatchIterator::new(vec![data1].into_iter().map(Ok), schema.clone());
10482 let dataset =
10483 Dataset::write_into_namespace(reader1, namespace.clone(), table_id.clone(), None)
10484 .await
10485 .unwrap();
10486
10487 assert_eq!(dataset.count_rows(None).await.unwrap(), 3);
10488 assert_eq!(dataset.version().version, 1);
10489
10490 let data2 = RecordBatch::try_new(
10492 schema.clone(),
10493 vec![
10494 Arc::new(Int32Array::from(vec![4, 5])),
10495 Arc::new(Int32Array::from(vec![40, 50])),
10496 ],
10497 )
10498 .unwrap();
10499
10500 let params_append = WriteParams {
10501 mode: WriteMode::Append,
10502 ..Default::default()
10503 };
10504
10505 let reader2 = RecordBatchIterator::new(vec![data2].into_iter().map(Ok), schema.clone());
10506 let dataset = Dataset::write_into_namespace(
10507 reader2,
10508 namespace.clone(),
10509 table_id.clone(),
10510 Some(params_append),
10511 )
10512 .await
10513 .unwrap();
10514
10515 assert_eq!(dataset.count_rows(None).await.unwrap(), 5);
10516 assert_eq!(dataset.version().version, 2);
10517
10518 let data3 = RecordBatch::try_new(
10520 schema.clone(),
10521 vec![
10522 Arc::new(Int32Array::from(vec![100, 200])),
10523 Arc::new(Int32Array::from(vec![1000, 2000])),
10524 ],
10525 )
10526 .unwrap();
10527
10528 let params_overwrite = WriteParams {
10529 mode: WriteMode::Overwrite,
10530 ..Default::default()
10531 };
10532
10533 let reader3 = RecordBatchIterator::new(vec![data3].into_iter().map(Ok), schema.clone());
10534 let dataset = Dataset::write_into_namespace(
10535 reader3,
10536 namespace.clone(),
10537 table_id.clone(),
10538 Some(params_overwrite),
10539 )
10540 .await
10541 .unwrap();
10542
10543 assert_eq!(dataset.count_rows(None).await.unwrap(), 2);
10544 assert_eq!(dataset.version().version, 3);
10545
10546 let result = dataset.scan().try_into_batch().await.unwrap();
10548 let a_col = result
10549 .column_by_name("a")
10550 .unwrap()
10551 .as_any()
10552 .downcast_ref::<Int32Array>()
10553 .unwrap();
10554 assert_eq!(a_col.values(), &[100, 200]);
10555 }
10556
10557 #[tokio::test]
10562 async fn test_declare_table_v1_mode() {
10563 use lance_namespace::models::{
10564 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
10565 };
10566
10567 let temp_dir = TempStdDir::default();
10568 let temp_path = temp_dir.to_str().unwrap();
10569
10570 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10572 .manifest_enabled(false)
10573 .build()
10574 .await
10575 .unwrap();
10576
10577 let mut declare_req = DeclareTableRequest::new();
10579 declare_req.id = Some(vec!["test_table".to_string()]);
10580 let response = namespace.declare_table(declare_req).await.unwrap();
10581
10582 assert!(response.location.is_some());
10584 let location = response.location.as_ref().unwrap();
10585 assert!(location.ends_with("test_table.lance"));
10586
10587 let mut exists_req = TableExistsRequest::new();
10589 exists_req.id = Some(vec!["test_table".to_string()]);
10590 assert!(namespace.table_exists(exists_req).await.is_ok());
10591
10592 let mut describe_req = DescribeTableRequest::new();
10594 describe_req.id = Some(vec!["test_table".to_string()]);
10595 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10596 assert!(describe_response.location.is_some());
10597 assert!(describe_response.version.is_none()); assert!(describe_response.schema.is_none()); assert_eq!(describe_response.is_only_declared, None);
10600
10601 let mut describe_req = DescribeTableRequest::new();
10602 describe_req.id = Some(vec!["test_table".to_string()]);
10603 describe_req.check_declared = Some(true);
10604 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10605 assert_eq!(describe_response.is_only_declared, Some(true));
10606
10607 let mut list_req = ListTablesRequest::new();
10608 list_req.id = Some(vec![]);
10609 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
10610 assert_eq!(list_response.tables, vec!["test_table".to_string()]);
10611
10612 list_req.include_declared = Some(false);
10613 let list_response = namespace.list_tables(list_req).await.unwrap();
10614 assert!(list_response.tables.is_empty());
10615 }
10616
10617 #[tokio::test]
10618 async fn test_insert_into_declared_table_promotes_it_from_declared_state() {
10619 use lance_namespace::models::{
10620 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest,
10621 };
10622
10623 let temp_dir = TempStdDir::default();
10624 let temp_path = temp_dir.to_str().unwrap();
10625
10626 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10627 .manifest_enabled(false)
10628 .build()
10629 .await
10630 .unwrap();
10631
10632 let mut declare_req = DeclareTableRequest::new();
10633 declare_req.id = Some(vec!["test_table".to_string()]);
10634 namespace.declare_table(declare_req).await.unwrap();
10635
10636 let schema = create_test_schema();
10637 let ipc_data = create_test_ipc_data(&schema);
10638 let mut insert_req = InsertIntoTableRequest::new();
10639 insert_req.id = Some(vec!["test_table".to_string()]);
10640 namespace
10641 .insert_into_table(insert_req, bytes::Bytes::from(ipc_data))
10642 .await
10643 .unwrap();
10644
10645 let mut describe_req = DescribeTableRequest::new();
10646 describe_req.id = Some(vec!["test_table".to_string()]);
10647 describe_req.load_detailed_metadata = Some(true);
10648 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10649
10650 assert_eq!(describe_response.is_only_declared, Some(false));
10651 assert_eq!(describe_response.version, Some(1));
10652 assert!(describe_response.schema.is_some());
10653
10654 let mut list_req = ListTablesRequest::new();
10655 list_req.id = Some(vec![]);
10656 list_req.include_declared = Some(false);
10657 assert_eq!(
10658 namespace.list_tables(list_req).await.unwrap().tables,
10659 vec!["test_table".to_string()]
10660 );
10661 }
10662
10663 #[tokio::test]
10664 async fn test_create_table_after_declare_table_v1_mode_creates_table() {
10665 use lance_namespace::models::{
10666 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
10667 };
10668
10669 let temp_dir = TempStdDir::default();
10670 let temp_path = temp_dir.to_str().unwrap();
10671
10672 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10673 .manifest_enabled(false)
10674 .build()
10675 .await
10676 .unwrap();
10677
10678 let mut declare_req = DeclareTableRequest::new();
10679 declare_req.id = Some(vec!["test_table".to_string()]);
10680 namespace.declare_table(declare_req).await.unwrap();
10681
10682 let mut create_req = CreateTableRequest::new();
10683 create_req.id = Some(vec!["test_table".to_string()]);
10684 let response = namespace
10685 .create_table(
10686 create_req,
10687 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10688 )
10689 .await
10690 .unwrap();
10691
10692 assert_eq!(response.version, Some(1));
10693
10694 let mut describe_req = DescribeTableRequest::new();
10695 describe_req.id = Some(vec!["test_table".to_string()]);
10696 describe_req.load_detailed_metadata = Some(true);
10697 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10698 assert_eq!(describe_response.is_only_declared, Some(false));
10699 assert_eq!(describe_response.version, Some(1));
10700
10701 let mut list_req = ListTablesRequest::new();
10702 list_req.id = Some(vec![]);
10703 list_req.include_declared = Some(false);
10704 assert_eq!(
10705 namespace.list_tables(list_req).await.unwrap().tables,
10706 vec!["test_table".to_string()]
10707 );
10708 }
10709
10710 #[tokio::test]
10711 async fn test_insert_into_declared_table_with_manifest_promotes_it() {
10712 use lance_namespace::models::{
10713 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest, ListTablesRequest,
10714 };
10715
10716 let temp_dir = TempStdDir::default();
10717 let temp_path = temp_dir.to_str().unwrap();
10718
10719 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10720 .manifest_enabled(true)
10721 .dir_listing_enabled(false)
10722 .build()
10723 .await
10724 .unwrap();
10725
10726 let mut declare_req = DeclareTableRequest::new();
10727 declare_req.id = Some(vec!["test_table".to_string()]);
10728 namespace.declare_table(declare_req).await.unwrap();
10729
10730 let mut insert_req = InsertIntoTableRequest::new();
10731 insert_req.id = Some(vec!["test_table".to_string()]);
10732 namespace
10733 .insert_into_table(
10734 insert_req,
10735 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10736 )
10737 .await
10738 .unwrap();
10739
10740 let mut describe_req = DescribeTableRequest::new();
10741 describe_req.id = Some(vec!["test_table".to_string()]);
10742 describe_req.load_detailed_metadata = Some(true);
10743 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10744 assert_eq!(describe_response.is_only_declared, Some(false));
10745 assert_eq!(describe_response.version, Some(1));
10746
10747 let mut list_req = ListTablesRequest::new();
10748 list_req.id = Some(vec![]);
10749 list_req.include_declared = Some(false);
10750 assert_eq!(
10751 namespace.list_tables(list_req).await.unwrap().tables,
10752 vec!["test_table".to_string()]
10753 );
10754 }
10755
10756 #[tokio::test]
10757 async fn test_create_table_after_declare_table_with_manifest_creates_table() {
10758 use lance_namespace::models::{
10759 CreateTableRequest, DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
10760 };
10761
10762 let temp_dir = TempStdDir::default();
10763 let temp_path = temp_dir.to_str().unwrap();
10764
10765 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10766 .manifest_enabled(true)
10767 .dir_listing_enabled(false)
10768 .build()
10769 .await
10770 .unwrap();
10771
10772 let mut declare_req = DeclareTableRequest::new();
10773 declare_req.id = Some(vec!["test_table".to_string()]);
10774 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10775 namespace.declare_table(declare_req).await.unwrap();
10776
10777 let mut create_req = CreateTableRequest::new();
10778 create_req.id = Some(vec!["test_table".to_string()]);
10779 create_req.mode = Some("Overwrite".to_string());
10780 let response = namespace
10781 .create_table(
10782 create_req,
10783 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10784 )
10785 .await
10786 .unwrap();
10787
10788 assert_eq!(response.version, Some(1));
10789 assert_eq!(
10790 response
10791 .properties
10792 .as_ref()
10793 .and_then(|properties| properties.get("owner")),
10794 Some(&"alice".to_string())
10795 );
10796
10797 let mut describe_req = DescribeTableRequest::new();
10798 describe_req.id = Some(vec!["test_table".to_string()]);
10799 describe_req.load_detailed_metadata = Some(true);
10800 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10801 assert_eq!(describe_response.is_only_declared, Some(false));
10802 assert_eq!(describe_response.version, Some(1));
10803 assert_eq!(
10804 describe_response
10805 .properties
10806 .as_ref()
10807 .and_then(|properties| properties.get("owner")),
10808 Some(&"alice".to_string())
10809 );
10810
10811 let mut list_req = ListTablesRequest::new();
10812 list_req.id = Some(vec![]);
10813 list_req.include_declared = Some(false);
10814 assert_eq!(
10815 namespace.list_tables(list_req).await.unwrap().tables,
10816 vec!["test_table".to_string()]
10817 );
10818 }
10819
10820 #[tokio::test]
10821 async fn test_create_table_after_declare_table_with_manifest_rejects_new_properties() {
10822 use lance_namespace::models::{CreateTableRequest, DeclareTableRequest};
10823
10824 let temp_dir = TempStdDir::default();
10825 let temp_path = temp_dir.to_str().unwrap();
10826
10827 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10828 .manifest_enabled(true)
10829 .dir_listing_enabled(false)
10830 .build()
10831 .await
10832 .unwrap();
10833
10834 let mut declare_req = DeclareTableRequest::new();
10835 declare_req.id = Some(vec!["test_table".to_string()]);
10836 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10837 namespace.declare_table(declare_req).await.unwrap();
10838
10839 let mut create_req = CreateTableRequest::new();
10840 create_req.id = Some(vec!["test_table".to_string()]);
10841 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
10842
10843 let result = namespace
10844 .create_table(
10845 create_req,
10846 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10847 )
10848 .await;
10849
10850 assert!(result.is_err());
10851 assert!(
10852 result
10853 .unwrap_err()
10854 .to_string()
10855 .contains("cannot set properties for already declared table")
10856 );
10857 }
10858
10859 #[tokio::test]
10860 async fn test_create_table_with_manifest_exist_ok_keeps_existing_table() {
10861 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
10862
10863 let temp_dir = TempStdDir::default();
10864 let temp_path = temp_dir.to_str().unwrap();
10865
10866 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10867 .manifest_enabled(true)
10868 .dir_listing_enabled(false)
10869 .build()
10870 .await
10871 .unwrap();
10872
10873 let mut create_req = CreateTableRequest::new();
10874 create_req.id = Some(vec!["test_table".to_string()]);
10875 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10876 namespace
10877 .create_table(
10878 create_req,
10879 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10880 )
10881 .await
10882 .unwrap();
10883
10884 let mut create_req = CreateTableRequest::new();
10885 create_req.id = Some(vec!["test_table".to_string()]);
10886 create_req.mode = Some("ExistOk".to_string());
10887 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
10888 let response = namespace
10889 .create_table(
10890 create_req,
10891 bytes::Bytes::from(create_single_row_test_ipc_data()),
10892 )
10893 .await
10894 .unwrap();
10895
10896 assert_eq!(
10897 response
10898 .properties
10899 .as_ref()
10900 .and_then(|properties| properties.get("owner")),
10901 Some(&"alice".to_string())
10902 );
10903 assert_eq!(
10904 open_dataset(&namespace, "test_table")
10905 .await
10906 .count_rows(None)
10907 .await
10908 .unwrap(),
10909 2
10910 );
10911
10912 let mut describe_req = DescribeTableRequest::new();
10913 describe_req.id = Some(vec!["test_table".to_string()]);
10914 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10915 assert_eq!(
10916 describe_response
10917 .properties
10918 .as_ref()
10919 .and_then(|properties| properties.get("owner")),
10920 Some(&"alice".to_string())
10921 );
10922 }
10923
10924 #[tokio::test]
10925 async fn test_create_table_with_manifest_overwrite_replaces_existing_table() {
10926 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
10927
10928 let temp_dir = TempStdDir::default();
10929 let temp_path = temp_dir.to_str().unwrap();
10930
10931 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10932 .manifest_enabled(true)
10933 .dir_listing_enabled(false)
10934 .build()
10935 .await
10936 .unwrap();
10937
10938 let mut create_req = CreateTableRequest::new();
10939 create_req.id = Some(vec!["test_table".to_string()]);
10940 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10941 namespace
10942 .create_table(
10943 create_req,
10944 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10945 )
10946 .await
10947 .unwrap();
10948
10949 let mut create_req = CreateTableRequest::new();
10950 create_req.id = Some(vec!["test_table".to_string()]);
10951 create_req.mode = Some("overwrite".to_string());
10952 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
10953 let response = namespace
10954 .create_table(
10955 create_req,
10956 bytes::Bytes::from(create_single_row_test_ipc_data()),
10957 )
10958 .await
10959 .unwrap();
10960
10961 assert_eq!(response.version, Some(2));
10962 assert_eq!(
10963 response
10964 .properties
10965 .as_ref()
10966 .and_then(|properties| properties.get("owner")),
10967 Some(&"bob".to_string())
10968 );
10969 assert_eq!(
10970 open_dataset(&namespace, "test_table")
10971 .await
10972 .count_rows(None)
10973 .await
10974 .unwrap(),
10975 1
10976 );
10977
10978 let mut describe_req = DescribeTableRequest::new();
10979 describe_req.id = Some(vec!["test_table".to_string()]);
10980 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10981 assert_eq!(
10982 describe_response
10983 .properties
10984 .as_ref()
10985 .and_then(|properties| properties.get("owner")),
10986 Some(&"bob".to_string())
10987 );
10988 }
10989
10990 #[tokio::test]
10991 async fn test_create_table_with_manifest_invalid_mode_rejected() {
10992 use lance_namespace::models::CreateTableRequest;
10993
10994 let temp_dir = TempStdDir::default();
10995 let temp_path = temp_dir.to_str().unwrap();
10996
10997 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10998 .manifest_enabled(true)
10999 .dir_listing_enabled(false)
11000 .build()
11001 .await
11002 .unwrap();
11003
11004 let mut create_req = CreateTableRequest::new();
11005 create_req.id = Some(vec!["test_table".to_string()]);
11006 create_req.mode = Some("append".to_string());
11007 let result = namespace
11008 .create_table(
11009 create_req,
11010 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11011 )
11012 .await;
11013
11014 assert!(result.is_err());
11015 assert!(
11016 result
11017 .unwrap_err()
11018 .to_string()
11019 .contains("Unsupported create_table mode")
11020 );
11021 }
11022
11023 #[tokio::test]
11024 async fn test_merge_insert_into_declared_table_v1_mode_creates_table() {
11025 use lance_namespace::models::{
11026 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11027 MergeInsertIntoTableRequest,
11028 };
11029
11030 let temp_dir = TempStdDir::default();
11031 let temp_path = temp_dir.to_str().unwrap();
11032
11033 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11034 .manifest_enabled(false)
11035 .build()
11036 .await
11037 .unwrap();
11038
11039 let mut declare_req = DeclareTableRequest::new();
11040 declare_req.id = Some(vec!["test_table".to_string()]);
11041 namespace.declare_table(declare_req).await.unwrap();
11042
11043 let mut merge_req = MergeInsertIntoTableRequest::new();
11044 merge_req.id = Some(vec!["test_table".to_string()]);
11045 merge_req.on = Some("id".to_string());
11046 let response = namespace
11047 .merge_insert_into_table(
11048 merge_req,
11049 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11050 )
11051 .await
11052 .unwrap();
11053
11054 assert_eq!(response.num_inserted_rows, Some(2));
11055 assert_eq!(response.num_updated_rows, Some(0));
11056
11057 let mut describe_req = DescribeTableRequest::new();
11058 describe_req.id = Some(vec!["test_table".to_string()]);
11059 describe_req.load_detailed_metadata = Some(true);
11060 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11061 assert_eq!(describe_response.is_only_declared, Some(false));
11062 assert_eq!(describe_response.version, Some(1));
11063
11064 let mut list_req = ListTablesRequest::new();
11065 list_req.id = Some(vec![]);
11066 list_req.include_declared = Some(false);
11067 assert_eq!(
11068 namespace.list_tables(list_req).await.unwrap().tables,
11069 vec!["test_table".to_string()]
11070 );
11071 }
11072
11073 #[tokio::test]
11074 async fn test_merge_insert_into_declared_table_with_manifest_creates_table() {
11075 use lance_namespace::models::{
11076 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11077 MergeInsertIntoTableRequest,
11078 };
11079
11080 let temp_dir = TempStdDir::default();
11081 let temp_path = temp_dir.to_str().unwrap();
11082
11083 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11084 .manifest_enabled(true)
11085 .dir_listing_enabled(false)
11086 .build()
11087 .await
11088 .unwrap();
11089
11090 let mut declare_req = DeclareTableRequest::new();
11091 declare_req.id = Some(vec!["test_table".to_string()]);
11092 namespace.declare_table(declare_req).await.unwrap();
11093
11094 let mut merge_req = MergeInsertIntoTableRequest::new();
11095 merge_req.id = Some(vec!["test_table".to_string()]);
11096 merge_req.on = Some("id".to_string());
11097 let response = namespace
11098 .merge_insert_into_table(
11099 merge_req,
11100 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11101 )
11102 .await
11103 .unwrap();
11104
11105 assert_eq!(response.num_inserted_rows, Some(2));
11106 assert_eq!(response.num_updated_rows, Some(0));
11107
11108 let mut describe_req = DescribeTableRequest::new();
11109 describe_req.id = Some(vec!["test_table".to_string()]);
11110 describe_req.load_detailed_metadata = Some(true);
11111 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11112 assert_eq!(describe_response.is_only_declared, Some(false));
11113 assert_eq!(describe_response.version, Some(1));
11114
11115 let mut list_req = ListTablesRequest::new();
11116 list_req.id = Some(vec![]);
11117 list_req.include_declared = Some(false);
11118 assert_eq!(
11119 namespace.list_tables(list_req).await.unwrap().tables,
11120 vec!["test_table".to_string()]
11121 );
11122 }
11123
11124 #[tokio::test]
11125 async fn test_declare_table_with_manifest() {
11126 use lance_namespace::models::{
11127 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
11128 };
11129
11130 let temp_dir = TempStdDir::default();
11131 let temp_path = temp_dir.to_str().unwrap();
11132
11133 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11135 .manifest_enabled(true)
11136 .dir_listing_enabled(false)
11137 .build()
11138 .await
11139 .unwrap();
11140
11141 let mut declare_req = DeclareTableRequest::new();
11143 declare_req.id = Some(vec!["test_table".to_string()]);
11144 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11145 let response = namespace.declare_table(declare_req).await.unwrap();
11146
11147 assert!(response.location.is_some());
11149 assert_eq!(
11150 response
11151 .properties
11152 .as_ref()
11153 .and_then(|properties| properties.get("owner")),
11154 Some(&"alice".to_string())
11155 );
11156
11157 let mut exists_req = TableExistsRequest::new();
11159 exists_req.id = Some(vec!["test_table".to_string()]);
11160 assert!(namespace.table_exists(exists_req).await.is_ok());
11161
11162 let mut describe_req = DescribeTableRequest::new();
11163 describe_req.id = Some(vec!["test_table".to_string()]);
11164 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11165 assert_eq!(describe_response.is_only_declared, None);
11166
11167 let mut describe_req = DescribeTableRequest::new();
11168 describe_req.id = Some(vec!["test_table".to_string()]);
11169 describe_req.check_declared = Some(true);
11170 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11171 assert_eq!(describe_response.is_only_declared, Some(true));
11172 assert_eq!(
11173 describe_response
11174 .properties
11175 .as_ref()
11176 .and_then(|properties| properties.get("owner")),
11177 Some(&"alice".to_string())
11178 );
11179
11180 let mut list_req = ListTablesRequest::new();
11181 list_req.id = Some(vec![]);
11182 assert_eq!(
11183 namespace
11184 .list_tables(list_req.clone())
11185 .await
11186 .unwrap()
11187 .tables,
11188 vec!["test_table".to_string()]
11189 );
11190 list_req.include_declared = Some(false);
11191 assert!(
11192 namespace
11193 .list_tables(list_req)
11194 .await
11195 .unwrap()
11196 .tables
11197 .is_empty()
11198 );
11199 }
11200
11201 #[tokio::test]
11202 async fn test_declare_table_with_manifest_marker_already_exists() {
11203 use lance_namespace::error::ErrorCode;
11206 use lance_namespace::models::DeclareTableRequest;
11207
11208 let temp_dir = TempStdDir::default();
11209 let temp_path = temp_dir.to_str().unwrap();
11210
11211 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11212 .manifest_enabled(true)
11213 .dir_listing_enabled(true)
11214 .build()
11215 .await
11216 .unwrap();
11217
11218 let table_dir = temp_dir.join("test_table.lance");
11219 std::fs::create_dir_all(&table_dir).unwrap();
11220 std::fs::write(table_dir.join(".lance-reserved"), b"reserved").unwrap();
11221
11222 let mut declare_req = DeclareTableRequest::new();
11223 declare_req.id = Some(vec!["test_table".to_string()]);
11224 let err = namespace
11225 .declare_table(declare_req)
11226 .await
11227 .expect_err("declare with existing marker must fail");
11228 let msg = err.to_string();
11229 assert!(
11230 msg.contains("already exists") || msg.contains("TableAlreadyExists"),
11231 "expected TableAlreadyExists, got: {msg}"
11232 );
11233 assert_eq!(
11234 mutation_error_code(err),
11235 ErrorCode::TableAlreadyExists,
11236 "expected TableAlreadyExists error code"
11237 );
11238 }
11239
11240 #[tokio::test]
11241 async fn test_declare_table_when_table_exists() {
11242 use lance_namespace::models::DeclareTableRequest;
11243
11244 let temp_dir = TempStdDir::default();
11245 let temp_path = temp_dir.to_str().unwrap();
11246
11247 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11248 .manifest_enabled(false)
11249 .build()
11250 .await
11251 .unwrap();
11252
11253 let schema = create_test_schema();
11255 let ipc_data = create_test_ipc_data(&schema);
11256 let mut create_req = CreateTableRequest::new();
11257 create_req.id = Some(vec!["test_table".to_string()]);
11258 namespace
11259 .create_table(create_req, bytes::Bytes::from(ipc_data))
11260 .await
11261 .unwrap();
11262
11263 let mut declare_req = DeclareTableRequest::new();
11265 declare_req.id = Some(vec!["test_table".to_string()]);
11266 let result = namespace.declare_table(declare_req).await;
11267 assert!(result.is_err());
11268 }
11269
11270 #[tokio::test]
11275 async fn test_deregister_table_v1_mode() {
11276 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
11277
11278 let temp_dir = TempStdDir::default();
11279 let temp_path = temp_dir.to_str().unwrap();
11280
11281 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11283 .manifest_enabled(false)
11284 .dir_listing_enabled(true)
11285 .build()
11286 .await
11287 .unwrap();
11288
11289 let schema = create_test_schema();
11291 let ipc_data = create_test_ipc_data(&schema);
11292 let mut create_req = CreateTableRequest::new();
11293 create_req.id = Some(vec!["test_table".to_string()]);
11294 namespace
11295 .create_table(create_req, bytes::Bytes::from(ipc_data))
11296 .await
11297 .unwrap();
11298
11299 let mut exists_req = TableExistsRequest::new();
11301 exists_req.id = Some(vec!["test_table".to_string()]);
11302 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
11303
11304 let mut deregister_req = DeregisterTableRequest::new();
11306 deregister_req.id = Some(vec!["test_table".to_string()]);
11307 let response = namespace.deregister_table(deregister_req).await.unwrap();
11308
11309 assert!(response.location.is_some());
11311 let location = response.location.as_ref().unwrap();
11312 assert!(location.contains("test_table"));
11313
11314 let result = namespace.table_exists(exists_req).await;
11316 assert!(result.is_err());
11317 assert!(result.unwrap_err().to_string().contains("deregistered"));
11318
11319 let dataset = Dataset::open(location).await;
11321 assert!(dataset.is_ok(), "Physical table data should still exist");
11322 }
11323
11324 #[tokio::test]
11325 async fn test_deregister_table_v1_already_deregistered() {
11326 use lance_namespace::models::DeregisterTableRequest;
11327
11328 let temp_dir = TempStdDir::default();
11329 let temp_path = temp_dir.to_str().unwrap();
11330
11331 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11332 .manifest_enabled(false)
11333 .dir_listing_enabled(true)
11334 .build()
11335 .await
11336 .unwrap();
11337
11338 let schema = create_test_schema();
11340 let ipc_data = create_test_ipc_data(&schema);
11341 let mut create_req = CreateTableRequest::new();
11342 create_req.id = Some(vec!["test_table".to_string()]);
11343 namespace
11344 .create_table(create_req, bytes::Bytes::from(ipc_data))
11345 .await
11346 .unwrap();
11347
11348 let mut deregister_req = DeregisterTableRequest::new();
11350 deregister_req.id = Some(vec!["test_table".to_string()]);
11351 namespace
11352 .deregister_table(deregister_req.clone())
11353 .await
11354 .unwrap();
11355
11356 let result = namespace.deregister_table(deregister_req).await;
11358 assert!(result.is_err());
11359 assert!(
11360 result
11361 .unwrap_err()
11362 .to_string()
11363 .contains("already deregistered")
11364 );
11365 }
11366
11367 #[tokio::test]
11372 async fn test_list_tables_skips_deregistered_v1() {
11373 use lance_namespace::models::DeregisterTableRequest;
11374
11375 let temp_dir = TempStdDir::default();
11376 let temp_path = temp_dir.to_str().unwrap();
11377
11378 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11379 .manifest_enabled(false)
11380 .dir_listing_enabled(true)
11381 .build()
11382 .await
11383 .unwrap();
11384
11385 let schema = create_test_schema();
11387 let ipc_data = create_test_ipc_data(&schema);
11388
11389 let mut create_req1 = CreateTableRequest::new();
11390 create_req1.id = Some(vec!["table1".to_string()]);
11391 namespace
11392 .create_table(create_req1, bytes::Bytes::from(ipc_data.clone()))
11393 .await
11394 .unwrap();
11395
11396 let mut create_req2 = CreateTableRequest::new();
11397 create_req2.id = Some(vec!["table2".to_string()]);
11398 namespace
11399 .create_table(create_req2, bytes::Bytes::from(ipc_data))
11400 .await
11401 .unwrap();
11402
11403 let mut list_req = ListTablesRequest::new();
11405 list_req.id = Some(vec![]);
11406 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
11407 assert_eq!(list_response.tables.len(), 2);
11408
11409 let mut deregister_req = DeregisterTableRequest::new();
11411 deregister_req.id = Some(vec!["table1".to_string()]);
11412 namespace.deregister_table(deregister_req).await.unwrap();
11413
11414 let list_response = namespace.list_tables(list_req).await.unwrap();
11416 assert_eq!(list_response.tables.len(), 1);
11417 assert!(list_response.tables.contains(&"table2".to_string()));
11418 assert!(!list_response.tables.contains(&"table1".to_string()));
11419 }
11420
11421 #[tokio::test]
11426 async fn test_describe_table_fails_for_deregistered_v1() {
11427 use lance_namespace::models::{DeregisterTableRequest, DescribeTableRequest};
11428
11429 let temp_dir = TempStdDir::default();
11430 let temp_path = temp_dir.to_str().unwrap();
11431
11432 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11433 .manifest_enabled(false)
11434 .dir_listing_enabled(true)
11435 .build()
11436 .await
11437 .unwrap();
11438
11439 let schema = create_test_schema();
11441 let ipc_data = create_test_ipc_data(&schema);
11442 let mut create_req = CreateTableRequest::new();
11443 create_req.id = Some(vec!["test_table".to_string()]);
11444 namespace
11445 .create_table(create_req, bytes::Bytes::from(ipc_data))
11446 .await
11447 .unwrap();
11448
11449 let mut describe_req = DescribeTableRequest::new();
11451 describe_req.id = Some(vec!["test_table".to_string()]);
11452 assert!(namespace.describe_table(describe_req.clone()).await.is_ok());
11453
11454 let mut deregister_req = DeregisterTableRequest::new();
11456 deregister_req.id = Some(vec!["test_table".to_string()]);
11457 namespace.deregister_table(deregister_req).await.unwrap();
11458
11459 let result = namespace.describe_table(describe_req).await;
11461 assert!(result.is_err());
11462 let err = result.unwrap_err();
11463 assert!(matches!(err, Error::Namespace { .. }));
11464 let err_msg = err.to_string();
11465 assert!(err_msg.contains("deregistered"));
11466 assert!(err_msg.contains("table id 'test_table'"));
11467 }
11468
11469 #[tokio::test]
11470 async fn test_table_exists_fails_for_deregistered_v1() {
11471 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
11472
11473 let temp_dir = TempStdDir::default();
11474 let temp_path = temp_dir.to_str().unwrap();
11475
11476 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11477 .manifest_enabled(false)
11478 .dir_listing_enabled(true)
11479 .build()
11480 .await
11481 .unwrap();
11482
11483 let schema = create_test_schema();
11485 let ipc_data = create_test_ipc_data(&schema);
11486 let mut create_req = CreateTableRequest::new();
11487 create_req.id = Some(vec!["test_table".to_string()]);
11488 namespace
11489 .create_table(create_req, bytes::Bytes::from(ipc_data))
11490 .await
11491 .unwrap();
11492
11493 let mut exists_req = TableExistsRequest::new();
11495 exists_req.id = Some(vec!["test_table".to_string()]);
11496 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
11497
11498 let mut deregister_req = DeregisterTableRequest::new();
11500 deregister_req.id = Some(vec!["test_table".to_string()]);
11501 namespace.deregister_table(deregister_req).await.unwrap();
11502
11503 let result = namespace.table_exists(exists_req).await;
11505 assert!(result.is_err());
11506 let err = result.unwrap_err();
11507 assert!(matches!(err, Error::Namespace { .. }));
11508 let err_msg = err.to_string();
11509 assert!(err_msg.contains("deregistered"));
11510 assert!(err_msg.contains("table id 'test_table'"));
11511 }
11512
11513 #[tokio::test]
11514 async fn test_atomic_table_status_check() {
11515 let temp_dir = TempStdDir::default();
11519 let temp_path = temp_dir.to_str().unwrap();
11520
11521 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11522 .manifest_enabled(false)
11523 .dir_listing_enabled(true)
11524 .build()
11525 .await
11526 .unwrap();
11527
11528 let schema = create_test_schema();
11530 let ipc_data = create_test_ipc_data(&schema);
11531 let mut create_req = CreateTableRequest::new();
11532 create_req.id = Some(vec!["test_table".to_string()]);
11533 namespace
11534 .create_table(create_req, bytes::Bytes::from(ipc_data))
11535 .await
11536 .unwrap();
11537
11538 let status = namespace.check_table_status("test_table").await.unwrap();
11540 assert!(status.exists);
11541 assert!(!status.is_deregistered);
11542 assert!(!status.has_reserved_file);
11543 }
11544
11545 #[tokio::test]
11546 async fn test_table_version_tracking_enabled_managed_versioning() {
11547 use lance_namespace::models::DescribeTableRequest;
11548
11549 let temp_dir = TempStdDir::default();
11550 let temp_path = temp_dir.to_str().unwrap();
11551
11552 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11554 .table_version_tracking_enabled(true)
11555 .build()
11556 .await
11557 .unwrap();
11558
11559 let schema = create_test_schema();
11561 let ipc_data = create_test_ipc_data(&schema);
11562 let mut create_req = CreateTableRequest::new();
11563 create_req.id = Some(vec!["test_table".to_string()]);
11564 namespace
11565 .create_table(create_req, bytes::Bytes::from(ipc_data))
11566 .await
11567 .unwrap();
11568
11569 let mut describe_req = DescribeTableRequest::new();
11571 describe_req.id = Some(vec!["test_table".to_string()]);
11572 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
11573
11574 assert_eq!(
11576 describe_resp.managed_versioning,
11577 Some(true),
11578 "managed_versioning should be true when table_version_tracking_enabled=true"
11579 );
11580 }
11581
11582 #[tokio::test]
11583 async fn test_table_version_tracking_disabled_no_managed_versioning() {
11584 use lance_namespace::models::DescribeTableRequest;
11585
11586 let temp_dir = TempStdDir::default();
11587 let temp_path = temp_dir.to_str().unwrap();
11588
11589 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11591 .table_version_tracking_enabled(false)
11592 .build()
11593 .await
11594 .unwrap();
11595
11596 let schema = create_test_schema();
11598 let ipc_data = create_test_ipc_data(&schema);
11599 let mut create_req = CreateTableRequest::new();
11600 create_req.id = Some(vec!["test_table".to_string()]);
11601 namespace
11602 .create_table(create_req, bytes::Bytes::from(ipc_data))
11603 .await
11604 .unwrap();
11605
11606 let mut describe_req = DescribeTableRequest::new();
11608 describe_req.id = Some(vec!["test_table".to_string()]);
11609 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
11610
11611 assert!(
11613 describe_resp.managed_versioning.is_none(),
11614 "managed_versioning should be None when table_version_tracking_enabled=false, got: {:?}",
11615 describe_resp.managed_versioning
11616 );
11617 }
11618
11619 #[tokio::test]
11620 async fn test_list_table_versions() {
11621 use arrow::array::{Int32Array, RecordBatchIterator};
11622 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11623 use arrow::record_batch::RecordBatch;
11624 use lance::dataset::{Dataset, WriteMode, WriteParams};
11625 use lance_namespace::models::{CreateNamespaceRequest, ListTableVersionsRequest};
11626
11627 let temp_dir = TempStrDir::default();
11628 let temp_path: &str = &temp_dir;
11629
11630 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11631 DirectoryNamespaceBuilder::new(temp_path)
11632 .table_version_tracking_enabled(true)
11633 .build()
11634 .await
11635 .unwrap(),
11636 );
11637
11638 let mut create_ns_req = CreateNamespaceRequest::new();
11640 create_ns_req.id = Some(vec!["workspace".to_string()]);
11641 namespace.create_namespace(create_ns_req).await.unwrap();
11642
11643 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11645 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
11646 "id",
11647 DataType::Int32,
11648 false,
11649 )]));
11650 let batch = RecordBatch::try_new(
11651 arrow_schema.clone(),
11652 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
11653 )
11654 .unwrap();
11655 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
11656 let write_params = WriteParams {
11657 mode: WriteMode::Create,
11658 ..Default::default()
11659 };
11660 let mut dataset = Dataset::write_into_namespace(
11661 batches,
11662 namespace.clone(),
11663 table_id.clone(),
11664 Some(write_params),
11665 )
11666 .await
11667 .unwrap();
11668
11669 let batch2 = RecordBatch::try_new(
11671 arrow_schema.clone(),
11672 vec![Arc::new(Int32Array::from(vec![100, 200]))],
11673 )
11674 .unwrap();
11675 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
11676 dataset.append(batches, None).await.unwrap();
11677
11678 let batch3 = RecordBatch::try_new(
11680 arrow_schema.clone(),
11681 vec![Arc::new(Int32Array::from(vec![300, 400]))],
11682 )
11683 .unwrap();
11684 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
11685 dataset.append(batches, None).await.unwrap();
11686
11687 let mut list_req = ListTableVersionsRequest::new();
11689 list_req.id = Some(table_id.clone());
11690 let list_resp = namespace.list_table_versions(list_req).await.unwrap();
11691
11692 assert_eq!(
11693 list_resp.versions.len(),
11694 3,
11695 "Should have 3 versions, got: {:?}",
11696 list_resp.versions
11697 );
11698
11699 for expected_version in 1..=3 {
11701 let version = list_resp
11702 .versions
11703 .iter()
11704 .find(|v| v.version == expected_version)
11705 .unwrap_or_else(|| panic!("Expected version {}", expected_version));
11706
11707 assert!(
11708 !version.manifest_path.is_empty(),
11709 "manifest_path should be set for version {}",
11710 expected_version
11711 );
11712 assert!(
11713 version.manifest_path.contains(".manifest"),
11714 "manifest_path should contain .manifest for version {}",
11715 expected_version
11716 );
11717 assert!(
11718 version.manifest_size.is_some(),
11719 "manifest_size should be set for version {}",
11720 expected_version
11721 );
11722 assert!(
11723 version.manifest_size.unwrap() > 0,
11724 "manifest_size should be > 0 for version {}",
11725 expected_version
11726 );
11727 assert!(
11728 version.timestamp_millis.is_some(),
11729 "timestamp_millis should be set for version {}",
11730 expected_version
11731 );
11732 }
11733 }
11734
11735 #[tokio::test]
11736 async fn test_describe_table_version() {
11737 use arrow::array::{Int32Array, RecordBatchIterator};
11738 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11739 use arrow::record_batch::RecordBatch;
11740 use lance::dataset::{Dataset, WriteMode, WriteParams};
11741 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
11742
11743 let temp_dir = TempStrDir::default();
11744 let temp_path: &str = &temp_dir;
11745
11746 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11747 DirectoryNamespaceBuilder::new(temp_path)
11748 .table_version_tracking_enabled(true)
11749 .build()
11750 .await
11751 .unwrap(),
11752 );
11753
11754 let mut create_ns_req = CreateNamespaceRequest::new();
11756 create_ns_req.id = Some(vec!["workspace".to_string()]);
11757 namespace.create_namespace(create_ns_req).await.unwrap();
11758
11759 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11761 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
11762 "id",
11763 DataType::Int32,
11764 false,
11765 )]));
11766 let batch = RecordBatch::try_new(
11767 arrow_schema.clone(),
11768 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
11769 )
11770 .unwrap();
11771 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
11772 let write_params = WriteParams {
11773 mode: WriteMode::Create,
11774 ..Default::default()
11775 };
11776 let mut dataset = Dataset::write_into_namespace(
11777 batches,
11778 namespace.clone(),
11779 table_id.clone(),
11780 Some(write_params),
11781 )
11782 .await
11783 .unwrap();
11784
11785 let batch2 = RecordBatch::try_new(
11787 arrow_schema.clone(),
11788 vec![Arc::new(Int32Array::from(vec![100, 200]))],
11789 )
11790 .unwrap();
11791 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
11792 dataset.append(batches, None).await.unwrap();
11793
11794 let mut describe_req = DescribeTableVersionRequest::new();
11796 describe_req.id = Some(table_id.clone());
11797 describe_req.version = Some(1);
11798 let describe_resp = namespace
11799 .describe_table_version(describe_req)
11800 .await
11801 .unwrap();
11802
11803 let version = &describe_resp.version;
11804 assert_eq!(version.version, 1);
11805 assert!(version.timestamp_millis.is_some());
11806 assert!(
11807 !version.manifest_path.is_empty(),
11808 "manifest_path should be set"
11809 );
11810 assert!(
11811 version.manifest_path.contains(".manifest"),
11812 "manifest_path should contain .manifest"
11813 );
11814 assert!(
11815 version.manifest_size.is_some(),
11816 "manifest_size should be set"
11817 );
11818 assert!(
11819 version.manifest_size.unwrap() > 0,
11820 "manifest_size should be > 0"
11821 );
11822
11823 let mut describe_req = DescribeTableVersionRequest::new();
11825 describe_req.id = Some(table_id.clone());
11826 describe_req.version = Some(2);
11827 let describe_resp = namespace
11828 .describe_table_version(describe_req)
11829 .await
11830 .unwrap();
11831
11832 let version = &describe_resp.version;
11833 assert_eq!(version.version, 2);
11834 assert!(version.timestamp_millis.is_some());
11835 assert!(
11836 !version.manifest_path.is_empty(),
11837 "manifest_path should be set"
11838 );
11839 assert!(
11840 version.manifest_size.is_some(),
11841 "manifest_size should be set"
11842 );
11843 assert!(
11844 version.manifest_size.unwrap() > 0,
11845 "manifest_size should be > 0"
11846 );
11847 }
11848
11849 #[tokio::test]
11850 async fn test_describe_table_version_latest() {
11851 use arrow::array::{Int32Array, RecordBatchIterator};
11852 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11853 use arrow::record_batch::RecordBatch;
11854 use lance::dataset::{Dataset, WriteMode, WriteParams};
11855 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
11856
11857 let temp_dir = TempStrDir::default();
11858 let temp_path: &str = &temp_dir;
11859
11860 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11861 DirectoryNamespaceBuilder::new(temp_path)
11862 .table_version_tracking_enabled(true)
11863 .build()
11864 .await
11865 .unwrap(),
11866 );
11867
11868 let mut create_ns_req = CreateNamespaceRequest::new();
11870 create_ns_req.id = Some(vec!["workspace".to_string()]);
11871 namespace.create_namespace(create_ns_req).await.unwrap();
11872
11873 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11875 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
11876 "id",
11877 DataType::Int32,
11878 false,
11879 )]));
11880 let batch = RecordBatch::try_new(
11881 arrow_schema.clone(),
11882 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
11883 )
11884 .unwrap();
11885 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
11886 let write_params = WriteParams {
11887 mode: WriteMode::Create,
11888 ..Default::default()
11889 };
11890 let mut dataset = Dataset::write_into_namespace(
11891 batches,
11892 namespace.clone(),
11893 table_id.clone(),
11894 Some(write_params),
11895 )
11896 .await
11897 .unwrap();
11898
11899 let batch2 = RecordBatch::try_new(
11901 arrow_schema.clone(),
11902 vec![Arc::new(Int32Array::from(vec![100, 200]))],
11903 )
11904 .unwrap();
11905 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
11906 dataset.append(batches, None).await.unwrap();
11907
11908 let batch3 = RecordBatch::try_new(
11910 arrow_schema.clone(),
11911 vec![Arc::new(Int32Array::from(vec![300, 400]))],
11912 )
11913 .unwrap();
11914 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
11915 dataset.append(batches, None).await.unwrap();
11916
11917 let mut describe_req = DescribeTableVersionRequest::new();
11919 describe_req.id = Some(table_id.clone());
11920 describe_req.version = None;
11921 let describe_resp = namespace
11922 .describe_table_version(describe_req)
11923 .await
11924 .unwrap();
11925
11926 assert_eq!(describe_resp.version.version, 3);
11928 }
11929
11930 #[tokio::test]
11931 async fn test_create_table_version() {
11932 use futures::TryStreamExt;
11933 use lance::dataset::builder::DatasetBuilder;
11934 use lance_namespace::models::CreateTableVersionRequest;
11935
11936 let temp_dir = TempStrDir::default();
11937 let temp_path: &str = &temp_dir;
11938
11939 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11940 DirectoryNamespaceBuilder::new(temp_path)
11941 .table_version_tracking_enabled(true)
11942 .build()
11943 .await
11944 .unwrap(),
11945 );
11946
11947 let schema = create_test_schema();
11949 let ipc_data = create_test_ipc_data(&schema);
11950 let mut create_req = CreateTableRequest::new();
11951 create_req.id = Some(vec!["test_table".to_string()]);
11952 namespace
11953 .create_table(create_req, bytes::Bytes::from(ipc_data))
11954 .await
11955 .unwrap();
11956
11957 let table_id = vec!["test_table".to_string()];
11959 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
11960 .await
11961 .unwrap()
11962 .load()
11963 .await
11964 .unwrap();
11965
11966 let versions_path = dataset.versions_dir();
11968 let manifest_metas: Vec<_> = dataset
11969 .object_store(None)
11970 .await
11971 .unwrap()
11972 .inner
11973 .list(Some(&versions_path))
11974 .try_collect()
11975 .await
11976 .unwrap();
11977
11978 let manifest_meta = manifest_metas
11979 .iter()
11980 .find(|m| {
11981 m.location
11982 .filename()
11983 .map(|f| f.ends_with(".manifest"))
11984 .unwrap_or(false)
11985 })
11986 .expect("No manifest file found");
11987
11988 let manifest_data = dataset
11990 .object_store(None)
11991 .await
11992 .unwrap()
11993 .inner
11994 .get(&manifest_meta.location)
11995 .await
11996 .unwrap()
11997 .bytes()
11998 .await
11999 .unwrap();
12000
12001 let staging_path = dataset.versions_dir().join("staging_manifest");
12003 dataset
12004 .object_store(None)
12005 .await
12006 .unwrap()
12007 .inner
12008 .put(&staging_path, manifest_data.into())
12009 .await
12010 .unwrap();
12011
12012 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12015 create_version_req.id = Some(table_id.clone());
12016 create_version_req.naming_scheme = Some("V2".to_string());
12017
12018 let result = namespace.create_table_version(create_version_req).await;
12019 assert!(
12020 result.is_ok(),
12021 "create_table_version should succeed: {:?}",
12022 result
12023 );
12024
12025 let response = result.unwrap();
12027 let version_info = response
12028 .version
12029 .expect("response should contain version info");
12030 let version_2_path = Path::parse(&version_info.manifest_path).unwrap();
12031 let head_result = dataset
12032 .object_store(None)
12033 .await
12034 .unwrap()
12035 .inner
12036 .head(&version_2_path)
12037 .await;
12038 assert!(
12039 head_result.is_ok(),
12040 "Version 2 manifest should exist at {}",
12041 version_2_path
12042 );
12043
12044 let staging_head_result = dataset
12046 .object_store(None)
12047 .await
12048 .unwrap()
12049 .inner
12050 .head(&staging_path)
12051 .await;
12052 assert!(
12053 staging_head_result.is_err(),
12054 "Staging manifest should have been deleted after create_table_version"
12055 );
12056 }
12057
12058 #[tokio::test]
12059 async fn test_create_table_version_idempotent() {
12060 use futures::TryStreamExt;
12063 use lance::dataset::builder::DatasetBuilder;
12064 use lance_namespace::models::CreateTableVersionRequest;
12065
12066 let temp_dir = TempStrDir::default();
12067 let temp_path: &str = &temp_dir;
12068
12069 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12070 DirectoryNamespaceBuilder::new(temp_path)
12071 .table_version_tracking_enabled(true)
12072 .build()
12073 .await
12074 .unwrap(),
12075 );
12076
12077 let schema = create_test_schema();
12078 let ipc_data = create_test_ipc_data(&schema);
12079 let mut create_req = CreateTableRequest::new();
12080 create_req.id = Some(vec!["test_table".to_string()]);
12081 namespace
12082 .create_table(create_req, bytes::Bytes::from(ipc_data))
12083 .await
12084 .unwrap();
12085
12086 let table_id = vec!["test_table".to_string()];
12087 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12088 .await
12089 .unwrap()
12090 .load()
12091 .await
12092 .unwrap();
12093
12094 let versions_path = dataset.versions_dir();
12095 let manifest_metas: Vec<_> = dataset
12096 .object_store(None)
12097 .await
12098 .unwrap()
12099 .inner
12100 .list(Some(&versions_path))
12101 .try_collect()
12102 .await
12103 .unwrap();
12104
12105 let manifest_meta = manifest_metas
12106 .iter()
12107 .find(|m| {
12108 m.location
12109 .filename()
12110 .map(|f| f.ends_with(".manifest"))
12111 .unwrap_or(false)
12112 })
12113 .expect("No manifest file found");
12114
12115 let manifest_data = dataset
12116 .object_store(None)
12117 .await
12118 .unwrap()
12119 .inner
12120 .get(&manifest_meta.location)
12121 .await
12122 .unwrap()
12123 .bytes()
12124 .await
12125 .unwrap();
12126
12127 let staging_path = dataset.versions_dir().join("staging_manifest");
12128 dataset
12129 .object_store(None)
12130 .await
12131 .unwrap()
12132 .inner
12133 .put(&staging_path, manifest_data.clone().into())
12134 .await
12135 .unwrap();
12136
12137 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12138 create_version_req.id = Some(table_id.clone());
12139 create_version_req.naming_scheme = Some("V2".to_string());
12140 let first = namespace
12141 .create_table_version(create_version_req)
12142 .await
12143 .expect("first create_table_version should succeed");
12144
12145 let retry_staging = dataset.versions_dir().join("staging_manifest_retry");
12147 dataset
12148 .object_store(None)
12149 .await
12150 .unwrap()
12151 .inner
12152 .put(&retry_staging, manifest_data.into())
12153 .await
12154 .unwrap();
12155
12156 let mut retry_req = CreateTableVersionRequest::new(2, retry_staging.to_string());
12157 retry_req.id = Some(table_id.clone());
12158 retry_req.naming_scheme = Some("V2".to_string());
12159 let second = namespace
12160 .create_table_version(retry_req)
12161 .await
12162 .expect("idempotent retry must succeed");
12163
12164 assert_eq!(
12165 first.version.as_ref().map(|v| v.version),
12166 second.version.as_ref().map(|v| v.version)
12167 );
12168 assert_eq!(
12169 first.version.as_ref().map(|v| &v.manifest_path),
12170 second.version.as_ref().map(|v| &v.manifest_path)
12171 );
12172 }
12173
12174 #[tokio::test]
12175 async fn test_create_table_version_conflict() {
12176 use futures::TryStreamExt;
12178 use lance::dataset::builder::DatasetBuilder;
12179 use lance_namespace::models::CreateTableVersionRequest;
12180
12181 let temp_dir = TempStrDir::default();
12182 let temp_path: &str = &temp_dir;
12183
12184 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12185 DirectoryNamespaceBuilder::new(temp_path)
12186 .table_version_tracking_enabled(true)
12187 .build()
12188 .await
12189 .unwrap(),
12190 );
12191
12192 let schema = create_test_schema();
12194 let ipc_data = create_test_ipc_data(&schema);
12195 let mut create_req = CreateTableRequest::new();
12196 create_req.id = Some(vec!["test_table".to_string()]);
12197 namespace
12198 .create_table(create_req, bytes::Bytes::from(ipc_data))
12199 .await
12200 .unwrap();
12201
12202 let table_id = vec!["test_table".to_string()];
12204 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12205 .await
12206 .unwrap()
12207 .load()
12208 .await
12209 .unwrap();
12210
12211 let versions_path = dataset.versions_dir();
12213 let manifest_metas: Vec<_> = dataset
12214 .object_store(None)
12215 .await
12216 .unwrap()
12217 .inner
12218 .list(Some(&versions_path))
12219 .try_collect()
12220 .await
12221 .unwrap();
12222
12223 let manifest_meta = manifest_metas
12224 .iter()
12225 .find(|m| {
12226 m.location
12227 .filename()
12228 .map(|f| f.ends_with(".manifest"))
12229 .unwrap_or(false)
12230 })
12231 .expect("No manifest file found");
12232
12233 let manifest_data = dataset
12235 .object_store(None)
12236 .await
12237 .unwrap()
12238 .inner
12239 .get(&manifest_meta.location)
12240 .await
12241 .unwrap()
12242 .bytes()
12243 .await
12244 .unwrap();
12245
12246 let staging_path = dataset.versions_dir().join("staging_manifest");
12248 dataset
12249 .object_store(None)
12250 .await
12251 .unwrap()
12252 .inner
12253 .put(&staging_path, manifest_data.into())
12254 .await
12255 .unwrap();
12256
12257 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12259 create_version_req.id = Some(table_id.clone());
12260 create_version_req.naming_scheme = Some("V2".to_string());
12261 let first_result = namespace.create_table_version(create_version_req).await;
12262 assert!(
12263 first_result.is_ok(),
12264 "First create_table_version for version 2 should succeed: {:?}",
12265 first_result
12266 );
12267
12268 let version_2_path = Path::parse(
12270 &first_result
12271 .unwrap()
12272 .version
12273 .expect("response should contain version info")
12274 .manifest_path,
12275 )
12276 .unwrap();
12277
12278 let conflict_staging = dataset.versions_dir().join("staging_manifest_conflict");
12280 dataset
12281 .object_store(None)
12282 .await
12283 .unwrap()
12284 .inner
12285 .put(
12286 &conflict_staging,
12287 bytes::Bytes::from_static(b"not-a-real-manifest").into(),
12288 )
12289 .await
12290 .unwrap();
12291
12292 let mut create_version_req =
12293 CreateTableVersionRequest::new(2, conflict_staging.to_string());
12294 create_version_req.id = Some(table_id.clone());
12295 create_version_req.naming_scheme = Some("V2".to_string());
12296
12297 let result = namespace.create_table_version(create_version_req).await;
12298 assert!(
12299 result.is_err(),
12300 "create_table_version should fail for existing version with different content"
12301 );
12302 let err = result.unwrap_err().to_string();
12303 assert!(
12304 err.contains("already exists") || err.contains("ConcurrentModification"),
12305 "expected ConcurrentModification, got: {err}"
12306 );
12307
12308 let head_result = dataset
12310 .object_store(None)
12311 .await
12312 .unwrap()
12313 .inner
12314 .head(&version_2_path)
12315 .await;
12316 assert!(
12317 head_result.is_ok(),
12318 "Version 2 manifest should still exist at {}",
12319 version_2_path
12320 );
12321 }
12322
12323 #[tokio::test]
12324 async fn test_create_table_version_cas_rejects_gap() {
12325 use futures::TryStreamExt;
12327 use lance::dataset::builder::DatasetBuilder;
12328 use lance_namespace::models::CreateTableVersionRequest;
12329
12330 let temp_dir = TempStrDir::default();
12331 let temp_path: &str = &temp_dir;
12332
12333 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12334 DirectoryNamespaceBuilder::new(temp_path)
12335 .table_version_tracking_enabled(true)
12336 .build()
12337 .await
12338 .unwrap(),
12339 );
12340
12341 let schema = create_test_schema();
12342 let ipc_data = create_test_ipc_data(&schema);
12343 let mut create_req = CreateTableRequest::new();
12344 create_req.id = Some(vec!["test_table".to_string()]);
12345 namespace
12346 .create_table(create_req, bytes::Bytes::from(ipc_data))
12347 .await
12348 .unwrap();
12349
12350 let table_id = vec!["test_table".to_string()];
12351 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12352 .await
12353 .unwrap()
12354 .load()
12355 .await
12356 .unwrap();
12357
12358 let versions_path = dataset.versions_dir();
12359 let manifest_metas: Vec<_> = dataset
12360 .object_store(None)
12361 .await
12362 .unwrap()
12363 .inner
12364 .list(Some(&versions_path))
12365 .try_collect()
12366 .await
12367 .unwrap();
12368 let manifest_meta = manifest_metas
12369 .iter()
12370 .find(|m| {
12371 m.location
12372 .filename()
12373 .map(|f| f.ends_with(".manifest"))
12374 .unwrap_or(false)
12375 })
12376 .expect("No manifest file found");
12377 let manifest_data = dataset
12378 .object_store(None)
12379 .await
12380 .unwrap()
12381 .inner
12382 .get(&manifest_meta.location)
12383 .await
12384 .unwrap()
12385 .bytes()
12386 .await
12387 .unwrap();
12388
12389 let staging_path = dataset.versions_dir().join("staging_gap");
12390 dataset
12391 .object_store(None)
12392 .await
12393 .unwrap()
12394 .inner
12395 .put(&staging_path, manifest_data.into())
12396 .await
12397 .unwrap();
12398
12399 let mut req = CreateTableVersionRequest::new(5, staging_path.to_string());
12401 req.id = Some(table_id);
12402 req.naming_scheme = Some("V2".to_string());
12403 let err = namespace
12404 .create_table_version(req)
12405 .await
12406 .expect_err("gap create must fail CAS");
12407 let msg = err.to_string();
12408 assert!(
12409 msg.contains("CAS") || msg.contains("ConcurrentModification"),
12410 "expected CAS ConcurrentModification, got: {msg}"
12411 );
12412 }
12413
12414 #[tokio::test]
12415 async fn test_create_table_version_branch_cas_requires_parent_version() {
12416 use futures::TryStreamExt;
12419 use lance_namespace::models::CreateTableVersionRequest;
12420
12421 let (namespace, _temp_dir) = create_test_namespace().await;
12422 create_scalar_table(&namespace, "users").await;
12423 let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
12424 append_scalar_version(&main_uri, 10).await; let mut main = open_dataset(&namespace, "users").await;
12427 let fork_version = main.version().version;
12428 assert_eq!(fork_version, 2);
12429 let branch_uri = main
12430 .create_branch("exp", fork_version, None)
12431 .await
12432 .unwrap()
12433 .uri()
12434 .to_string();
12435
12436 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
12437 let versions_dir = branch_ds.versions_dir();
12438 let store = branch_ds.object_store(None).await.unwrap();
12439 let manifests: Vec<_> = store
12440 .inner
12441 .list(Some(&versions_dir))
12442 .try_collect()
12443 .await
12444 .unwrap();
12445 for meta in &manifests {
12446 if meta
12447 .location
12448 .filename()
12449 .is_some_and(|f| f.ends_with(".manifest"))
12450 {
12451 store.inner.delete(&meta.location).await.unwrap();
12452 }
12453 }
12454 let remaining_manifests = store
12457 .inner
12458 .list(Some(&versions_dir))
12459 .try_collect::<Vec<_>>()
12460 .await
12461 .unwrap()
12462 .into_iter()
12463 .filter(|m| {
12464 m.location
12465 .filename()
12466 .is_some_and(|f| f.ends_with(".manifest"))
12467 })
12468 .count();
12469 assert_eq!(
12470 remaining_manifests, 0,
12471 "branch version chain should be empty after deleting manifests"
12472 );
12473
12474 let main_ds = open_dataset(&namespace, "users").await;
12476 let main_versions = main_ds.versions_dir();
12477 let main_store = main_ds.object_store(None).await.unwrap();
12478 let source_meta = main_store
12479 .inner
12480 .list(Some(&main_versions))
12481 .try_collect::<Vec<_>>()
12482 .await
12483 .unwrap()
12484 .into_iter()
12485 .find(|m| {
12486 m.location
12487 .filename()
12488 .is_some_and(|f| f.ends_with(".manifest"))
12489 })
12490 .expect("main should have a manifest");
12491 let source_bytes = main_store
12492 .inner
12493 .get(&source_meta.location)
12494 .await
12495 .unwrap()
12496 .bytes()
12497 .await
12498 .unwrap();
12499
12500 let staging_wrong = versions_dir.clone().join("staging_wrong");
12501 store
12502 .inner
12503 .put(&staging_wrong, source_bytes.clone().into())
12504 .await
12505 .unwrap();
12506 let err = namespace
12507 .create_table_version(CreateTableVersionRequest {
12508 id: Some(vec!["users".to_string()]),
12509 version: 1,
12510 manifest_path: staging_wrong.to_string(),
12511 naming_scheme: Some("V2".to_string()),
12512 branch: Some("exp".to_string()),
12513 ..Default::default()
12514 })
12515 .await
12516 .expect_err("bootstrap at v1 must fail when parent_version is 2");
12517 let msg = err.to_string();
12518 assert!(
12519 msg.contains("CAS") || msg.contains("ConcurrentModification"),
12520 "expected CAS ConcurrentModification, got: {msg}"
12521 );
12522
12523 let staging_ok = versions_dir.join("staging_ok");
12524 store
12525 .inner
12526 .put(&staging_ok, source_bytes.into())
12527 .await
12528 .unwrap();
12529 let resp = namespace
12530 .create_table_version(CreateTableVersionRequest {
12531 id: Some(vec!["users".to_string()]),
12532 version: 2,
12533 manifest_path: staging_ok.to_string(),
12534 naming_scheme: Some("V2".to_string()),
12535 branch: Some("exp".to_string()),
12536 ..Default::default()
12537 })
12538 .await
12539 .expect("bootstrap at parent_version must succeed");
12540 assert_eq!(resp.version.as_ref().map(|v| v.version), Some(2));
12541 }
12542
12543 #[tokio::test]
12544 async fn test_create_table_version_table_not_found() {
12545 use lance_namespace::models::CreateTableVersionRequest;
12546
12547 let temp_dir = TempStdDir::default();
12548 let temp_path = temp_dir.to_str().unwrap();
12549
12550 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12551 .table_version_tracking_enabled(true)
12552 .build()
12553 .await
12554 .unwrap();
12555
12556 let mut create_version_req =
12558 CreateTableVersionRequest::new(1, "/some/staging/path".to_string());
12559 create_version_req.id = Some(vec!["non_existent_table".to_string()]);
12560
12561 let result = namespace.create_table_version(create_version_req).await;
12562 assert!(
12563 result.is_err(),
12564 "create_table_version should fail for non-existent table"
12565 );
12566 let err_msg = result.unwrap_err().to_string();
12567 assert!(
12568 err_msg.contains("Table not found"),
12569 "Error should mention table not found, got: {}",
12570 err_msg
12571 );
12572 }
12573
12574 mod e2e_table_version_tracking {
12576 use super::*;
12577 use std::sync::atomic::{AtomicUsize, Ordering};
12578
12579 struct TrackingNamespace {
12581 inner: DirectoryNamespace,
12582 create_table_version_count: AtomicUsize,
12583 describe_table_version_count: AtomicUsize,
12584 list_table_versions_count: AtomicUsize,
12585 }
12586
12587 impl TrackingNamespace {
12588 fn new(inner: DirectoryNamespace) -> Self {
12589 Self {
12590 inner,
12591 create_table_version_count: AtomicUsize::new(0),
12592 describe_table_version_count: AtomicUsize::new(0),
12593 list_table_versions_count: AtomicUsize::new(0),
12594 }
12595 }
12596
12597 fn create_table_version_calls(&self) -> usize {
12598 self.create_table_version_count.load(Ordering::SeqCst)
12599 }
12600
12601 fn describe_table_version_calls(&self) -> usize {
12602 self.describe_table_version_count.load(Ordering::SeqCst)
12603 }
12604
12605 fn list_table_versions_calls(&self) -> usize {
12606 self.list_table_versions_count.load(Ordering::SeqCst)
12607 }
12608 }
12609
12610 impl std::fmt::Debug for TrackingNamespace {
12611 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12612 f.debug_struct("TrackingNamespace")
12613 .field(
12614 "create_table_version_calls",
12615 &self.create_table_version_calls(),
12616 )
12617 .finish()
12618 }
12619 }
12620
12621 #[async_trait]
12622 impl LanceNamespace for TrackingNamespace {
12623 async fn create_namespace(
12624 &self,
12625 request: CreateNamespaceRequest,
12626 ) -> Result<CreateNamespaceResponse> {
12627 self.inner.create_namespace(request).await
12628 }
12629
12630 async fn describe_namespace(
12631 &self,
12632 request: DescribeNamespaceRequest,
12633 ) -> Result<DescribeNamespaceResponse> {
12634 self.inner.describe_namespace(request).await
12635 }
12636
12637 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
12638 self.inner.namespace_exists(request).await
12639 }
12640
12641 async fn list_namespaces(
12642 &self,
12643 request: ListNamespacesRequest,
12644 ) -> Result<ListNamespacesResponse> {
12645 self.inner.list_namespaces(request).await
12646 }
12647
12648 async fn drop_namespace(
12649 &self,
12650 request: DropNamespaceRequest,
12651 ) -> Result<DropNamespaceResponse> {
12652 self.inner.drop_namespace(request).await
12653 }
12654
12655 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
12656 self.inner.list_tables(request).await
12657 }
12658
12659 async fn describe_table(
12660 &self,
12661 request: DescribeTableRequest,
12662 ) -> Result<DescribeTableResponse> {
12663 self.inner.describe_table(request).await
12664 }
12665
12666 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
12667 self.inner.table_exists(request).await
12668 }
12669
12670 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
12671 self.inner.drop_table(request).await
12672 }
12673
12674 async fn create_table(
12675 &self,
12676 request: CreateTableRequest,
12677 request_data: Bytes,
12678 ) -> Result<CreateTableResponse> {
12679 self.inner.create_table(request, request_data).await
12680 }
12681
12682 async fn declare_table(
12683 &self,
12684 request: DeclareTableRequest,
12685 ) -> Result<DeclareTableResponse> {
12686 self.inner.declare_table(request).await
12687 }
12688
12689 async fn list_table_versions(
12690 &self,
12691 request: ListTableVersionsRequest,
12692 ) -> Result<ListTableVersionsResponse> {
12693 self.list_table_versions_count
12694 .fetch_add(1, Ordering::SeqCst);
12695 self.inner.list_table_versions(request).await
12696 }
12697
12698 async fn create_table_version(
12699 &self,
12700 request: CreateTableVersionRequest,
12701 ) -> Result<CreateTableVersionResponse> {
12702 self.create_table_version_count
12703 .fetch_add(1, Ordering::SeqCst);
12704 self.inner.create_table_version(request).await
12705 }
12706
12707 async fn describe_table_version(
12708 &self,
12709 request: DescribeTableVersionRequest,
12710 ) -> Result<DescribeTableVersionResponse> {
12711 self.describe_table_version_count
12712 .fetch_add(1, Ordering::SeqCst);
12713 self.inner.describe_table_version(request).await
12714 }
12715
12716 async fn batch_delete_table_versions(
12717 &self,
12718 request: BatchDeleteTableVersionsRequest,
12719 ) -> Result<BatchDeleteTableVersionsResponse> {
12720 self.inner.batch_delete_table_versions(request).await
12721 }
12722
12723 fn namespace_id(&self) -> String {
12724 self.inner.namespace_id()
12725 }
12726 }
12727
12728 #[tokio::test]
12729 async fn test_describe_table_returns_managed_versioning() {
12730 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableRequest};
12731
12732 let temp_dir = TempStdDir::default();
12733 let temp_path = temp_dir.to_str().unwrap();
12734
12735 let ns = DirectoryNamespaceBuilder::new(temp_path)
12737 .table_version_tracking_enabled(true)
12738 .manifest_enabled(true)
12739 .build()
12740 .await
12741 .unwrap();
12742
12743 let mut create_ns_req = CreateNamespaceRequest::new();
12745 create_ns_req.id = Some(vec!["workspace".to_string()]);
12746 ns.create_namespace(create_ns_req).await.unwrap();
12747
12748 let schema = create_test_schema();
12750 let ipc_data = create_test_ipc_data(&schema);
12751 let mut create_req = CreateTableRequest::new();
12752 create_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
12753 ns.create_table(create_req, bytes::Bytes::from(ipc_data))
12754 .await
12755 .unwrap();
12756
12757 let mut describe_req = DescribeTableRequest::new();
12759 describe_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
12760 let describe_resp = ns.describe_table(describe_req).await.unwrap();
12761
12762 assert_eq!(
12764 describe_resp.managed_versioning,
12765 Some(true),
12766 "managed_versioning should be true when table_version_tracking_enabled=true"
12767 );
12768 }
12769
12770 #[tokio::test]
12771 async fn test_external_manifest_store_invokes_namespace_apis() {
12772 use arrow::array::{Int32Array, StringArray};
12773 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12774 use arrow::record_batch::RecordBatch;
12775 use lance::Dataset;
12776 use lance::dataset::builder::DatasetBuilder;
12777 use lance::dataset::{WriteMode, WriteParams};
12778 use lance_namespace::models::CreateNamespaceRequest;
12779
12780 let temp_dir = TempStdDir::default();
12781 let temp_path = temp_dir.to_str().unwrap();
12782
12783 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
12785 .table_version_tracking_enabled(true)
12786 .manifest_enabled(true)
12787 .build()
12788 .await
12789 .unwrap();
12790
12791 let tracking_ns = Arc::new(TrackingNamespace::new(inner_ns));
12792 let ns: Arc<dyn LanceNamespace> = tracking_ns.clone();
12793
12794 let mut create_ns_req = CreateNamespaceRequest::new();
12796 create_ns_req.id = Some(vec!["workspace".to_string()]);
12797 ns.create_namespace(create_ns_req).await.unwrap();
12798
12799 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12801
12802 let arrow_schema = Arc::new(ArrowSchema::new(vec![
12804 Field::new("id", DataType::Int32, false),
12805 Field::new("name", DataType::Utf8, true),
12806 ]));
12807 let batch = RecordBatch::try_new(
12808 arrow_schema.clone(),
12809 vec![
12810 Arc::new(Int32Array::from(vec![1, 2, 3])),
12811 Arc::new(StringArray::from(vec!["a", "b", "c"])),
12812 ],
12813 )
12814 .unwrap();
12815
12816 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
12818 let write_params = WriteParams {
12819 mode: WriteMode::Create,
12820 ..Default::default()
12821 };
12822 let mut dataset = Dataset::write_into_namespace(
12823 batches,
12824 ns.clone(),
12825 table_id.clone(),
12826 Some(write_params),
12827 )
12828 .await
12829 .unwrap();
12830 assert_eq!(dataset.version().version, 1);
12831
12832 assert_eq!(
12834 tracking_ns.create_table_version_calls(),
12835 1,
12836 "create_table_version should have been called once during initial write_into_namespace"
12837 );
12838
12839 let append_batch = RecordBatch::try_new(
12841 arrow_schema.clone(),
12842 vec![
12843 Arc::new(Int32Array::from(vec![4, 5, 6])),
12844 Arc::new(StringArray::from(vec!["d", "e", "f"])),
12845 ],
12846 )
12847 .unwrap();
12848 let append_batches = RecordBatchIterator::new(vec![Ok(append_batch)], arrow_schema);
12849 dataset.append(append_batches, None).await.unwrap();
12850
12851 assert_eq!(
12852 tracking_ns.create_table_version_calls(),
12853 2,
12854 "create_table_version should have been called twice (once for create, once for append)"
12855 );
12856
12857 let initial_list_calls = tracking_ns.list_table_versions_calls();
12859 let latest_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
12860 .await
12861 .unwrap()
12862 .load()
12863 .await
12864 .unwrap();
12865 assert_eq!(latest_dataset.version().version, 2);
12866 assert_eq!(
12867 tracking_ns.list_table_versions_calls(),
12868 initial_list_calls + 1,
12869 "list_table_versions should have been called exactly once during checkout_latest"
12870 );
12871
12872 let initial_describe_calls = tracking_ns.describe_table_version_calls();
12874 let v1_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
12875 .await
12876 .unwrap()
12877 .with_version(1)
12878 .load()
12879 .await
12880 .unwrap();
12881 assert_eq!(v1_dataset.version().version, 1);
12882 assert_eq!(
12883 tracking_ns.describe_table_version_calls(),
12884 initial_describe_calls + 1,
12885 "describe_table_version should have been called exactly once during checkout to version 1"
12886 );
12887 }
12888
12889 #[tokio::test]
12890 async fn test_dataset_commit_with_external_manifest_store() {
12891 use arrow::array::{Int32Array, StringArray};
12892 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12893 use arrow::record_batch::RecordBatch;
12894 use futures::TryStreamExt;
12895 use lance::dataset::{Dataset, WriteMode, WriteParams};
12896 use lance_namespace::models::CreateNamespaceRequest;
12897 use lance_table::io::commit::ManifestNamingScheme;
12898
12899 let temp_dir = TempStdDir::default();
12900 let temp_path = temp_dir.to_str().unwrap();
12901
12902 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
12904 .table_version_tracking_enabled(true)
12905 .manifest_enabled(true)
12906 .build()
12907 .await
12908 .unwrap();
12909
12910 let tracking_ns: Arc<dyn LanceNamespace> = Arc::new(TrackingNamespace::new(inner_ns));
12911
12912 let mut create_ns_req = CreateNamespaceRequest::new();
12914 create_ns_req.id = Some(vec!["workspace".to_string()]);
12915 tracking_ns.create_namespace(create_ns_req).await.unwrap();
12916
12917 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12919 let arrow_schema = Arc::new(ArrowSchema::new(vec![
12920 Field::new("id", DataType::Int32, false),
12921 Field::new("name", DataType::Utf8, true),
12922 ]));
12923 let batch = RecordBatch::try_new(
12924 arrow_schema.clone(),
12925 vec![
12926 Arc::new(Int32Array::from(vec![1, 2, 3])),
12927 Arc::new(StringArray::from(vec!["a", "b", "c"])),
12928 ],
12929 )
12930 .unwrap();
12931 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12932 let write_params = WriteParams {
12933 mode: WriteMode::Create,
12934 ..Default::default()
12935 };
12936 let dataset = Dataset::write_into_namespace(
12937 batches,
12938 tracking_ns.clone(),
12939 table_id.clone(),
12940 Some(write_params),
12941 )
12942 .await
12943 .unwrap();
12944 assert_eq!(dataset.version().version, 1);
12945
12946 let batch2 = RecordBatch::try_new(
12948 arrow_schema.clone(),
12949 vec![
12950 Arc::new(Int32Array::from(vec![4, 5, 6])),
12951 Arc::new(StringArray::from(vec!["d", "e", "f"])),
12952 ],
12953 )
12954 .unwrap();
12955 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
12956 let write_params = WriteParams {
12957 mode: WriteMode::Append,
12958 ..Default::default()
12959 };
12960 Dataset::write_into_namespace(
12961 batches,
12962 tracking_ns.clone(),
12963 table_id.clone(),
12964 Some(write_params),
12965 )
12966 .await
12967 .unwrap();
12968
12969 let manifest_metas: Vec<_> = dataset
12972 .object_store(None)
12973 .await
12974 .unwrap()
12975 .inner
12976 .list(Some(&dataset.versions_dir()))
12977 .try_collect()
12978 .await
12979 .unwrap();
12980 let version_2_found = manifest_metas.iter().any(|m| {
12981 m.location
12982 .filename()
12983 .map(|f| {
12984 f.ends_with(".manifest")
12985 && ManifestNamingScheme::V2.parse_version(f) == Some(2)
12986 })
12987 .unwrap_or(false)
12988 });
12989 assert!(
12990 version_2_found,
12991 "Version 2 manifest should exist in versions directory"
12992 );
12993 }
12994
12995 async fn create_ns_with_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
12997 use arrow::array::{Int32Array, StringArray};
12998 use arrow::ipc::writer::StreamWriter;
12999
13000 let (namespace, temp_dir) = create_test_namespace().await;
13001
13002 let schema = create_test_schema();
13003 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13004 let arrow_schema = Arc::new(arrow_schema);
13005
13006 let id_array = Int32Array::from(vec![1, 2, 3]);
13007 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
13008 let batch = arrow::record_batch::RecordBatch::try_new(
13009 arrow_schema.clone(),
13010 vec![Arc::new(id_array), Arc::new(name_array)],
13011 )
13012 .unwrap();
13013
13014 let mut buffer = Vec::new();
13015 {
13016 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13017 writer.write(&batch).unwrap();
13018 writer.finish().unwrap();
13019 }
13020
13021 let mut request = CreateTableRequest::new();
13022 let table_id = vec!["test_ops_table".to_string()];
13023 request.id = Some(table_id.clone());
13024
13025 namespace
13026 .create_table(request, Bytes::from(buffer))
13027 .await
13028 .unwrap();
13029
13030 (namespace, temp_dir, table_id)
13031 }
13032
13033 #[tokio::test]
13034 async fn test_count_table_rows_basic() {
13035 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13036
13037 let request = CountTableRowsRequest {
13038 id: Some(table_id),
13039 version: None,
13040 predicate: None,
13041 ..Default::default()
13042 };
13043
13044 let count = namespace.count_table_rows(request).await.unwrap();
13045 assert_eq!(count, 3);
13046 }
13047
13048 #[tokio::test]
13049 async fn test_count_table_rows_with_predicate() {
13050 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13051
13052 let request = CountTableRowsRequest {
13053 id: Some(table_id),
13054 version: None,
13055 predicate: Some("id > 1".to_string()),
13056 ..Default::default()
13057 };
13058
13059 let count = namespace.count_table_rows(request).await.unwrap();
13060 assert_eq!(count, 2);
13061 }
13062
13063 #[tokio::test]
13064 async fn test_query_table_invalid_distance_type() {
13065 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13066
13067 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13068 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13069 multi_vector: None,
13070 });
13071
13072 let request = QueryTableRequest {
13073 id: Some(table_id),
13074 k: 2,
13075 vector,
13076 vector_column: Some("vector".to_string()),
13077 distance_type: Some("invalid_metric".to_string()),
13078 filter: None,
13079 offset: None,
13080 version: None,
13081 ..Default::default()
13082 };
13083
13084 let result = namespace.query_table(request).await;
13085 assert!(result.is_err());
13086 let err_msg = result.unwrap_err().to_string();
13087 assert!(
13088 err_msg.contains("Unknown distance type"),
13089 "Expected error about unknown distance type, got: {}",
13090 err_msg
13091 );
13092 }
13093
13094 #[tokio::test]
13095 async fn test_insert_into_table_append() {
13096 use arrow::array::{Int32Array, StringArray};
13097 use arrow::ipc::writer::StreamWriter;
13098
13099 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13100
13101 let schema = create_test_schema();
13103 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13104 let arrow_schema = Arc::new(arrow_schema);
13105
13106 let id_array = Int32Array::from(vec![4, 5]);
13107 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13108 let batch = arrow::record_batch::RecordBatch::try_new(
13109 arrow_schema.clone(),
13110 vec![Arc::new(id_array), Arc::new(name_array)],
13111 )
13112 .unwrap();
13113
13114 let mut buffer = Vec::new();
13115 {
13116 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13117 writer.write(&batch).unwrap();
13118 writer.finish().unwrap();
13119 }
13120
13121 let request = InsertIntoTableRequest {
13122 id: Some(table_id.clone()),
13123 mode: Some("append".to_string()),
13124 ..Default::default()
13125 };
13126
13127 let response = namespace
13128 .insert_into_table(request, Bytes::from(buffer))
13129 .await
13130 .unwrap();
13131 assert!(response.transaction_id.is_none());
13132
13133 let count_req = CountTableRowsRequest {
13135 id: Some(table_id),
13136 version: None,
13137 predicate: None,
13138 ..Default::default()
13139 };
13140 let count = namespace.count_table_rows(count_req).await.unwrap();
13141 assert_eq!(count, 5);
13142 }
13143
13144 #[tokio::test]
13145 async fn test_insert_into_table_overwrite() {
13146 use arrow::array::{Int32Array, StringArray};
13147 use arrow::ipc::writer::StreamWriter;
13148
13149 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13150
13151 let schema = create_test_schema();
13152 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13153 let arrow_schema = Arc::new(arrow_schema);
13154
13155 let id_array = Int32Array::from(vec![10, 20]);
13156 let name_array = StringArray::from(vec!["X", "Y"]);
13157 let batch = arrow::record_batch::RecordBatch::try_new(
13158 arrow_schema.clone(),
13159 vec![Arc::new(id_array), Arc::new(name_array)],
13160 )
13161 .unwrap();
13162
13163 let mut buffer = Vec::new();
13164 {
13165 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13166 writer.write(&batch).unwrap();
13167 writer.finish().unwrap();
13168 }
13169
13170 let request = InsertIntoTableRequest {
13171 id: Some(table_id.clone()),
13172 mode: Some("overwrite".to_string()),
13173 ..Default::default()
13174 };
13175
13176 namespace
13177 .insert_into_table(request, Bytes::from(buffer))
13178 .await
13179 .unwrap();
13180
13181 let count_req = CountTableRowsRequest {
13183 id: Some(table_id),
13184 version: None,
13185 predicate: None,
13186 ..Default::default()
13187 };
13188 let count = namespace.count_table_rows(count_req).await.unwrap();
13189 assert_eq!(count, 2);
13190 }
13191
13192 #[tokio::test]
13193 async fn test_insert_into_table_empty_data() {
13194 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13195
13196 let request = InsertIntoTableRequest {
13197 id: Some(table_id),
13198 mode: None,
13199 ..Default::default()
13200 };
13201
13202 let result = namespace.insert_into_table(request, Bytes::new()).await;
13203 assert!(result.is_err());
13204 assert!(
13205 result
13206 .unwrap_err()
13207 .to_string()
13208 .contains("Arrow IPC stream) is required")
13209 );
13210 }
13211
13212 #[tokio::test]
13213 async fn test_insert_into_table_with_storage_options() {
13214 use arrow::array::{Int32Array, StringArray};
13215 use arrow::ipc::writer::StreamWriter;
13216
13217 let temp_dir = TempStdDir::default();
13218
13219 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
13221 .storage_option("allow_http", "true")
13222 .build()
13223 .await
13224 .unwrap();
13225
13226 let schema = create_test_schema();
13228 let ipc_data = create_test_ipc_data(&schema);
13229 let mut create_req = CreateTableRequest::new();
13230 let table_id = vec!["so_table".to_string()];
13231 create_req.id = Some(table_id.clone());
13232 namespace
13233 .create_table(create_req, Bytes::from(ipc_data))
13234 .await
13235 .unwrap();
13236
13237 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13239 let arrow_schema = Arc::new(arrow_schema);
13240
13241 let id_array = Int32Array::from(vec![10, 20]);
13242 let name_array = StringArray::from(vec!["X", "Y"]);
13243 let batch = arrow::record_batch::RecordBatch::try_new(
13244 arrow_schema.clone(),
13245 vec![Arc::new(id_array), Arc::new(name_array)],
13246 )
13247 .unwrap();
13248
13249 let mut buffer = Vec::new();
13250 {
13251 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13252 writer.write(&batch).unwrap();
13253 writer.finish().unwrap();
13254 }
13255
13256 let request = InsertIntoTableRequest {
13257 id: Some(table_id.clone()),
13258 mode: Some("append".to_string()),
13259 ..Default::default()
13260 };
13261
13262 let response = namespace
13263 .insert_into_table(request, Bytes::from(buffer))
13264 .await
13265 .unwrap();
13266 assert!(response.transaction_id.is_none());
13267
13268 let count_req = CountTableRowsRequest {
13270 id: Some(table_id),
13271 version: None,
13272 predicate: None,
13273 ..Default::default()
13274 };
13275 let count = namespace.count_table_rows(count_req).await.unwrap();
13276 assert_eq!(count, 2);
13277 }
13278
13279 #[tokio::test]
13280 async fn test_query_table_basic() {
13281 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13282
13283 let request = QueryTableRequest {
13284 id: Some(table_id),
13285 k: 10,
13286 filter: None,
13287 offset: None,
13288 version: None,
13289 ..Default::default()
13290 };
13291
13292 let bytes = namespace.query_table(request).await.unwrap();
13293
13294 let cursor = Cursor::new(bytes.to_vec());
13296 let reader = FileReader::try_new(cursor, None).unwrap();
13297 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13298 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13299 assert_eq!(total_rows, 3);
13300 }
13301
13302 #[tokio::test]
13303 async fn test_query_table_with_filter() {
13304 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13305
13306 let request = QueryTableRequest {
13307 id: Some(table_id),
13308 k: 10,
13309 filter: Some("id <= 2".to_string()),
13310 offset: None,
13311 version: None,
13312 ..Default::default()
13313 };
13314
13315 let bytes = namespace.query_table(request).await.unwrap();
13316
13317 let cursor = Cursor::new(bytes.to_vec());
13318 let reader = FileReader::try_new(cursor, None).unwrap();
13319 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13320 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13321 assert_eq!(total_rows, 2);
13322 }
13323
13324 #[tokio::test]
13325 async fn test_query_table_with_limit_and_offset() {
13326 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13327
13328 let request = QueryTableRequest {
13329 id: Some(table_id),
13330 k: 2,
13331 filter: None,
13332 offset: Some(1),
13333 version: None,
13334 ..Default::default()
13335 };
13336
13337 let bytes = namespace.query_table(request).await.unwrap();
13338
13339 let cursor = Cursor::new(bytes.to_vec());
13340 let reader = FileReader::try_new(cursor, None).unwrap();
13341 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13342 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13343 assert_eq!(total_rows, 2);
13344 }
13345
13346 #[tokio::test]
13347 async fn test_query_table_no_limit() {
13348 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13349
13350 let request = QueryTableRequest {
13352 id: Some(table_id),
13353 k: 0,
13354 filter: None,
13355 offset: None,
13356 version: None,
13357 ..Default::default()
13358 };
13359
13360 let bytes = namespace.query_table(request).await.unwrap();
13361
13362 let cursor = Cursor::new(bytes.to_vec());
13363 let reader = FileReader::try_new(cursor, None).unwrap();
13364 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13365 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13366 assert_eq!(total_rows, 3);
13367 }
13368
13369 #[tokio::test]
13370 async fn test_query_table_with_columns() {
13371 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13372
13373 let columns = Box::new(lance_namespace::models::QueryTableRequestColumns {
13374 column_names: Some(vec!["id".to_string()]),
13375 column_aliases: None,
13376 });
13377
13378 let request = QueryTableRequest {
13379 id: Some(table_id),
13380 k: 10,
13381 filter: None,
13382 offset: None,
13383 version: None,
13384 columns: Some(columns),
13385 ..Default::default()
13386 };
13387
13388 let bytes = namespace.query_table(request).await.unwrap();
13389
13390 let cursor = Cursor::new(bytes.to_vec());
13391 let reader = FileReader::try_new(cursor, None).unwrap();
13392 let schema = reader.schema();
13393 assert_eq!(schema.fields().len(), 1);
13394 assert_eq!(schema.field(0).name(), "id");
13395 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13396 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13397 assert_eq!(total_rows, 3);
13398 }
13399
13400 #[tokio::test]
13401 async fn test_count_table_rows_with_version() {
13402 use arrow::array::{Int32Array, StringArray};
13403 use arrow::ipc::writer::StreamWriter;
13404
13405 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13406
13407 let schema = create_test_schema();
13409 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13410 let arrow_schema = Arc::new(arrow_schema);
13411
13412 let id_array = Int32Array::from(vec![4, 5]);
13413 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13414 let batch = arrow::record_batch::RecordBatch::try_new(
13415 arrow_schema.clone(),
13416 vec![Arc::new(id_array), Arc::new(name_array)],
13417 )
13418 .unwrap();
13419
13420 let mut buffer = Vec::new();
13421 {
13422 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13423 writer.write(&batch).unwrap();
13424 writer.finish().unwrap();
13425 }
13426
13427 let request = InsertIntoTableRequest {
13428 id: Some(table_id.clone()),
13429 mode: None,
13430 ..Default::default()
13431 };
13432 namespace
13433 .insert_into_table(request, Bytes::from(buffer))
13434 .await
13435 .unwrap();
13436
13437 let count_req = CountTableRowsRequest {
13439 id: Some(table_id.clone()),
13440 version: Some(1),
13441 predicate: None,
13442 ..Default::default()
13443 };
13444 let count = namespace.count_table_rows(count_req).await.unwrap();
13445 assert_eq!(count, 3);
13446
13447 let count_req = CountTableRowsRequest {
13449 id: Some(table_id),
13450 version: None,
13451 predicate: None,
13452 ..Default::default()
13453 };
13454 let count = namespace.count_table_rows(count_req).await.unwrap();
13455 assert_eq!(count, 5);
13456 }
13457
13458 #[tokio::test]
13459 async fn test_query_table_with_version() {
13460 use arrow::array::{Int32Array, StringArray};
13461 use arrow::ipc::writer::StreamWriter;
13462
13463 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13464
13465 let schema = create_test_schema();
13467 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13468 let arrow_schema = Arc::new(arrow_schema);
13469
13470 let id_array = Int32Array::from(vec![4, 5]);
13471 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13472 let batch = arrow::record_batch::RecordBatch::try_new(
13473 arrow_schema.clone(),
13474 vec![Arc::new(id_array), Arc::new(name_array)],
13475 )
13476 .unwrap();
13477
13478 let mut buffer = Vec::new();
13479 {
13480 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13481 writer.write(&batch).unwrap();
13482 writer.finish().unwrap();
13483 }
13484
13485 let request = InsertIntoTableRequest {
13486 id: Some(table_id.clone()),
13487 mode: None,
13488 ..Default::default()
13489 };
13490 namespace
13491 .insert_into_table(request, Bytes::from(buffer))
13492 .await
13493 .unwrap();
13494
13495 let request = QueryTableRequest {
13497 id: Some(table_id.clone()),
13498 k: 100,
13499 filter: None,
13500 offset: None,
13501 version: Some(1),
13502 ..Default::default()
13503 };
13504
13505 let bytes = namespace.query_table(request).await.unwrap();
13506 let cursor = Cursor::new(bytes.to_vec());
13507 let reader = FileReader::try_new(cursor, None).unwrap();
13508 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13509 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13510 assert_eq!(total_rows, 3);
13511
13512 let request = QueryTableRequest {
13514 id: Some(table_id),
13515 k: 100,
13516 filter: None,
13517 offset: None,
13518 version: None,
13519 ..Default::default()
13520 };
13521
13522 let bytes = namespace.query_table(request).await.unwrap();
13523 let cursor = Cursor::new(bytes.to_vec());
13524 let reader = FileReader::try_new(cursor, None).unwrap();
13525 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13526 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13527 assert_eq!(total_rows, 5);
13528 }
13529
13530 async fn create_ns_with_vector_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
13533 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
13534 use arrow::ipc::writer::StreamWriter;
13535
13536 let (namespace, temp_dir) = create_test_namespace().await;
13537
13538 let arrow_schema = Arc::new(arrow::datatypes::Schema::new(vec![
13540 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
13541 arrow::datatypes::Field::new(
13542 "vector",
13543 arrow::datatypes::DataType::FixedSizeList(
13544 Arc::new(arrow::datatypes::Field::new(
13545 "item",
13546 arrow::datatypes::DataType::Float32,
13547 true,
13548 )),
13549 4,
13550 ),
13551 true,
13552 ),
13553 ]));
13554
13555 let id_array = Int32Array::from(vec![1, 2, 3]);
13556 let values = Float32Array::from(vec![
13557 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ]);
13561 let vector_array = FixedSizeListArray::try_new(
13562 Arc::new(arrow::datatypes::Field::new(
13563 "item",
13564 arrow::datatypes::DataType::Float32,
13565 true,
13566 )),
13567 4,
13568 Arc::new(values),
13569 None,
13570 )
13571 .unwrap();
13572
13573 let batch = arrow::record_batch::RecordBatch::try_new(
13574 arrow_schema.clone(),
13575 vec![Arc::new(id_array), Arc::new(vector_array)],
13576 )
13577 .unwrap();
13578
13579 let mut buffer = Vec::new();
13580 {
13581 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13582 writer.write(&batch).unwrap();
13583 writer.finish().unwrap();
13584 }
13585
13586 let table_name = "vector_table";
13588 let table_uri = format!("{}/{}.lance", temp_dir.to_str().unwrap(), table_name);
13589 let reader = arrow::record_batch::RecordBatchIterator::new(
13590 vec![Ok(batch)],
13591 arrow_schema.clone(),
13592 );
13593 Dataset::write(reader, &table_uri, None).await.unwrap();
13594
13595 let table_id = vec![table_name.to_string()];
13596 (namespace, temp_dir, table_id)
13597 }
13598
13599 #[tokio::test]
13600 async fn test_query_table_vector_search() {
13601 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13602
13603 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13604 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13605 multi_vector: None,
13606 });
13607
13608 let request = QueryTableRequest {
13609 id: Some(table_id),
13610 k: 2,
13611 vector,
13612 filter: None,
13613 offset: None,
13614 version: None,
13615 ..Default::default()
13616 };
13617
13618 let bytes = namespace.query_table(request).await.unwrap();
13619
13620 let cursor = Cursor::new(bytes.to_vec());
13621 let reader = FileReader::try_new(cursor, None).unwrap();
13622 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13623 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13624 assert_eq!(total_rows, 2);
13625 }
13626
13627 #[tokio::test]
13628 async fn test_query_table_vector_search_with_distance_type() {
13629 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13630
13631 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13632 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13633 multi_vector: None,
13634 });
13635
13636 let request = QueryTableRequest {
13637 id: Some(table_id),
13638 k: 3,
13639 vector,
13640 filter: None,
13641 offset: None,
13642 version: None,
13643 distance_type: Some("cosine".to_string()),
13644 ..Default::default()
13645 };
13646
13647 let bytes = namespace.query_table(request).await.unwrap();
13648
13649 let cursor = Cursor::new(bytes.to_vec());
13650 let reader = FileReader::try_new(cursor, None).unwrap();
13651 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13652 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13653 assert_eq!(total_rows, 3);
13654 }
13655
13656 #[tokio::test]
13657 async fn test_query_table_vector_search_with_filter() {
13658 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13659
13660 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13661 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13662 multi_vector: None,
13663 });
13664
13665 let request = QueryTableRequest {
13666 id: Some(table_id),
13667 k: 10,
13668 vector,
13669 filter: Some("id <= 2".to_string()),
13670 offset: None,
13671 version: None,
13672 ..Default::default()
13673 };
13674
13675 let bytes = namespace.query_table(request).await.unwrap();
13676
13677 let cursor = Cursor::new(bytes.to_vec());
13678 let reader = FileReader::try_new(cursor, None).unwrap();
13679 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13680 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13681 assert!(total_rows <= 2);
13682 }
13683
13684 #[tokio::test]
13685 async fn test_query_table_vector_search_with_nprobes_and_refine() {
13686 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13687
13688 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13689 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
13690 multi_vector: None,
13691 });
13692
13693 let request = QueryTableRequest {
13694 id: Some(table_id),
13695 k: 2,
13696 vector,
13697 filter: None,
13698 offset: None,
13699 version: None,
13700 nprobes: Some(1),
13701 refine_factor: Some(1),
13702 prefilter: Some(true),
13703 ..Default::default()
13704 };
13705
13706 let bytes = namespace.query_table(request).await.unwrap();
13707
13708 let cursor = Cursor::new(bytes.to_vec());
13709 let reader = FileReader::try_new(cursor, None).unwrap();
13710 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13711 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13712 assert_eq!(total_rows, 2);
13713 }
13714
13715 #[tokio::test]
13716 async fn test_namespace_id() {
13717 let (namespace, _temp_dir) = create_test_namespace().await;
13718 let id = namespace.namespace_id();
13719 assert!(id.contains("DirectoryNamespace"));
13720 assert!(id.contains("root"));
13721 }
13722
13723 #[tokio::test]
13724 async fn test_query_table_empty_table() {
13725 let (namespace, _temp_dir) = create_test_namespace().await;
13726
13727 let schema = create_test_schema();
13729 let ipc_data = create_test_ipc_data(&schema);
13730 let mut create_request = CreateTableRequest::new();
13731 create_request.id = Some(vec!["empty_table".to_string()]);
13732 namespace
13733 .create_table(create_request, bytes::Bytes::from(ipc_data))
13734 .await
13735 .unwrap();
13736
13737 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13739 single_vector: None,
13740 multi_vector: None,
13741 });
13742 let request = QueryTableRequest {
13743 id: Some(vec!["empty_table".to_string()]),
13744 k: 10,
13745 vector,
13746 ..Default::default()
13747 };
13748 let bytes = namespace.query_table(request).await.unwrap();
13749
13750 let cursor = Cursor::new(bytes.to_vec());
13751 let reader = FileReader::try_new(cursor, None).unwrap();
13752 let batches: Vec<_> = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
13753 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13754 assert_eq!(total_rows, 0, "empty table should yield no rows");
13755 }
13756
13757 #[tokio::test]
13758 async fn test_query_table_with_plain_filter_no_vector() {
13759 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13760
13761 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13763 single_vector: None,
13764 multi_vector: None,
13765 });
13766 let request = QueryTableRequest {
13767 id: Some(table_id),
13768 k: 0,
13769 vector,
13770 filter: Some("id > 1".to_string()),
13771 ..Default::default()
13772 };
13773 let bytes = namespace.query_table(request).await.unwrap();
13774
13775 let cursor = Cursor::new(bytes.to_vec());
13776 let reader = FileReader::try_new(cursor, None).unwrap();
13777 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13778 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13779 assert!(total_rows > 0);
13780 assert!(total_rows < 3);
13781 }
13782
13783 #[tokio::test]
13786 async fn test_update_full_table() {
13787 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13788
13789 let base_version = open_dataset(&namespace, &table_id[0])
13791 .await
13792 .version()
13793 .version;
13794
13795 let request = UpdateTableRequest {
13796 id: Some(table_id.clone()),
13797 updates: vec![vec!["name".to_string(), "'updated'".to_string()]],
13798 predicate: None,
13799 ..Default::default()
13800 };
13801
13802 let response = namespace.update_table(request).await.unwrap();
13803 assert_eq!(response.updated_rows, 3);
13804 assert!(response.version as u64 > base_version);
13805
13806 let count_req = CountTableRowsRequest {
13808 id: Some(table_id),
13809 version: None,
13810 predicate: Some("name = 'updated'".to_string()),
13811 ..Default::default()
13812 };
13813 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 3);
13814 }
13815
13816 #[tokio::test]
13817 async fn test_update_with_predicate() {
13818 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13819
13820 let request = UpdateTableRequest {
13821 id: Some(table_id.clone()),
13822 updates: vec![vec!["name".to_string(), "'matched'".to_string()]],
13823 predicate: Some("id > 1".to_string()),
13824 ..Default::default()
13825 };
13826
13827 let response = namespace.update_table(request).await.unwrap();
13828 assert_eq!(response.updated_rows, 2);
13829
13830 let untouched = CountTableRowsRequest {
13832 id: Some(table_id.clone()),
13833 version: None,
13834 predicate: Some("name = 'Alice'".to_string()),
13835 ..Default::default()
13836 };
13837 assert_eq!(namespace.count_table_rows(untouched).await.unwrap(), 1);
13838
13839 let touched = CountTableRowsRequest {
13840 id: Some(table_id),
13841 version: None,
13842 predicate: Some("name = 'matched'".to_string()),
13843 ..Default::default()
13844 };
13845 assert_eq!(namespace.count_table_rows(touched).await.unwrap(), 2);
13846 }
13847
13848 #[tokio::test]
13849 async fn test_update_invalid_expression_returns_invalid_input() {
13850 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13851
13852 let request = UpdateTableRequest {
13853 id: Some(table_id),
13854 updates: vec![vec!["name".to_string(), "no_such_column + 1".to_string()]],
13856 predicate: None,
13857 ..Default::default()
13858 };
13859
13860 let err = namespace.update_table(request).await.unwrap_err();
13861 let msg = err.to_string();
13862 assert!(
13863 msg.contains("Invalid input"),
13864 "expected Invalid input error, got: {}",
13865 msg
13866 );
13867 }
13868
13869 #[tokio::test]
13870 async fn test_update_rejects_duplicate_columns() {
13871 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13872
13873 let request = UpdateTableRequest {
13874 id: Some(table_id),
13875 updates: vec![
13876 vec!["name".to_string(), "'a'".to_string()],
13877 vec!["name".to_string(), "'b'".to_string()],
13878 ],
13879 predicate: None,
13880 ..Default::default()
13881 };
13882
13883 let err = namespace.update_table(request).await.unwrap_err();
13884 let msg = err.to_string();
13885 assert!(
13886 msg.contains("Invalid input") && msg.contains("more than once"),
13887 "expected duplicate column InvalidInput error, got: {}",
13888 msg
13889 );
13890 }
13891
13892 #[tokio::test]
13893 async fn test_delete_with_predicate() {
13894 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13895
13896 let request = DeleteFromTableRequest {
13897 id: Some(table_id.clone()),
13898 predicate: "id > 1".to_string(),
13899 ..Default::default()
13900 };
13901
13902 let response = namespace.delete_from_table(request).await.unwrap();
13903 assert!(response.version.is_some());
13904
13905 let count_req = CountTableRowsRequest {
13906 id: Some(table_id),
13907 version: None,
13908 predicate: None,
13909 ..Default::default()
13910 };
13911 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 1);
13913 }
13914
13915 #[tokio::test]
13916 async fn test_delete_empty_predicate_returns_invalid_input() {
13917 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13918
13919 let request = DeleteFromTableRequest {
13920 id: Some(table_id),
13921 predicate: " ".to_string(),
13922 ..Default::default()
13923 };
13924
13925 let err = namespace.delete_from_table(request).await.unwrap_err();
13926 let msg = err.to_string();
13927 assert!(
13928 msg.contains("Invalid input") && msg.contains("non-empty predicate"),
13929 "expected non-empty predicate InvalidInput error, got: {}",
13930 msg
13931 );
13932 }
13933
13934 #[tokio::test]
13935 async fn test_delete_table_not_found() {
13936 let (namespace, _temp_dir) = create_test_namespace().await;
13937
13938 let request = DeleteFromTableRequest {
13939 id: Some(vec!["does_not_exist".to_string()]),
13940 predicate: "id = 1".to_string(),
13941 ..Default::default()
13942 };
13943
13944 let err = namespace.delete_from_table(request).await.unwrap_err();
13945 let msg = err.to_string();
13946 assert!(
13947 msg.contains("Table not found"),
13948 "expected TableNotFound for missing table, got: {}",
13949 msg
13950 );
13951 }
13952
13953 #[tokio::test]
13954 async fn test_delete_invalid_predicate_returns_invalid_input() {
13955 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13956
13957 let request = DeleteFromTableRequest {
13961 id: Some(table_id),
13962 predicate: "no_such_column = 1".to_string(),
13963 ..Default::default()
13964 };
13965
13966 let err = namespace.delete_from_table(request).await.unwrap_err();
13967 let lance_core::Error::Namespace { source, .. } = &err else {
13968 panic!("expected a Namespace error, got: {}", err);
13969 };
13970 let ns_err = source
13971 .downcast_ref::<NamespaceError>()
13972 .expect("expected a NamespaceError source");
13973 assert_eq!(
13974 ns_err.code(),
13975 lance_namespace::ErrorCode::InvalidInput,
13976 "expected InvalidInput for an invalid delete predicate, got: {}",
13977 err
13978 );
13979 }
13980 }
13981
13982 #[tokio::test]
13983 async fn test_list_all_tables() {
13984 use lance_namespace::models::ListTablesRequest;
13985
13986 let (namespace, _temp_dir) = create_test_namespace().await;
13987 create_scalar_table(&namespace, "alpha").await;
13988 create_scalar_table(&namespace, "beta").await;
13989
13990 let request = ListTablesRequest {
13991 id: Some(vec![]),
13992 page_token: None,
13993 limit: None,
13994 ..Default::default()
13995 };
13996 let response = namespace.list_all_tables(request).await.unwrap();
13997 let mut tables = response.tables;
13998 tables.sort();
13999 assert_eq!(tables, vec!["alpha", "beta"]);
14000 }
14001
14002 #[tokio::test]
14003 async fn test_restore_table() {
14004 use lance_namespace::models::RestoreTableRequest;
14005
14006 let (namespace, _temp_dir) = create_test_namespace().await;
14007 create_scalar_table(&namespace, "users").await;
14008
14009 create_scalar_index(&namespace, "users", "users_id_idx").await;
14011
14012 let dataset = open_dataset(&namespace, "users").await;
14013 let current_version = dataset.version().version;
14014 assert!(current_version >= 2, "Should have at least 2 versions");
14015
14016 let mut restore_req = RestoreTableRequest::new(1);
14018 restore_req.id = Some(vec!["users".to_string()]);
14019 let response = namespace.restore_table(restore_req).await.unwrap();
14020
14021 assert!(
14023 response.transaction_id.is_some(),
14024 "restore_table should return a transaction_id"
14025 );
14026
14027 let dataset_after = open_dataset(&namespace, "users").await;
14029 assert!(
14030 dataset_after.version().version > current_version,
14031 "Restore should create a new version"
14032 );
14033 }
14034
14035 #[tokio::test]
14036 async fn test_alter_table_add_columns() {
14037 use lance_namespace::models::{
14038 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
14039 };
14040
14041 let (namespace, _temp_dir) = create_test_namespace().await;
14042
14043 let schema = create_test_schema();
14045 let ipc_data = create_test_ipc_data(&schema);
14046 let mut create_request = CreateTableRequest::new();
14047 create_request.id = Some(vec!["test_table".to_string()]);
14048 namespace
14049 .create_table(create_request, bytes::Bytes::from(ipc_data))
14050 .await
14051 .unwrap();
14052
14053 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
14055 new_col.expression = Some(Some("id * 2".to_string()));
14056 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
14057 add_request.id = Some(vec!["test_table".to_string()]);
14058
14059 let response = namespace
14060 .alter_table_add_columns(add_request)
14061 .await
14062 .unwrap();
14063 assert!(
14064 response.version > 1,
14065 "Version should increment after adding columns"
14066 );
14067
14068 let mut describe_request = DescribeTableRequest::new();
14070 describe_request.id = Some(vec!["test_table".to_string()]);
14071 describe_request.load_detailed_metadata = Some(true);
14072 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14073 assert!(describe_response.schema.is_some());
14074
14075 let resp_schema = describe_response.schema.unwrap();
14076 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14077 assert!(
14078 field_names.contains(&"doubled_id"),
14079 "Column 'doubled_id' should exist, got: {:?}",
14080 field_names
14081 );
14082 }
14083
14084 #[tokio::test]
14085 async fn test_update_table_schema_metadata() {
14086 use lance_namespace::models::UpdateTableSchemaMetadataRequest;
14087
14088 let (namespace, _temp_dir) = create_test_namespace().await;
14089 create_scalar_table(&namespace, "products").await;
14090
14091 let mut metadata = HashMap::new();
14092 metadata.insert("owner".to_string(), "team_a".to_string());
14093 metadata.insert("version".to_string(), "1.0".to_string());
14094
14095 let mut req = UpdateTableSchemaMetadataRequest::new();
14096 req.id = Some(vec!["products".to_string()]);
14097 req.metadata = Some(metadata.clone());
14098
14099 let response = namespace.update_table_schema_metadata(req).await.unwrap();
14100
14101 assert!(response.metadata.is_some());
14102 let returned = response.metadata.unwrap();
14103 assert_eq!(returned.get("owner"), Some(&"team_a".to_string()));
14104 assert_eq!(returned.get("version"), Some(&"1.0".to_string()));
14105 assert!(
14106 response.transaction_id.is_some(),
14107 "update_table_schema_metadata should return a transaction_id"
14108 );
14109 }
14110
14111 #[tokio::test]
14112 async fn test_alter_table_add_columns_missing_id() {
14113 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
14114
14115 let (namespace, _temp_dir) = create_test_namespace().await;
14116
14117 let new_col = AddColumnsEntry::new("col".to_string());
14118 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
14119 let result = namespace.alter_table_add_columns(request).await;
14120 assert!(result.is_err(), "Should fail when table ID is missing");
14121 }
14122
14123 #[tokio::test]
14124 async fn test_alter_table_alter_columns_rename() {
14125 use lance_namespace::models::{
14126 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
14127 };
14128
14129 let (namespace, _temp_dir) = create_test_namespace().await;
14130
14131 let schema = create_test_schema();
14133 let ipc_data = create_test_ipc_data(&schema);
14134 let mut create_request = CreateTableRequest::new();
14135 create_request.id = Some(vec!["test_table".to_string()]);
14136 namespace
14137 .create_table(create_request, bytes::Bytes::from(ipc_data))
14138 .await
14139 .unwrap();
14140
14141 let mut entry = AlterColumnsEntry::new("name".to_string());
14143 entry.rename = Some(Some("full_name".to_string()));
14144 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
14145 alter_request.id = Some(vec!["test_table".to_string()]);
14146
14147 let response = namespace
14148 .alter_table_alter_columns(alter_request)
14149 .await
14150 .unwrap();
14151 assert!(
14152 response.version > 1,
14153 "Version should increment after altering columns"
14154 );
14155
14156 let mut describe_request = DescribeTableRequest::new();
14158 describe_request.id = Some(vec!["test_table".to_string()]);
14159 describe_request.load_detailed_metadata = Some(true);
14160 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14161 assert!(describe_response.schema.is_some());
14162
14163 let resp_schema = describe_response.schema.unwrap();
14164 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14165 assert!(
14166 field_names.contains(&"full_name"),
14167 "Column should be renamed to 'full_name', got: {:?}",
14168 field_names
14169 );
14170 assert!(
14171 !field_names.contains(&"name"),
14172 "Old column 'name' should not exist, got: {:?}",
14173 field_names
14174 );
14175 }
14176
14177 #[tokio::test]
14178 async fn test_get_table_stats() {
14179 use lance_namespace::models::GetTableStatsRequest;
14180
14181 let (namespace, _temp_dir) = create_test_namespace().await;
14182 create_scalar_table(&namespace, "items").await;
14183 create_scalar_index(&namespace, "items", "items_id_idx").await;
14184
14185 let mut req = GetTableStatsRequest::new();
14186 req.id = Some(vec!["items".to_string()]);
14187
14188 let response = namespace.get_table_stats(req).await.unwrap();
14189 assert_eq!(response.num_rows, 3);
14190 assert_eq!(response.num_indices, 1);
14191 }
14192
14193 #[tokio::test]
14194 async fn test_explain_table_query_plan() {
14195 use lance_namespace::models::QueryTableRequestVector;
14196 use lance_namespace::models::{ExplainTableQueryPlanRequest, QueryTableRequest};
14197
14198 let (namespace, _temp_dir) = create_test_namespace().await;
14199 create_scalar_table(&namespace, "catalog").await;
14200
14201 let mut query = QueryTableRequest::new(1, QueryTableRequestVector::new());
14202 query.filter = Some("id > 1".to_string());
14203 query.columns = Some(Box::new(QueryTableRequestColumns {
14204 column_names: Some(vec!["id".to_string(), "name".to_string()]),
14205 column_aliases: None,
14206 }));
14207 query.with_row_id = Some(true);
14208
14209 let mut req = ExplainTableQueryPlanRequest::new(query);
14210 req.id = Some(vec!["catalog".to_string()]);
14211
14212 let plan_str = namespace.explain_table_query_plan(req).await.unwrap();
14213 assert_plan_contains_all(
14214 &plan_str,
14215 &[
14216 "ProjectionExec: expr=[id@0 as id, name@2 as name",
14217 "projection=[name], source=stream(_rowid)",
14218 "LanceRead: uri=",
14219 "projection=[id]",
14220 "row_id=true, row_addr=false",
14221 "full_filter=id > Int32(1)",
14222 "refine_filter=id > Int32(1)",
14223 ],
14224 "Filtered explain plan should preserve late materialization and filter pushdown",
14225 );
14226 }
14227
14228 #[tokio::test]
14229 async fn test_alter_table_alter_columns_missing_id() {
14230 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
14231
14232 let (namespace, _temp_dir) = create_test_namespace().await;
14233
14234 let entry = AlterColumnsEntry::new("name".to_string());
14235 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
14236 let result = namespace.alter_table_alter_columns(request).await;
14237 assert!(result.is_err(), "Should fail when table ID is missing");
14238 }
14239
14240 #[tokio::test]
14241 async fn test_alter_table_drop_columns() {
14242 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
14243
14244 let (namespace, _temp_dir) = create_test_namespace().await;
14245
14246 let schema = create_test_schema();
14248 let ipc_data = create_test_ipc_data(&schema);
14249 let mut create_request = CreateTableRequest::new();
14250 create_request.id = Some(vec!["test_table".to_string()]);
14251 namespace
14252 .create_table(create_request, bytes::Bytes::from(ipc_data))
14253 .await
14254 .unwrap();
14255
14256 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
14258 drop_request.id = Some(vec!["test_table".to_string()]);
14259
14260 let response = namespace
14261 .alter_table_drop_columns(drop_request)
14262 .await
14263 .unwrap();
14264 assert!(
14265 response.version > 1,
14266 "Version should increment after dropping columns"
14267 );
14268
14269 let mut describe_request = DescribeTableRequest::new();
14271 describe_request.id = Some(vec!["test_table".to_string()]);
14272 describe_request.load_detailed_metadata = Some(true);
14273 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14274 assert!(describe_response.schema.is_some());
14275
14276 let resp_schema = describe_response.schema.unwrap();
14277 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14278 assert!(
14279 !field_names.contains(&"name"),
14280 "Column 'name' should be dropped, got: {:?}",
14281 field_names
14282 );
14283 assert!(
14284 field_names.contains(&"id"),
14285 "Column 'id' should still exist, got: {:?}",
14286 field_names
14287 );
14288 }
14289
14290 #[tokio::test]
14291 async fn test_analyze_table_query_plan() {
14292 use lance_namespace::models::AnalyzeTableQueryPlanRequest;
14293 use lance_namespace::models::QueryTableRequestVector;
14294
14295 let (namespace, _temp_dir) = create_test_namespace().await;
14296 create_scalar_table(&namespace, "catalog").await;
14297
14298 let mut req = AnalyzeTableQueryPlanRequest::new(1, QueryTableRequestVector::new());
14299 req.id = Some(vec!["catalog".to_string()]);
14300 req.filter = Some("id > 0".to_string());
14301 req.columns = Some(Box::new(QueryTableRequestColumns {
14302 column_names: Some(vec!["id".to_string(), "name".to_string()]),
14303 column_aliases: None,
14304 }));
14305 req.with_row_id = Some(true);
14306
14307 let analysis_str = namespace.analyze_table_query_plan(req).await.unwrap();
14308 assert_plan_contains_all(
14309 &analysis_str,
14310 &[
14311 "AnalyzeExec verbose=true",
14312 "ProjectionExec: elapsed=",
14313 "expr=[id@0 as id, name@2 as name",
14314 "projection=[name], source=stream(_rowid)",
14315 "LanceRead: elapsed=",
14316 "projection=[id]",
14317 "row_id=true, row_addr=false",
14318 "full_filter=id > Int32(0)",
14319 "refine_filter=id > Int32(0)",
14320 "metrics=[output_rows=",
14321 ],
14322 "Filtered analyze plan should preserve late materialization and filter pushdown",
14323 );
14324 }
14325
14326 #[tokio::test]
14327 async fn test_dir_listing_no_extra_calls_without_migration() {
14328 let temp_dir = TempStdDir::default();
14329 let temp_path = temp_dir.to_str().unwrap();
14330 let root_uri = file_object_store_uri(temp_path);
14331 let listing_count = Arc::new(AtomicUsize::new(0));
14332 let session = build_listing_counting_session(listing_count.clone());
14333
14334 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
14336 .session(session.clone())
14337 .manifest_enabled(false)
14338 .dir_listing_enabled(true)
14339 .build()
14340 .await
14341 .unwrap();
14342
14343 let schema = create_test_schema();
14344 let ipc_data = create_test_ipc_data(&schema);
14345 let mut create_req = CreateTableRequest::new();
14346 create_req.id = Some(vec!["test_table".to_string()]);
14347 dir_only_ns
14348 .create_table(create_req, Bytes::from(ipc_data))
14349 .await
14350 .unwrap();
14351
14352 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
14354 .session(session)
14355 .manifest_enabled(true)
14356 .dir_listing_enabled(true)
14357 .dir_listing_to_manifest_migration_enabled(false)
14358 .build()
14359 .await
14360 .unwrap();
14361
14362 listing_count.store(0, Ordering::SeqCst);
14364
14365 let mut exists_req = TableExistsRequest::new();
14367 exists_req.id = Some(vec!["test_table".to_string()]);
14368 hybrid_ns.table_exists(exists_req).await.unwrap();
14369
14370 let count = listing_count.load(Ordering::SeqCst);
14371 assert_eq!(
14372 count, 1,
14373 "Expected exactly 1 listing call for table_exists \
14374 without migration mode, but got {}",
14375 count
14376 );
14377
14378 listing_count.store(0, Ordering::SeqCst);
14380
14381 let mut describe_req = DescribeTableRequest::new();
14382 describe_req.id = Some(vec!["test_table".to_string()]);
14383 hybrid_ns.describe_table(describe_req).await.unwrap();
14384
14385 let count = listing_count.load(Ordering::SeqCst);
14386 assert_eq!(
14387 count, 1,
14388 "Expected exactly 1 listing call for describe_table \
14389 without migration mode, but got {}",
14390 count
14391 );
14392 }
14393
14394 #[tokio::test]
14395 async fn test_build_and_root_reads_do_not_create_manifest() {
14396 let temp_dir = TempStdDir::default();
14397 let temp_path = temp_dir.to_str().unwrap();
14398 let manifest_path = std::path::Path::new(temp_path).join("__manifest");
14399
14400 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
14401 .manifest_enabled(false)
14402 .dir_listing_enabled(true)
14403 .build()
14404 .await
14405 .unwrap();
14406 create_scalar_table(&dir_only_ns, "catalog").await;
14407 assert!(!manifest_path.exists());
14408
14409 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14410 .manifest_enabled(true)
14411 .dir_listing_enabled(true)
14412 .build()
14413 .await
14414 .unwrap();
14415 assert!(!manifest_path.exists());
14416
14417 let mut exists_req = TableExistsRequest::new();
14418 exists_req.id = Some(vec!["catalog".to_string()]);
14419 namespace.table_exists(exists_req).await.unwrap();
14420 assert!(!manifest_path.exists());
14421
14422 let mut describe_req = DescribeTableRequest::new();
14423 describe_req.id = Some(vec!["catalog".to_string()]);
14424 namespace.describe_table(describe_req).await.unwrap();
14425 assert!(!manifest_path.exists());
14426
14427 let list_response = namespace
14428 .list_tables(ListTablesRequest {
14429 id: Some(vec![]),
14430 ..Default::default()
14431 })
14432 .await
14433 .unwrap();
14434 assert_eq!(list_response.tables, vec!["catalog".to_string()]);
14435 assert!(!manifest_path.exists());
14436
14437 let mut list_namespaces_req = ListNamespacesRequest::new();
14438 list_namespaces_req.id = Some(vec!["workspace".to_string()]);
14439 let err = namespace
14440 .list_namespaces(list_namespaces_req)
14441 .await
14442 .unwrap_err();
14443 assert!(err.to_string().contains("__manifest"));
14444 assert!(!manifest_path.exists());
14445
14446 let err = namespace
14447 .list_tables(ListTablesRequest {
14448 id: Some(vec!["workspace".to_string()]),
14449 ..Default::default()
14450 })
14451 .await
14452 .unwrap_err();
14453 assert!(err.to_string().contains("__manifest"));
14454 assert!(!manifest_path.exists());
14455
14456 let mut child_describe_req = DescribeTableRequest::new();
14457 child_describe_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
14458 let err = namespace
14459 .describe_table(child_describe_req)
14460 .await
14461 .unwrap_err();
14462 assert!(err.to_string().contains("__manifest"));
14463 assert!(!manifest_path.exists());
14464
14465 let mut child_exists_req = TableExistsRequest::new();
14466 child_exists_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
14467 let err = namespace.table_exists(child_exists_req).await.unwrap_err();
14468 assert!(err.to_string().contains("__manifest"));
14469 assert!(!manifest_path.exists());
14470
14471 let mut create_ns_req = CreateNamespaceRequest::new();
14472 create_ns_req.id = Some(vec!["workspace".to_string()]);
14473 namespace.create_namespace(create_ns_req).await.unwrap();
14474 assert!(manifest_path.exists());
14475 }
14476
14477 #[tokio::test]
14478 async fn test_migrate_updates_read_opened_legacy_manifest() {
14479 let temp_dir = TempStdDir::default();
14480 let temp_path = temp_dir.to_str().unwrap();
14481 create_legacy_manifest_without_primary_key_metadata(temp_path).await;
14482 assert!(!manifest_has_primary_key_metadata(temp_path).await);
14483
14484 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14485 .manifest_enabled(true)
14486 .dir_listing_enabled(true)
14487 .build()
14488 .await
14489 .unwrap();
14490 assert!(!manifest_has_primary_key_metadata(temp_path).await);
14491
14492 let migrated = namespace.migrate().await.unwrap();
14493 assert_eq!(migrated, 0);
14494 assert!(manifest_has_primary_key_metadata(temp_path).await);
14495 }
14496
14497 #[tokio::test]
14498 async fn test_describe_declared_table_checks_versions_only_when_requested() {
14499 let temp_dir = TempStdDir::default();
14500 let temp_path = temp_dir.to_str().unwrap();
14501 let root_uri = file_object_store_uri(temp_path);
14502 let listing_count = Arc::new(AtomicUsize::new(0));
14503 let session = build_listing_counting_session(listing_count.clone());
14504
14505 let namespace = DirectoryNamespaceBuilder::new(root_uri)
14506 .session(session)
14507 .manifest_enabled(false)
14508 .dir_listing_enabled(true)
14509 .build()
14510 .await
14511 .unwrap();
14512
14513 let mut declare_req = DeclareTableRequest::new();
14514 declare_req.id = Some(vec!["test_table".to_string()]);
14515 namespace.declare_table(declare_req).await.unwrap();
14516
14517 listing_count.store(0, Ordering::SeqCst);
14518
14519 let mut describe_req = DescribeTableRequest::new();
14520 describe_req.id = Some(vec!["test_table".to_string()]);
14521 let describe_response = namespace.describe_table(describe_req).await.unwrap();
14522
14523 assert_eq!(describe_response.is_only_declared, None);
14524 assert_eq!(
14525 listing_count.load(Ordering::SeqCst),
14526 1,
14527 "Default describe_table should only list the table directory"
14528 );
14529
14530 listing_count.store(0, Ordering::SeqCst);
14531
14532 let mut describe_req = DescribeTableRequest::new();
14533 describe_req.id = Some(vec!["test_table".to_string()]);
14534 describe_req.check_declared = Some(true);
14535 let describe_response = namespace.describe_table(describe_req).await.unwrap();
14536
14537 assert_eq!(describe_response.is_only_declared, Some(true));
14538 assert_eq!(
14539 listing_count.load(Ordering::SeqCst),
14540 2,
14541 "check_declared describe_table should list the table directory and _versions"
14542 );
14543 }
14544
14545 #[tokio::test]
14546 async fn test_dir_listing_extra_calls_with_migration() {
14547 let temp_dir = TempStdDir::default();
14548 let temp_path = temp_dir.to_str().unwrap();
14549 let root_uri = file_object_store_uri(temp_path);
14550 let listing_count = Arc::new(AtomicUsize::new(0));
14551 let session = build_listing_counting_session(listing_count.clone());
14552
14553 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
14555 .session(session.clone())
14556 .manifest_enabled(false)
14557 .dir_listing_enabled(true)
14558 .build()
14559 .await
14560 .unwrap();
14561
14562 let schema = create_test_schema();
14563 let ipc_data = create_test_ipc_data(&schema);
14564 let mut create_req = CreateTableRequest::new();
14565 create_req.id = Some(vec!["test_table".to_string()]);
14566 dir_only_ns
14567 .create_table(create_req, Bytes::from(ipc_data))
14568 .await
14569 .unwrap();
14570
14571 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
14572 .session(session)
14573 .manifest_enabled(true)
14574 .dir_listing_enabled(true)
14575 .dir_listing_to_manifest_migration_enabled(true)
14576 .build()
14577 .await
14578 .unwrap();
14579
14580 listing_count.store(0, Ordering::SeqCst);
14588
14589 let mut exists_req = TableExistsRequest::new();
14590 exists_req.id = Some(vec!["test_table".to_string()]);
14591 hybrid_ns.table_exists(exists_req).await.unwrap();
14592
14593 let count = listing_count.load(Ordering::SeqCst);
14594 assert_eq!(
14595 count, 2,
14596 "Expected 2 listing calls for table_exists with migration mode \
14597 (absent-__manifest probe + table directory fallback), but got {}",
14598 count
14599 );
14600
14601 listing_count.store(0, Ordering::SeqCst);
14604
14605 let mut describe_req = DescribeTableRequest::new();
14606 describe_req.id = Some(vec!["test_table".to_string()]);
14607 hybrid_ns.describe_table(describe_req).await.unwrap();
14608
14609 let count = listing_count.load(Ordering::SeqCst);
14610 assert_eq!(
14611 count, 2,
14612 "Expected 2 listing calls for describe_table with migration mode \
14613 (absent-__manifest probe + table directory fallback), but got {}",
14614 count
14615 );
14616 }
14617
14618 #[tokio::test]
14619 async fn test_manifest_reload_observes_new_version_from_other_namespace() {
14620 let temp_dir = TempStdDir::default();
14621 let temp_path = temp_dir.to_str().unwrap();
14622
14623 let namespace_a = DirectoryNamespaceBuilder::new(temp_path)
14624 .manifest_enabled(true)
14625 .dir_listing_enabled(false)
14626 .build()
14627 .await
14628 .unwrap();
14629 create_scalar_table(&namespace_a, "alpha").await;
14630
14631 let namespace_b = DirectoryNamespaceBuilder::new(temp_path)
14632 .manifest_enabled(true)
14633 .dir_listing_enabled(false)
14634 .build()
14635 .await
14636 .unwrap();
14637 create_scalar_table(&namespace_b, "beta").await;
14638
14639 let response = namespace_a
14640 .list_tables(ListTablesRequest {
14641 id: Some(vec![]),
14642 ..Default::default()
14643 })
14644 .await
14645 .unwrap();
14646
14647 let mut tables = response.tables;
14648 tables.sort();
14649 assert_eq!(tables, vec!["alpha", "beta"]);
14650 }
14651
14652 #[tokio::test]
14653 async fn test_migration_not_found_errors_include_table_id() {
14654 let temp_dir = TempStdDir::default();
14655 let temp_path = temp_dir.to_str().unwrap();
14656
14657 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14658 .manifest_enabled(true)
14659 .dir_listing_enabled(true)
14660 .dir_listing_to_manifest_migration_enabled(true)
14661 .build()
14662 .await
14663 .unwrap();
14664
14665 let mut exists_req = TableExistsRequest::new();
14666 exists_req.id = Some(vec!["missing_table".to_string()]);
14667 let err = namespace.table_exists(exists_req).await.unwrap_err();
14668 assert!(matches!(err, Error::Namespace { .. }));
14669 let err_msg = err.to_string();
14670 assert!(err_msg.contains("Table not found"));
14671 assert!(err_msg.contains("table id 'missing_table'"));
14672
14673 let mut describe_req = DescribeTableRequest::new();
14674 describe_req.id = Some(vec!["missing_table".to_string()]);
14675 let err = namespace.describe_table(describe_req).await.unwrap_err();
14676 assert!(matches!(err, Error::Namespace { .. }));
14677 let err_msg = err.to_string();
14678 assert!(err_msg.contains("Table not found"));
14679 assert!(err_msg.contains("table id 'missing_table'"));
14680 }
14681
14682 #[tokio::test]
14683 async fn test_manifest_not_found_errors_include_full_table_id() {
14684 use lance_namespace::models::CreateNamespaceRequest;
14685
14686 let temp_dir = TempStdDir::default();
14687 let temp_path = temp_dir.to_str().unwrap();
14688
14689 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14690 .manifest_enabled(true)
14691 .dir_listing_enabled(true)
14692 .build()
14693 .await
14694 .unwrap();
14695
14696 let mut create_ns_req = CreateNamespaceRequest::new();
14697 create_ns_req.id = Some(vec!["workspace".to_string()]);
14698 namespace.create_namespace(create_ns_req).await.unwrap();
14699
14700 let missing_table_id = vec!["workspace".to_string(), "missing_table".to_string()];
14701
14702 let mut exists_req = TableExistsRequest::new();
14703 exists_req.id = Some(missing_table_id.clone());
14704 let err = namespace.table_exists(exists_req).await.unwrap_err();
14705 assert!(matches!(err, Error::Namespace { .. }));
14706 let err_msg = err.to_string();
14707 assert!(err_msg.contains("Table not found"));
14708 assert!(err_msg.contains("table id 'workspace$missing_table'"));
14709
14710 let mut describe_req = DescribeTableRequest::new();
14711 describe_req.id = Some(missing_table_id);
14712 let err = namespace.describe_table(describe_req).await.unwrap_err();
14713 assert!(matches!(err, Error::Namespace { .. }));
14714 let err_msg = err.to_string();
14715 assert!(err_msg.contains("Table not found"));
14716 assert!(err_msg.contains("table id 'workspace$missing_table'"));
14717 }
14718
14719 async fn create_tagged_test_table(
14722 versions: u32,
14723 ) -> (Arc<DirectoryNamespace>, TempStdDir, Vec<String>) {
14724 use arrow::array::{Int32Array, RecordBatchIterator};
14725 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
14726 use arrow::record_batch::RecordBatch;
14727 use lance::dataset::{Dataset, WriteMode, WriteParams};
14728
14729 assert!(versions >= 1, "versions must be at least 1");
14730
14731 let temp_dir = TempStdDir::default();
14732 let temp_path = temp_dir.to_str().unwrap();
14733
14734 let namespace = Arc::new(
14735 DirectoryNamespaceBuilder::new(temp_path)
14736 .build()
14737 .await
14738 .unwrap(),
14739 );
14740 let table_id = vec!["tag_table".to_string()];
14741 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
14742 "id",
14743 DataType::Int32,
14744 false,
14745 )]));
14746 let initial_batch = RecordBatch::try_new(
14747 arrow_schema.clone(),
14748 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
14749 )
14750 .unwrap();
14751 let batches = RecordBatchIterator::new(vec![Ok(initial_batch)], arrow_schema.clone());
14752 let write_params = WriteParams {
14753 mode: WriteMode::Create,
14754 ..Default::default()
14755 };
14756
14757 let mut dataset = Dataset::write_into_namespace(
14758 batches,
14759 namespace.clone() as Arc<dyn LanceNamespace>,
14760 table_id.clone(),
14761 Some(write_params),
14762 )
14763 .await
14764 .unwrap();
14765
14766 for i in 1..versions {
14767 let value_start = (i as i32) * 10;
14768 let batch = RecordBatch::try_new(
14769 arrow_schema.clone(),
14770 vec![Arc::new(Int32Array::from(vec![
14771 value_start,
14772 value_start + 1,
14773 ]))],
14774 )
14775 .unwrap();
14776 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
14777 dataset.append(batches, None).await.unwrap();
14778 }
14779
14780 (namespace, temp_dir, table_id)
14781 }
14782
14783 fn namespace_code(err: &Error) -> Option<ErrorCode> {
14785 match err {
14786 Error::Namespace { source, .. } => {
14787 source.downcast_ref::<NamespaceError>().map(|e| e.code())
14788 }
14789 _ => None,
14790 }
14791 }
14792
14793 #[tokio::test]
14794 async fn test_create_and_list_branches() {
14795 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14796
14797 namespace
14798 .create_table_branch(CreateTableBranchRequest {
14799 id: Some(table_id.clone()),
14800 name: "dev".to_string(),
14801 ..Default::default()
14802 })
14803 .await
14804 .unwrap();
14805 namespace
14806 .create_table_branch(CreateTableBranchRequest {
14807 id: Some(table_id.clone()),
14808 name: "staging".to_string(),
14809 ..Default::default()
14810 })
14811 .await
14812 .unwrap();
14813
14814 let resp = namespace
14815 .list_table_branches(ListTableBranchesRequest {
14816 id: Some(table_id.clone()),
14817 ..Default::default()
14818 })
14819 .await
14820 .unwrap();
14821 assert_eq!(
14822 resp.branches.len(),
14823 2,
14824 "expected 2 branches, got: {:?}",
14825 resp.branches
14826 );
14827 assert!(resp.branches.contains_key("dev"));
14828 assert!(resp.branches.contains_key("staging"));
14829 assert!(resp.page_token.is_none());
14830
14831 namespace
14833 .delete_table_branch(DeleteTableBranchRequest {
14834 id: Some(table_id.clone()),
14835 name: "dev".to_string(),
14836 ..Default::default()
14837 })
14838 .await
14839 .unwrap();
14840
14841 let resp = namespace
14842 .list_table_branches(ListTableBranchesRequest {
14843 id: Some(table_id),
14844 ..Default::default()
14845 })
14846 .await
14847 .unwrap();
14848 assert_eq!(resp.branches.len(), 1, "expected 1 branch after delete");
14849 assert!(!resp.branches.contains_key("dev"));
14850 assert!(resp.branches.contains_key("staging"));
14851 }
14852
14853 #[tokio::test]
14854 async fn test_create_branch_from_version() {
14855 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14856
14857 namespace
14859 .create_table_branch(CreateTableBranchRequest {
14860 id: Some(table_id.clone()),
14861 name: "from-v1".to_string(),
14862 from_version: Some(1),
14863 ..Default::default()
14864 })
14865 .await
14866 .unwrap();
14867
14868 let resp = namespace
14869 .list_table_branches(ListTableBranchesRequest {
14870 id: Some(table_id),
14871 ..Default::default()
14872 })
14873 .await
14874 .unwrap();
14875 let branch = resp
14876 .branches
14877 .get("from-v1")
14878 .expect("forked branch should be listed");
14879 assert_eq!(
14880 branch.parent_version, 1,
14881 "branch should fork from version 1"
14882 );
14883 assert!(
14884 branch.parent_branch.is_none(),
14885 "a branch forked from main has no parent branch"
14886 );
14887 }
14888
14889 #[tokio::test]
14894 async fn test_create_branch_from_other_branch() {
14895 use lance::dataset::builder::DatasetBuilder;
14896
14897 let (namespace, _temp_dir) = create_test_namespace().await;
14898 create_scalar_table(&namespace, "users").await; create_branch_with_commits(&namespace, "users", "dev", 1).await;
14901 let main_ds = open_dataset(&namespace, "users").await;
14903 append_scalar_version(main_ds.uri(), 500).await; namespace
14906 .create_table_branch(CreateTableBranchRequest {
14907 id: Some(vec!["users".to_string()]),
14908 name: "child".to_string(),
14909 from_branch: Some("dev".to_string()),
14910 from_version: Some(2),
14911 ..Default::default()
14912 })
14913 .await
14914 .unwrap();
14915
14916 let child_ds = DatasetBuilder::from_uri(main_ds.uri())
14917 .with_branch("child", None)
14918 .load()
14919 .await
14920 .unwrap();
14921 let ids = scan_id_column(&child_ds).await;
14922 assert!(
14923 ids.contains(&100) && ids.contains(&101),
14924 "child must contain dev's appended rows, got: {:?}",
14925 ids
14926 );
14927 assert!(
14928 !ids.contains(&500),
14929 "child must not contain main's diverged rows, got: {:?}",
14930 ids
14931 );
14932
14933 let listed = namespace
14935 .list_table_branches(ListTableBranchesRequest {
14936 id: Some(vec!["users".to_string()]),
14937 ..Default::default()
14938 })
14939 .await
14940 .unwrap();
14941 assert_eq!(
14942 listed
14943 .branches
14944 .get("child")
14945 .unwrap()
14946 .parent_branch
14947 .as_deref(),
14948 Some("dev")
14949 );
14950 }
14951
14952 #[tokio::test]
14953 async fn test_create_existing_branch_conflict() {
14954 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14955
14956 namespace
14957 .create_table_branch(CreateTableBranchRequest {
14958 id: Some(table_id.clone()),
14959 name: "dev".to_string(),
14960 ..Default::default()
14961 })
14962 .await
14963 .unwrap();
14964
14965 let err = namespace
14966 .create_table_branch(CreateTableBranchRequest {
14967 id: Some(table_id),
14968 name: "dev".to_string(),
14969 ..Default::default()
14970 })
14971 .await
14972 .unwrap_err();
14973 assert_eq!(
14974 namespace_code(&err),
14975 Some(ErrorCode::TableBranchAlreadyExists),
14976 "expected TableBranchAlreadyExists, got: {}",
14977 err
14978 );
14979 assert!(
14980 err.to_string().to_lowercase().contains("already exists"),
14981 "expected already-exists message, got: {}",
14982 err
14983 );
14984 }
14985
14986 #[tokio::test]
14987 async fn test_delete_unknown_branch() {
14988 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14989
14990 let err = namespace
14991 .delete_table_branch(DeleteTableBranchRequest {
14992 id: Some(table_id),
14993 name: "does-not-exist".to_string(),
14994 ..Default::default()
14995 })
14996 .await
14997 .unwrap_err();
14998 assert_eq!(
14999 namespace_code(&err),
15000 Some(ErrorCode::TableBranchNotFound),
15001 "expected TableBranchNotFound, got: {}",
15002 err
15003 );
15004 assert!(
15005 err.to_string().to_lowercase().contains("not found"),
15006 "expected not-found message, got: {}",
15007 err
15008 );
15009 }
15010
15011 #[tokio::test]
15012 async fn test_delete_referenced_branch_conflict() {
15013 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15014
15015 namespace
15017 .create_table_branch(CreateTableBranchRequest {
15018 id: Some(table_id.clone()),
15019 name: "parent".to_string(),
15020 ..Default::default()
15021 })
15022 .await
15023 .unwrap();
15024 namespace
15025 .create_table_branch(CreateTableBranchRequest {
15026 id: Some(table_id.clone()),
15027 name: "child".to_string(),
15028 from_branch: Some("parent".to_string()),
15029 ..Default::default()
15030 })
15031 .await
15032 .unwrap();
15033
15034 let listed = namespace
15036 .list_table_branches(ListTableBranchesRequest {
15037 id: Some(table_id.clone()),
15038 ..Default::default()
15039 })
15040 .await
15041 .unwrap();
15042 let child = listed
15043 .branches
15044 .get("child")
15045 .expect("child branch should be listed");
15046 assert_eq!(
15047 child.parent_branch.as_deref(),
15048 Some("parent"),
15049 "child should record parent branch as its fork point"
15050 );
15051 assert!(
15052 child.parent_version >= 1,
15053 "child should record the parent version it forked from, got {}",
15054 child.parent_version
15055 );
15056
15057 let err = namespace
15060 .delete_table_branch(DeleteTableBranchRequest {
15061 id: Some(table_id),
15062 name: "parent".to_string(),
15063 ..Default::default()
15064 })
15065 .await
15066 .unwrap_err();
15067 assert_eq!(
15068 namespace_code(&err),
15069 Some(ErrorCode::InvalidInput),
15070 "expected InvalidInput for deleting a referenced branch, got: {}",
15071 err
15072 );
15073 assert!(
15074 err.to_string().to_lowercase().contains("referenced"),
15075 "error should explain the branch is still referenced, got: {}",
15076 err
15077 );
15078 }
15079
15080 #[tokio::test]
15081 async fn test_branch_name_required() {
15082 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15083
15084 let create_err = namespace
15085 .create_table_branch(CreateTableBranchRequest {
15086 id: Some(table_id.clone()),
15087 name: String::new(),
15088 ..Default::default()
15089 })
15090 .await
15091 .unwrap_err();
15092 assert_eq!(
15093 namespace_code(&create_err),
15094 Some(ErrorCode::InvalidInput),
15095 "empty name on create should be InvalidInput, got: {}",
15096 create_err
15097 );
15098 assert!(
15099 create_err
15100 .to_string()
15101 .to_lowercase()
15102 .contains("must not be empty")
15103 );
15104
15105 let delete_err = namespace
15106 .delete_table_branch(DeleteTableBranchRequest {
15107 id: Some(table_id),
15108 name: String::new(),
15109 ..Default::default()
15110 })
15111 .await
15112 .unwrap_err();
15113 assert_eq!(
15114 namespace_code(&delete_err),
15115 Some(ErrorCode::InvalidInput),
15116 "empty name on delete should be InvalidInput, got: {}",
15117 delete_err
15118 );
15119 assert!(
15120 delete_err
15121 .to_string()
15122 .to_lowercase()
15123 .contains("must not be empty")
15124 );
15125 }
15126
15127 #[tokio::test]
15128 async fn test_create_branch_rejects_negative_from_version() {
15129 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15130
15131 let err = namespace
15132 .create_table_branch(CreateTableBranchRequest {
15133 id: Some(table_id),
15134 name: "dev".to_string(),
15135 from_version: Some(-1),
15136 ..Default::default()
15137 })
15138 .await
15139 .unwrap_err();
15140 assert_eq!(
15141 namespace_code(&err),
15142 Some(ErrorCode::InvalidInput),
15143 "negative from_version should be InvalidInput, got: {}",
15144 err
15145 );
15146 assert!(err.to_string().to_lowercase().contains("from_version"));
15147 }
15148
15149 #[tokio::test]
15150 async fn test_create_branch_nonexistent_from_version() {
15151 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15152
15153 let err = namespace
15156 .create_table_branch(CreateTableBranchRequest {
15157 id: Some(table_id),
15158 name: "dev".to_string(),
15159 from_version: Some(999),
15160 ..Default::default()
15161 })
15162 .await
15163 .unwrap_err();
15164 assert_eq!(
15165 namespace_code(&err),
15166 Some(ErrorCode::InvalidInput),
15167 "non-existent from_version should map to InvalidInput, got: {}",
15168 err
15169 );
15170 assert!(
15171 err.to_string().to_lowercase().contains("does not exist"),
15172 "error should name the missing source, got: {}",
15173 err
15174 );
15175 }
15176
15177 #[tokio::test]
15178 async fn test_create_and_list_tags() {
15179 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15180
15181 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15182 req.id = Some(table_id.clone());
15183 namespace.create_table_tag(req).await.unwrap();
15184
15185 let mut req = CreateTableTagRequest::new("v2".to_string(), 2);
15186 req.id = Some(table_id.clone());
15187 namespace.create_table_tag(req).await.unwrap();
15188
15189 let mut list_req = ListTableTagsRequest::new();
15190 list_req.id = Some(table_id);
15191 let resp = namespace.list_table_tags(list_req).await.unwrap();
15192
15193 assert_eq!(resp.tags.len(), 2, "expected 2 tags, got: {:?}", resp.tags);
15194 assert_eq!(resp.tags.get("v1").unwrap().version, 1);
15195 assert_eq!(resp.tags.get("v2").unwrap().version, 2);
15196 assert!(resp.page_token.is_none());
15197 }
15198
15199 #[tokio::test]
15200 async fn test_create_existing_tag_conflict() {
15201 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15202
15203 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15204 req.id = Some(table_id.clone());
15205 namespace.create_table_tag(req).await.unwrap();
15206
15207 let mut req = CreateTableTagRequest::new("v1".to_string(), 2);
15208 req.id = Some(table_id);
15209 let err = namespace.create_table_tag(req).await.unwrap_err();
15210 assert!(
15211 err.to_string().to_lowercase().contains("already exists"),
15212 "expected already-exists error, got: {}",
15213 err
15214 );
15215 }
15216
15217 #[tokio::test]
15218 async fn test_get_tag_version() {
15219 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15220
15221 let mut req = CreateTableTagRequest::new("release".to_string(), 2);
15222 req.id = Some(table_id.clone());
15223 namespace.create_table_tag(req).await.unwrap();
15224
15225 let mut get_req = GetTableTagVersionRequest::new("release".to_string());
15226 get_req.id = Some(table_id);
15227 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15228 assert_eq!(resp.version, 2);
15229 assert_eq!(resp.branch, None);
15230 }
15231
15232 #[tokio::test]
15233 async fn test_get_unknown_tag() {
15234 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15235
15236 let mut get_req = GetTableTagVersionRequest::new("does-not-exist".to_string());
15237 get_req.id = Some(table_id);
15238 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
15239 assert!(
15240 err.to_string().to_lowercase().contains("not found"),
15241 "expected not-found error, got: {}",
15242 err
15243 );
15244 }
15245
15246 #[tokio::test]
15247 async fn test_update_tag_to_new_version() {
15248 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15249
15250 let mut req = CreateTableTagRequest::new("rolling".to_string(), 1);
15251 req.id = Some(table_id.clone());
15252 namespace.create_table_tag(req).await.unwrap();
15253
15254 let mut update_req = UpdateTableTagRequest::new("rolling".to_string(), 3);
15255 update_req.id = Some(table_id.clone());
15256 namespace.update_table_tag(update_req).await.unwrap();
15257
15258 let mut get_req = GetTableTagVersionRequest::new("rolling".to_string());
15259 get_req.id = Some(table_id);
15260 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15261 assert_eq!(resp.version, 3);
15262 }
15263
15264 #[tokio::test]
15265 async fn test_update_unknown_tag() {
15266 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15267
15268 let mut update_req = UpdateTableTagRequest::new("ghost".to_string(), 1);
15269 update_req.id = Some(table_id);
15270 let err = namespace.update_table_tag(update_req).await.unwrap_err();
15271 assert!(
15272 err.to_string().to_lowercase().contains("not found"),
15273 "expected not-found error, got: {}",
15274 err
15275 );
15276 }
15277
15278 #[tokio::test]
15279 async fn test_delete_tag() {
15280 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15281
15282 let mut req = CreateTableTagRequest::new("doomed".to_string(), 1);
15283 req.id = Some(table_id.clone());
15284 namespace.create_table_tag(req).await.unwrap();
15285
15286 let mut delete_req = DeleteTableTagRequest::new("doomed".to_string());
15287 delete_req.id = Some(table_id.clone());
15288 namespace.delete_table_tag(delete_req).await.unwrap();
15289
15290 let mut list_req = ListTableTagsRequest::new();
15291 list_req.id = Some(table_id.clone());
15292 let resp = namespace.list_table_tags(list_req).await.unwrap();
15293 assert!(resp.tags.is_empty(), "tag should be removed after delete");
15294
15295 let mut get_req = GetTableTagVersionRequest::new("doomed".to_string());
15297 get_req.id = Some(table_id);
15298 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
15299 assert!(err.to_string().to_lowercase().contains("not found"));
15300 }
15301
15302 #[tokio::test]
15303 async fn test_delete_unknown_tag() {
15304 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15305
15306 let mut delete_req = DeleteTableTagRequest::new("nope".to_string());
15307 delete_req.id = Some(table_id);
15308 let err = namespace.delete_table_tag(delete_req).await.unwrap_err();
15309 assert!(
15310 err.to_string().to_lowercase().contains("not found"),
15311 "expected not-found error, got: {}",
15312 err
15313 );
15314 }
15315
15316 #[tokio::test]
15317 async fn test_create_tag_invalid_version() {
15318 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15319
15320 let mut req = CreateTableTagRequest::new("v0".to_string(), 0);
15322 req.id = Some(table_id.clone());
15323 let err = namespace.create_table_tag(req).await.unwrap_err();
15324 assert!(
15325 err.to_string().to_lowercase().contains("positive"),
15326 "expected positive-version error, got: {}",
15327 err
15328 );
15329
15330 let mut req = CreateTableTagRequest::new(String::new(), 1);
15332 req.id = Some(table_id);
15333 let err = namespace.create_table_tag(req).await.unwrap_err();
15334 assert!(
15335 err.to_string().to_lowercase().contains("must not be empty"),
15336 "expected empty-tag-name error, got: {}",
15337 err
15338 );
15339 }
15340
15341 #[tokio::test]
15342 async fn test_create_tag_table_not_found() {
15343 let (namespace, _temp_dir) = create_test_namespace().await;
15344
15345 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15346 req.id = Some(vec!["does_not_exist".to_string()]);
15347 let err = namespace.create_table_tag(req).await.unwrap_err();
15348 let msg = err.to_string();
15349 assert!(
15350 msg.contains("Table") && msg.to_lowercase().contains("not found"),
15351 "expected TableNotFound error, got: {}",
15352 err
15353 );
15354 }
15355 #[tokio::test]
15356 async fn test_alter_table_drop_columns_missing_id() {
15357 use lance_namespace::models::AlterTableDropColumnsRequest;
15358
15359 let (namespace, _temp_dir) = create_test_namespace().await;
15360
15361 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
15362 let result = namespace.alter_table_drop_columns(request).await;
15363 assert!(result.is_err(), "Should fail when table ID is missing");
15364 }
15365
15366 #[tokio::test]
15367 async fn test_alter_table_drop_columns_nonexistent_table() {
15368 use lance_namespace::models::AlterTableDropColumnsRequest;
15369
15370 let (namespace, _temp_dir) = create_test_namespace().await;
15371
15372 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
15373 request.id = Some(vec!["nonexistent".to_string()]);
15374 let result = namespace.alter_table_drop_columns(request).await;
15375 assert!(result.is_err(), "Should fail when table does not exist");
15376 }
15377
15378 #[tokio::test]
15379 async fn test_create_branch_on_managed_dataset_succeeds() {
15380 use lance::dataset::builder::DatasetBuilder;
15381
15382 let temp = TempStdDir::default();
15383 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
15384 let table_id = vec!["t".to_string()];
15385 let mut main = create_managed_table(&ns, &table_id).await;
15386
15387 let fork_version = main.version().version;
15388 let branch = main
15389 .create_branch("exp", fork_version, None)
15390 .await
15391 .expect("create_branch failed");
15392 assert_eq!(branch.manifest.branch.as_deref(), Some("exp"));
15393 assert_eq!(scan_id_column(&branch).await, vec![1, 2]);
15394
15395 let reopened = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
15396 .await
15397 .unwrap()
15398 .with_branch("exp", None)
15399 .load()
15400 .await
15401 .expect("reopen branch failed");
15402 assert_eq!(scan_id_column(&reopened).await, vec![1, 2]);
15403 }
15404
15405 #[tokio::test]
15406 async fn test_alter_transaction_set_status() {
15407 use lance_namespace::models::{
15408 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15409 DescribeTransactionRequest,
15410 };
15411
15412 let (namespace, _temp_dir) = create_test_namespace().await;
15413 create_scalar_table(&namespace, "users").await;
15414 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15415 .await
15416 .expect("create_scalar_index should return a transaction id");
15417
15418 let describe_resp = namespace
15420 .describe_transaction(DescribeTransactionRequest {
15421 id: Some(vec!["users".to_string(), txn_id.clone()]),
15422 ..Default::default()
15423 })
15424 .await
15425 .unwrap();
15426 assert_eq!(describe_resp.status, "SUCCEEDED");
15427
15428 let response = namespace
15430 .alter_transaction(AlterTransactionRequest {
15431 id: Some(vec!["users".to_string(), txn_id.clone()]),
15432 actions: vec![AlterTransactionAction {
15433 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15434 status: Some("Canceled".to_string()),
15435 })),
15436 set_property_action: None,
15437 unset_property_action: None,
15438 }],
15439 ..Default::default()
15440 })
15441 .await
15442 .unwrap();
15443 assert_eq!(response.status, "Canceled");
15444 assert!(response.properties.is_some());
15445 let props = response.properties.unwrap();
15446 assert_eq!(props.get("uuid"), Some(&txn_id));
15447 assert_eq!(props.get("operation"), Some(&"CreateIndex".to_string()));
15448 }
15449
15450 #[tokio::test]
15451 async fn test_alter_transaction_set_property() {
15452 use lance_namespace::models::{
15453 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15454 };
15455
15456 let (namespace, _temp_dir) = create_test_namespace().await;
15457 create_scalar_table(&namespace, "users").await;
15458 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15459 .await
15460 .expect("create_scalar_index should return a transaction id");
15461
15462 let response = namespace
15463 .alter_transaction(AlterTransactionRequest {
15464 id: Some(vec!["users".to_string(), txn_id.clone()]),
15465 actions: vec![AlterTransactionAction {
15466 set_status_action: None,
15467 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15468 key: Some("custom_key".to_string()),
15469 value: Some("custom_value".to_string()),
15470 mode: None,
15471 })),
15472 unset_property_action: None,
15473 }],
15474 ..Default::default()
15475 })
15476 .await
15477 .unwrap();
15478 assert_eq!(response.status, "SUCCEEDED");
15479 let props = response.properties.unwrap();
15480 assert_eq!(props.get("custom_key"), Some(&"custom_value".to_string()));
15481 }
15482
15483 #[tokio::test]
15484 async fn test_alter_transaction_set_property_fail_mode() {
15485 use lance_namespace::models::{
15486 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15487 };
15488
15489 let (namespace, _temp_dir) = create_test_namespace().await;
15490 create_scalar_table(&namespace, "users").await;
15491 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15492 .await
15493 .expect("create_scalar_index should return a transaction id");
15494
15495 namespace
15497 .alter_transaction(AlterTransactionRequest {
15498 id: Some(vec!["users".to_string(), txn_id.clone()]),
15499 actions: vec![AlterTransactionAction {
15500 set_status_action: None,
15501 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15502 key: Some("custom_key".to_string()),
15503 value: Some("initial_value".to_string()),
15504 mode: None,
15505 })),
15506 unset_property_action: None,
15507 }],
15508 ..Default::default()
15509 })
15510 .await
15511 .unwrap();
15512
15513 let result = namespace
15516 .alter_transaction(AlterTransactionRequest {
15517 id: Some(vec!["users".to_string(), txn_id.clone()]),
15518 actions: vec![AlterTransactionAction {
15519 set_status_action: None,
15520 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15521 key: Some("custom_key".to_string()),
15522 value: Some("new_value".to_string()),
15523 mode: Some("Fail".to_string()),
15524 })),
15525 unset_property_action: None,
15526 }],
15527 ..Default::default()
15528 })
15529 .await;
15530 assert!(result.is_err());
15531 }
15532
15533 #[tokio::test]
15534 async fn test_alter_transaction_unset_property() {
15535 use lance_namespace::models::{
15536 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15537 AlterTransactionUnsetProperty,
15538 };
15539
15540 let (namespace, _temp_dir) = create_test_namespace().await;
15541 create_scalar_table(&namespace, "users").await;
15542 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15543 .await
15544 .expect("create_scalar_index should return a transaction id");
15545
15546 let response = namespace
15548 .alter_transaction(AlterTransactionRequest {
15549 id: Some(vec!["users".to_string(), txn_id.clone()]),
15550 actions: vec![
15551 AlterTransactionAction {
15552 set_status_action: None,
15553 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15554 key: Some("temp_key".to_string()),
15555 value: Some("temp_value".to_string()),
15556 mode: None,
15557 })),
15558 unset_property_action: None,
15559 },
15560 AlterTransactionAction {
15561 set_status_action: None,
15562 set_property_action: None,
15563 unset_property_action: Some(Box::new(AlterTransactionUnsetProperty {
15564 key: Some("temp_key".to_string()),
15565 mode: None,
15566 })),
15567 },
15568 ],
15569 ..Default::default()
15570 })
15571 .await
15572 .unwrap();
15573 assert_eq!(response.status, "SUCCEEDED");
15574 let props = response.properties.unwrap();
15575 assert!(!props.contains_key("temp_key"));
15576 }
15577
15578 #[tokio::test]
15579 async fn test_alter_transaction_invalid_status() {
15580 use lance_namespace::models::{
15581 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15582 };
15583
15584 let (namespace, _temp_dir) = create_test_namespace().await;
15585 create_scalar_table(&namespace, "users").await;
15586 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15587 .await
15588 .expect("create_scalar_index should return a transaction id");
15589
15590 let result = namespace
15591 .alter_transaction(AlterTransactionRequest {
15592 id: Some(vec!["users".to_string(), txn_id.clone()]),
15593 actions: vec![AlterTransactionAction {
15594 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15595 status: Some("InvalidStatus".to_string()),
15596 })),
15597 set_property_action: None,
15598 unset_property_action: None,
15599 }],
15600 ..Default::default()
15601 })
15602 .await;
15603 assert!(result.is_err());
15604 }
15605
15606 #[tokio::test]
15607 async fn test_alter_transaction_not_found() {
15608 use lance_namespace::models::{
15609 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15610 };
15611
15612 let (namespace, _temp_dir) = create_test_namespace().await;
15613 create_scalar_table(&namespace, "users").await;
15614
15615 let result = namespace
15617 .alter_transaction(AlterTransactionRequest {
15618 id: Some(vec!["users".to_string(), "non_existent_txn".to_string()]),
15619 actions: vec![AlterTransactionAction {
15620 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15621 status: Some("Canceled".to_string()),
15622 })),
15623 set_property_action: None,
15624 unset_property_action: None,
15625 }],
15626 ..Default::default()
15627 })
15628 .await;
15629 assert!(result.is_err());
15630 }
15631
15632 #[tokio::test]
15633 async fn test_alter_transaction_missing_id() {
15634 use lance_namespace::models::{
15635 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15636 };
15637
15638 let (namespace, _temp_dir) = create_test_namespace().await;
15639
15640 let result = namespace
15642 .alter_transaction(AlterTransactionRequest {
15643 id: None,
15644 actions: vec![AlterTransactionAction {
15645 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15646 status: Some("Canceled".to_string()),
15647 })),
15648 set_property_action: None,
15649 unset_property_action: None,
15650 }],
15651 ..Default::default()
15652 })
15653 .await;
15654 assert!(result.is_err());
15655
15656 let result = namespace
15658 .alter_transaction(AlterTransactionRequest {
15659 id: Some(vec!["users".to_string()]),
15660 actions: vec![AlterTransactionAction {
15661 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15662 status: Some("Canceled".to_string()),
15663 })),
15664 set_property_action: None,
15665 unset_property_action: None,
15666 }],
15667 ..Default::default()
15668 })
15669 .await;
15670 assert!(result.is_err());
15671 }
15672
15673 #[tokio::test]
15674 async fn test_alter_transaction_persists_changes() {
15675 use lance_namespace::models::{
15676 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15677 AlterTransactionSetStatus, DescribeTransactionRequest,
15678 };
15679
15680 let (namespace, _temp_dir) = create_test_namespace().await;
15681 create_scalar_table(&namespace, "users").await;
15682 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
15683
15684 let txn_id = transaction_id.expect("scalar index should produce a transaction id");
15685
15686 namespace
15688 .alter_transaction(AlterTransactionRequest {
15689 id: Some(vec!["users".to_string(), txn_id.clone()]),
15690 actions: vec![
15691 AlterTransactionAction {
15692 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15693 status: Some("Canceled".to_string()),
15694 })),
15695 set_property_action: None,
15696 unset_property_action: None,
15697 },
15698 AlterTransactionAction {
15699 set_status_action: None,
15700 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15701 key: Some("owner".to_string()),
15702 value: Some("alice".to_string()),
15703 mode: None,
15704 })),
15705 unset_property_action: None,
15706 },
15707 ],
15708 ..Default::default()
15709 })
15710 .await
15711 .unwrap();
15712
15713 let describe_resp = namespace
15716 .describe_transaction(DescribeTransactionRequest {
15717 id: Some(vec!["users".to_string(), txn_id.clone()]),
15718 ..Default::default()
15719 })
15720 .await
15721 .unwrap();
15722 let props = describe_resp.properties.expect("properties should be set");
15723 assert_eq!(props.get("owner"), Some(&"alice".to_string()));
15724 assert!(!props.contains_key("_status"));
15728
15729 let follow_up = namespace
15730 .alter_transaction(AlterTransactionRequest {
15731 id: Some(vec!["users".to_string(), txn_id.clone()]),
15732 actions: vec![],
15733 ..Default::default()
15734 })
15735 .await
15736 .unwrap();
15737 assert_eq!(follow_up.status, "Canceled");
15738 let follow_up_props = follow_up.properties.unwrap();
15739 assert_eq!(follow_up_props.get("owner"), Some(&"alice".to_string()));
15740 }
15741}