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::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
38use lance_linalg::distance::MetricType;
39use lance_table::io::commit::{ManifestNamingScheme, VERSIONS_DIR};
40use object_store::ObjectStoreExt;
41use object_store::path::Path;
42use object_store::{Error as ObjectStoreError, ObjectStore as OSObjectStore, PutMode, PutOptions};
43use std::collections::HashMap;
44use std::io::Cursor;
45use std::sync::{Arc, Mutex};
46use tokio::sync::OnceCell;
47
48use crate::context::DynamicContextProvider;
49use lance_namespace::models::{
50 AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest,
51 AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse,
52 AlterTransactionRequest, AlterTransactionResponse, AnalyzeTableQueryPlanRequest,
53 BatchDeleteTableVersionsRequest, BatchDeleteTableVersionsResponse,
54 BranchContents as ModelBranchContents, CountTableRowsRequest, CreateNamespaceRequest,
55 CreateNamespaceResponse, CreateTableBranchRequest, CreateTableBranchResponse,
56 CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest, CreateTableResponse,
57 CreateTableScalarIndexResponse, CreateTableTagRequest, CreateTableTagResponse,
58 CreateTableVersionRequest, CreateTableVersionResponse, DeclareTableRequest,
59 DeclareTableResponse, DeleteFromTableRequest, DeleteFromTableResponse,
60 DeleteTableBranchRequest, DeleteTableBranchResponse, DeleteTableTagRequest,
61 DeleteTableTagResponse, DescribeNamespaceRequest, DescribeNamespaceResponse,
62 DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, DescribeTableRequest,
63 DescribeTableResponse, DescribeTableVersionRequest, DescribeTableVersionResponse,
64 DescribeTransactionRequest, DescribeTransactionResponse, DropNamespaceRequest,
65 DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, DropTableRequest,
66 DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary,
67 GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest,
68 GetTableTagVersionResponse, Identity, IndexContent, InsertIntoTableRequest,
69 InsertIntoTableResponse, ListNamespacesRequest, ListNamespacesResponse,
70 ListTableBranchesRequest, ListTableBranchesResponse, ListTableIndicesRequest,
71 ListTableIndicesResponse, ListTableTagsRequest, ListTableTagsResponse,
72 ListTableVersionsRequest, ListTableVersionsResponse, ListTablesRequest, ListTablesResponse,
73 MergeInsertIntoTableRequest, MergeInsertIntoTableResponse, NamespaceExistsRequest,
74 QueryTableRequest, QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest,
75 RestoreTableResponse, TableExistsRequest, TableVersion, TagContents as ModelTagContents,
76 UpdateTableRequest, UpdateTableResponse, UpdateTableSchemaMetadataRequest,
77 UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, UpdateTableTagResponse,
78};
79
80use lance_core::{Error, Result, box_error};
81use lance_namespace::LanceNamespace;
82use lance_namespace::error::NamespaceError;
83use lance_namespace::schema::arrow_schema_to_json;
84
85use crate::credentials::{
86 CredentialVendor, create_credential_vendor_for_location, has_credential_vendor_config,
87};
88
89#[derive(Debug, Default)]
94pub struct OpsMetrics {
95 counters: Mutex<HashMap<String, u64>>,
96}
97
98impl OpsMetrics {
99 pub fn increment(&self, operation: &str) {
101 if let Ok(mut counters) = self.counters.lock() {
102 *counters.entry(operation.to_string()).or_insert(0) += 1;
103 }
104 }
105
106 pub fn retrieve(&self) -> HashMap<String, u64> {
108 self.counters.lock().map(|c| c.clone()).unwrap_or_default()
109 }
110
111 pub fn reset(&self) {
113 if let Ok(mut counters) = self.counters.lock() {
114 counters.clear();
115 }
116 }
117}
118
119pub(crate) fn build_sql_expressions(
122 new_columns: &[lance_namespace::models::AddColumnsEntry],
123) -> Result<Vec<(String, String)>> {
124 new_columns
125 .iter()
126 .map(|col| {
127 let expression = col.expression.clone().and_then(|opt| opt).ok_or_else(|| {
130 Error::invalid_input(format!(
131 "Expression is required for new column '{}'",
132 col.name
133 ))
134 })?;
135 Ok((col.name.clone(), expression))
136 })
137 .collect()
138}
139
140pub(crate) fn build_column_alterations(
143 alterations: &[lance_namespace::models::AlterColumnsEntry],
144) -> Result<Vec<lance::dataset::ColumnAlteration>> {
145 alterations
146 .iter()
147 .map(|entry| {
148 let mut alteration = lance::dataset::ColumnAlteration::new(entry.path.clone());
149 if let Some(Some(rename)) = &entry.rename {
151 alteration = alteration.rename(rename.clone());
152 }
153 if let Some(Some(nullable)) = entry.nullable {
155 alteration = alteration.set_nullable(nullable);
156 }
157 if let Some(data_type) = &entry.data_type
159 && !data_type.is_null()
160 {
161 let type_str = data_type.as_str().ok_or_else(|| {
162 Error::invalid_input(format!(
163 "data_type for column '{}' must be a JSON string, got: {}",
164 entry.path, data_type
165 ))
166 })?;
167 let json_type =
168 lance_namespace::models::JsonArrowDataType::new(type_str.to_string());
169 let dt =
170 lance_namespace::schema::convert_json_arrow_type(&json_type).map_err(|e| {
171 Error::invalid_input(format!(
172 "Failed to parse data_type '{}' for column '{}': {}",
173 type_str, entry.path, e
174 ))
175 })?;
176 alteration = alteration.cast_to(dt);
177 }
178 Ok(alteration)
179 })
180 .collect()
181}
182
183pub(crate) struct TableStatus {
188 pub(crate) exists: bool,
190 pub(crate) is_deregistered: bool,
192 pub(crate) has_reserved_file: bool,
194}
195
196enum DirectoryIndexParams {
197 Scalar {
198 index_type: IndexType,
199 params: ScalarIndexParams,
200 },
201 Inverted(InvertedIndexParams),
202 Vector {
203 index_type: IndexType,
204 params: VectorIndexParams,
205 },
206}
207
208impl DirectoryIndexParams {
209 fn index_type(&self) -> IndexType {
210 match self {
211 Self::Scalar { index_type, .. } | Self::Vector { index_type, .. } => *index_type,
212 Self::Inverted(_) => IndexType::Inverted,
213 }
214 }
215
216 fn params(&self) -> &dyn IndexParams {
217 match self {
218 Self::Scalar { params, .. } => params,
219 Self::Inverted(params) => params,
220 Self::Vector { params, .. } => params,
221 }
222 }
223}
224
225#[derive(Clone)]
260pub struct DirectoryNamespaceBuilder {
261 root: String,
262 storage_options: Option<HashMap<String, String>>,
263 session: Option<Arc<Session>>,
264 manifest_enabled: bool,
265 dir_listing_enabled: bool,
266 inline_optimization_enabled: bool,
267 table_version_tracking_enabled: bool,
268 dir_listing_to_manifest_migration_enabled: bool,
273 credential_vendor_properties: HashMap<String, String>,
274 context_provider: Option<Arc<dyn DynamicContextProvider>>,
275 commit_retries: Option<u32>,
276 vend_input_storage_options: bool,
279 vend_input_storage_options_refresh_interval_millis: Option<u64>,
284 ops_metrics_enabled: bool,
286}
287
288impl std::fmt::Debug for DirectoryNamespaceBuilder {
289 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
290 f.debug_struct("DirectoryNamespaceBuilder")
291 .field("root", &self.root)
292 .field("storage_options", &self.storage_options)
293 .field("manifest_enabled", &self.manifest_enabled)
294 .field("dir_listing_enabled", &self.dir_listing_enabled)
295 .field(
296 "inline_optimization_enabled",
297 &self.inline_optimization_enabled,
298 )
299 .field(
300 "table_version_tracking_enabled",
301 &self.table_version_tracking_enabled,
302 )
303 .field(
304 "dir_listing_to_manifest_migration_enabled",
305 &self.dir_listing_to_manifest_migration_enabled,
306 )
307 .field(
308 "context_provider",
309 &self.context_provider.as_ref().map(|_| "Some(...)"),
310 )
311 .field(
312 "vend_input_storage_options",
313 &self.vend_input_storage_options,
314 )
315 .field(
316 "vend_input_storage_options_refresh_interval_millis",
317 &self.vend_input_storage_options_refresh_interval_millis,
318 )
319 .field("ops_metrics_enabled", &self.ops_metrics_enabled)
320 .finish()
321 }
322}
323
324impl DirectoryNamespaceBuilder {
325 pub fn new(root: impl Into<String>) -> Self {
331 Self {
332 root: root.into().trim_end_matches('/').to_string(),
333 storage_options: None,
334 session: None,
335 manifest_enabled: true,
336 dir_listing_enabled: true, inline_optimization_enabled: true,
338 table_version_tracking_enabled: false, dir_listing_to_manifest_migration_enabled: false, credential_vendor_properties: HashMap::new(),
341 context_provider: None,
342 commit_retries: None,
343 vend_input_storage_options: false,
344 vend_input_storage_options_refresh_interval_millis: None,
345 ops_metrics_enabled: false,
346 }
347 }
348
349 pub fn manifest_enabled(mut self, enabled: bool) -> Self {
354 self.manifest_enabled = enabled;
355 self
356 }
357
358 pub fn dir_listing_enabled(mut self, enabled: bool) -> Self {
363 self.dir_listing_enabled = enabled;
364 self
365 }
366
367 pub fn dir_listing_to_manifest_migration_enabled(mut self, enabled: bool) -> Self {
374 self.dir_listing_to_manifest_migration_enabled = enabled;
375 self
376 }
377
378 pub fn inline_optimization_enabled(mut self, enabled: bool) -> Self {
383 self.inline_optimization_enabled = enabled;
384 self
385 }
386
387 pub fn table_version_tracking_enabled(mut self, enabled: bool) -> Self {
395 self.table_version_tracking_enabled = enabled;
396 self
397 }
398
399 pub fn from_properties(
467 properties: HashMap<String, String>,
468 session: Option<Arc<Session>>,
469 ) -> Result<Self> {
470 let root = properties.get("root").cloned().ok_or_else(|| {
472 lance_core::Error::from(NamespaceError::InvalidInput {
473 message: "Missing required property 'root' for directory namespace".to_string(),
474 })
475 })?;
476
477 let storage_options: HashMap<String, String> = properties
479 .iter()
480 .filter_map(|(k, v)| {
481 k.strip_prefix("storage.")
482 .map(|key| (key.to_string(), v.clone()))
483 })
484 .collect();
485
486 let storage_options = if storage_options.is_empty() {
487 None
488 } else {
489 Some(storage_options)
490 };
491
492 let manifest_enabled = properties
494 .get("manifest_enabled")
495 .and_then(|v| v.parse::<bool>().ok())
496 .unwrap_or(true);
497
498 let dir_listing_enabled = properties
500 .get("dir_listing_enabled")
501 .and_then(|v| v.parse::<bool>().ok())
502 .unwrap_or(true);
503
504 let inline_optimization_enabled = properties
506 .get("inline_optimization_enabled")
507 .and_then(|v| v.parse::<bool>().ok())
508 .unwrap_or(true);
509
510 let table_version_tracking_enabled = properties
512 .get("table_version_tracking_enabled")
513 .and_then(|v| v.parse::<bool>().ok())
514 .unwrap_or(false);
515
516 let dir_listing_to_manifest_migration_enabled = properties
518 .get("dir_listing_to_manifest_migration_enabled")
519 .and_then(|v| v.parse::<bool>().ok())
520 .unwrap_or(false);
521
522 let credential_vendor_properties: HashMap<String, String> = properties
526 .iter()
527 .filter_map(|(k, v)| {
528 k.strip_prefix("credential_vendor.")
529 .map(|key| (key.to_string(), v.clone()))
530 })
531 .collect();
532
533 let commit_retries = properties
534 .get("commit_retries")
535 .and_then(|v| v.parse::<u32>().ok());
536
537 let vend_input_storage_options = properties
539 .get("vend_input_storage_options")
540 .and_then(|v| v.parse::<bool>().ok())
541 .unwrap_or(false);
542
543 let vend_input_storage_options_refresh_interval_millis = properties
545 .get("vend_input_storage_options_refresh_interval_millis")
546 .and_then(|v| v.parse::<u64>().ok());
547
548 let ops_metrics_enabled = properties
550 .get("ops_metrics_enabled")
551 .and_then(|v| v.parse::<bool>().ok())
552 .unwrap_or(false);
553
554 Ok(Self {
555 root: root.trim_end_matches('/').to_string(),
556 storage_options,
557 session,
558 manifest_enabled,
559 dir_listing_enabled,
560 inline_optimization_enabled,
561 table_version_tracking_enabled,
562 dir_listing_to_manifest_migration_enabled,
563 credential_vendor_properties,
564 context_provider: None,
565 commit_retries,
566 vend_input_storage_options,
567 vend_input_storage_options_refresh_interval_millis,
568 ops_metrics_enabled,
569 })
570 }
571
572 pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
579 self.storage_options
580 .get_or_insert_with(HashMap::new)
581 .insert(key.into(), value.into());
582 self
583 }
584
585 pub fn storage_options(mut self, options: HashMap<String, String>) -> Self {
591 self.storage_options
592 .get_or_insert_with(HashMap::new)
593 .extend(options);
594 self
595 }
596
597 pub fn session(mut self, session: Arc<Session>) -> Self {
607 self.session = Some(session);
608 self
609 }
610
611 pub fn commit_retries(mut self, retries: u32) -> Self {
614 self.commit_retries = Some(retries);
615 self
616 }
617
618 pub fn credential_vendor_property(
646 mut self,
647 key: impl Into<String>,
648 value: impl Into<String>,
649 ) -> Self {
650 self.credential_vendor_properties
651 .insert(key.into(), value.into());
652 self
653 }
654
655 pub fn credential_vendor_properties(mut self, properties: HashMap<String, String>) -> Self {
663 self.credential_vendor_properties.extend(properties);
664 self
665 }
666
667 pub fn context_provider(mut self, provider: Arc<dyn DynamicContextProvider>) -> Self {
677 self.context_provider = Some(provider);
678 self
679 }
680
681 pub fn vend_input_storage_options(mut self, enabled: bool) -> Self {
690 self.vend_input_storage_options = enabled;
691 self
692 }
693
694 pub fn vend_input_storage_options_refresh_interval_millis(
706 mut self,
707 interval_millis: u64,
708 ) -> Self {
709 self.vend_input_storage_options_refresh_interval_millis = Some(interval_millis);
710 self
711 }
712
713 pub fn ops_metrics_enabled(mut self, enabled: bool) -> Self {
721 self.ops_metrics_enabled = enabled;
722 self
723 }
724
725 pub async fn build(self) -> Result<DirectoryNamespace> {
738 let (object_store, base_path) =
739 Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?;
740
741 let manifest_ns = if self.manifest_enabled {
742 match manifest::ManifestNamespace::open_from_directory(
743 self.root.clone(),
744 self.storage_options.clone(),
745 self.session.clone(),
746 object_store.clone(),
747 base_path.clone(),
748 self.dir_listing_enabled,
749 self.inline_optimization_enabled,
750 self.commit_retries,
751 )
752 .await
753 {
754 Ok(ns) => Some(Arc::new(ns)),
755 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
756 return Err(e);
760 }
761 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => {
762 log::debug!("Manifest namespace does not exist yet: {}", e);
763 None
764 }
765 Err(e) => return Err(e),
766 }
767 } else {
768 None
769 };
770 let manifest_cell = OnceCell::new();
771 if let Some(manifest_ns) = manifest_ns {
772 let _ = manifest_cell.set(manifest_ns);
773 }
774
775 let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties)
777 {
778 create_credential_vendor_for_location(&self.root, &self.credential_vendor_properties)
779 .await?
780 .map(Arc::from)
781 } else {
782 None
783 };
784
785 let ops_metrics = if self.ops_metrics_enabled {
786 Some(Arc::new(OpsMetrics::default()))
787 } else {
788 None
789 };
790
791 Ok(DirectoryNamespace {
792 root: self.root,
793 storage_options: self.storage_options,
794 session: self.session,
795 object_store,
796 base_path,
797 manifest_ns: manifest_cell,
798 write_manifest_ns: OnceCell::new(),
799 manifest_enabled: self.manifest_enabled,
800 dir_listing_enabled: self.dir_listing_enabled,
801 inline_optimization_enabled: self.inline_optimization_enabled,
802 commit_retries: self.commit_retries,
803 dir_listing_to_manifest_migration_enabled: self
804 .dir_listing_to_manifest_migration_enabled,
805 table_version_tracking_enabled: self.table_version_tracking_enabled,
806 credential_vendor,
807 context_provider: self.context_provider,
808 vend_input_storage_options: self.vend_input_storage_options,
809 vend_input_storage_options_refresh_interval_millis: self
810 .vend_input_storage_options_refresh_interval_millis,
811 ops_metrics,
812 })
813 }
814
815 async fn initialize_object_store(
817 root: &str,
818 storage_options: &Option<HashMap<String, String>>,
819 session: &Option<Arc<Session>>,
820 ) -> Result<(Arc<ObjectStore>, Path)> {
821 let accessor = storage_options.clone().map(|opts| {
823 Arc::new(lance_io::object_store::StorageOptionsAccessor::with_static_options(opts))
824 });
825 let params = ObjectStoreParams {
826 storage_options_accessor: accessor,
827 ..Default::default()
828 };
829
830 let registry = if let Some(session) = session {
832 session.store_registry()
833 } else {
834 Arc::new(ObjectStoreRegistry::default())
835 };
836
837 let (object_store, base_path) = ObjectStore::from_uri_and_params(registry, root, ¶ms)
839 .await
840 .map_err(|e| {
841 lance_core::Error::from(NamespaceError::Internal {
842 message: format!("Failed to create object store: {:?}", e),
843 })
844 })?;
845
846 Ok((object_store, base_path))
847 }
848}
849
850pub struct DirectoryNamespace {
874 root: String,
875 storage_options: Option<HashMap<String, String>>,
876 session: Option<Arc<Session>>,
877 object_store: Arc<ObjectStore>,
878 base_path: Path,
879 manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
880 write_manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
881 manifest_enabled: bool,
882 dir_listing_enabled: bool,
883 inline_optimization_enabled: bool,
884 commit_retries: Option<u32>,
885 dir_listing_to_manifest_migration_enabled: bool,
889 table_version_tracking_enabled: bool,
892 credential_vendor: Option<Arc<dyn CredentialVendor>>,
895 #[allow(dead_code)]
898 context_provider: Option<Arc<dyn DynamicContextProvider>>,
899 vend_input_storage_options: bool,
901 vend_input_storage_options_refresh_interval_millis: Option<u64>,
904 ops_metrics: Option<Arc<OpsMetrics>>,
906}
907
908impl std::fmt::Debug for DirectoryNamespace {
909 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
910 write!(f, "{}", self.namespace_id())
911 }
912}
913
914impl std::fmt::Display for DirectoryNamespace {
915 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
916 write!(f, "{}", self.namespace_id())
917 }
918}
919
920struct TableDeleteEntry {
923 table_id: Option<Vec<String>>,
924 ranges: Vec<(i64, i64)>,
925}
926
927#[derive(Debug, Clone, Default)]
937struct TransactionAlteration {
938 status: Option<String>,
940 properties: HashMap<String, String>,
944 removed_properties: std::collections::HashSet<String>,
948}
949
950impl TransactionAlteration {
951 const F_STATUS: &'static str = "status";
953 const F_PROPERTIES: &'static str = "properties";
954 const F_REMOVED_PROPERTIES: &'static str = "removed_properties";
955
956 fn to_json_bytes(&self) -> serde_json::Result<Vec<u8>> {
963 serde_json::to_vec(&serde_json::json!({
964 Self::F_STATUS: self.status,
965 Self::F_PROPERTIES: self.properties,
966 Self::F_REMOVED_PROPERTIES: self.removed_properties,
967 }))
968 }
969
970 fn from_json_slice(bytes: &[u8]) -> serde_json::Result<Self> {
975 let mut obj: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes)?;
976 Ok(Self {
977 status: serde_json::from_value(
978 obj.remove(Self::F_STATUS)
979 .unwrap_or(serde_json::Value::Null),
980 )?,
981 properties: serde_json::from_value(
982 obj.remove(Self::F_PROPERTIES)
983 .unwrap_or(serde_json::Value::Null),
984 )
985 .unwrap_or_default(),
986 removed_properties: serde_json::from_value(
987 obj.remove(Self::F_REMOVED_PROPERTIES)
988 .unwrap_or(serde_json::Value::Null),
989 )
990 .unwrap_or_default(),
991 })
992 }
993}
994
995impl DirectoryNamespace {
996 fn manifest_ns_for_read(&self) -> Option<&Arc<manifest::ManifestNamespace>> {
997 self.write_manifest_ns
998 .get()
999 .or_else(|| self.manifest_ns.get())
1000 }
1001
1002 async fn manifest_ns_for_write(&self) -> Result<Option<Arc<manifest::ManifestNamespace>>> {
1003 if !self.manifest_enabled {
1004 return Ok(None);
1005 }
1006
1007 let manifest_ns = self
1008 .write_manifest_ns
1009 .get_or_try_init(|| async {
1010 manifest::ManifestNamespace::from_directory(
1011 self.root.clone(),
1012 self.storage_options.clone(),
1013 self.session.clone(),
1014 self.object_store.clone(),
1015 self.base_path.clone(),
1016 self.dir_listing_enabled,
1017 self.inline_optimization_enabled,
1018 self.commit_retries,
1019 )
1020 .await
1021 .map(Arc::new)
1022 })
1023 .await?;
1024 Ok(Some(manifest_ns.clone()))
1025 }
1026
1027 fn child_namespace_requires_manifest_error(&self) -> Error {
1028 if self.manifest_enabled {
1029 NamespaceError::NamespaceNotFound {
1030 message: "Child namespace reads require an existing __manifest dataset".to_string(),
1031 }
1032 .into()
1033 } else {
1034 NamespaceError::Unsupported {
1035 message: "Child namespaces are only supported when manifest mode is enabled"
1036 .to_string(),
1037 }
1038 .into()
1039 }
1040 }
1041
1042 fn apply_pagination(
1055 names: &mut Vec<String>,
1056 page_token: Option<String>,
1057 limit: Option<i32>,
1058 ) -> Option<String> {
1059 names.sort();
1061
1062 if let Some(start_after) = page_token {
1064 if let Some(index) = names
1065 .iter()
1066 .position(|name| name.as_str() > start_after.as_str())
1067 {
1068 names.drain(0..index);
1069 } else {
1070 names.clear();
1071 }
1072 }
1073
1074 if let Some(limit) = limit
1076 && limit >= 0
1077 {
1078 let limit = limit as usize;
1079 if names.len() > limit {
1080 let next_page_token = if limit > 0 {
1081 Some(names[limit - 1].clone())
1082 } else {
1083 None
1084 };
1085 names.truncate(limit);
1086 return next_page_token;
1087 }
1088 }
1089
1090 None
1091 }
1092
1093 async fn list_directory_tables(&self) -> Result<Vec<String>> {
1095 let mut tables = Vec::new();
1096 let entries = self
1097 .object_store
1098 .read_dir(self.base_path.clone())
1099 .await
1100 .map_err(|e| {
1101 lance_core::Error::from(NamespaceError::Internal {
1102 message: format!("Failed to list directory: {:?}", e),
1103 })
1104 })?;
1105
1106 for entry in entries {
1107 let path = entry.trim_end_matches('/');
1108 if !path.ends_with(".lance") {
1109 continue;
1110 }
1111
1112 let table_name = &path[..path.len() - 6];
1113
1114 let status = self.check_table_status(table_name).await;
1116 if status.is_deregistered {
1117 continue;
1118 }
1119
1120 tables.push(table_name.to_string());
1121 }
1122
1123 Ok(tables)
1124 }
1125
1126 fn validate_root_namespace_id(id: &Option<Vec<String>>) -> Result<()> {
1128 if let Some(id) = id
1129 && !id.is_empty()
1130 {
1131 return Err(NamespaceError::Unsupported {
1132 message: format!(
1133 "Directory namespace only supports root namespace operations, but got namespace ID: {:?}. Expected empty ID.",
1134 id
1135 ),
1136 }
1137 .into());
1138 }
1139 Ok(())
1140 }
1141
1142 fn table_name_from_id(id: &Option<Vec<String>>) -> Result<String> {
1144 let id = id.as_ref().ok_or_else(|| {
1145 lance_core::Error::from(NamespaceError::InvalidInput {
1146 message: "Directory namespace table ID cannot be empty".to_string(),
1147 })
1148 })?;
1149
1150 if id.len() != 1 {
1151 return Err(NamespaceError::Unsupported {
1152 message: format!(
1153 "Multi-level table IDs are only supported when manifest mode is enabled, but got: {:?}",
1154 id
1155 ),
1156 }
1157 .into());
1158 }
1159
1160 Ok(id[0].clone())
1161 }
1162
1163 fn format_table_id(table_id: &[String]) -> String {
1164 format!(
1165 "table id '{}'",
1166 manifest::ManifestNamespace::str_object_id(table_id)
1167 )
1168 }
1169
1170 fn format_table_id_from_request(id: &Option<Vec<String>>) -> String {
1171 id.as_ref()
1172 .map(|table_id| Self::format_table_id(table_id))
1173 .unwrap_or_else(|| "table id '<unknown>'".to_string())
1174 }
1175
1176 async fn resolve_table_location(&self, id: &Option<Vec<String>>) -> Result<String> {
1177 let mut describe_req = DescribeTableRequest::new();
1178 describe_req.id = id.clone();
1179 describe_req.load_detailed_metadata = Some(false);
1180
1181 let describe_resp = self.describe_table_impl(describe_req).await?;
1183
1184 describe_resp.location.ok_or_else(|| {
1185 lance_core::Error::from(NamespaceError::TableNotFound {
1186 message: format!("Table location not found for: {:?}", id),
1187 })
1188 })
1189 }
1190
1191 fn map_tag_error(err: lance_core::Error, tag: &str, table_uri: &str) -> lance_core::Error {
1194 match err {
1195 lance_core::Error::RefNotFound { .. } => NamespaceError::TableTagNotFound {
1196 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1197 }
1198 .into(),
1199 lance_core::Error::RefConflict { .. } => NamespaceError::TableTagAlreadyExists {
1200 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1201 }
1202 .into(),
1203 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1204 message: format!("invalid tag '{}': {}", tag, message),
1205 }
1206 .into(),
1207 lance_core::Error::VersionNotFound { message } => {
1208 NamespaceError::TableVersionNotFound {
1209 message: format!(
1210 "version referenced by tag '{}' not found for table at '{}': {}",
1211 tag, table_uri, message
1212 ),
1213 }
1214 .into()
1215 }
1216 other => NamespaceError::Internal {
1217 message: format!(
1218 "tag operation failed for tag '{}' on table at '{}': {}",
1219 tag, table_uri, other
1220 ),
1221 }
1222 .into(),
1223 }
1224 }
1225
1226 fn map_branch_error(
1232 err: lance_core::Error,
1233 branch: &str,
1234 table_uri: &str,
1235 ) -> lance_core::Error {
1236 match err {
1237 lance_core::Error::RefNotFound { .. } => NamespaceError::TableBranchNotFound {
1238 message: format!("branch '{}' for table at '{}'", branch, table_uri),
1239 }
1240 .into(),
1241 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1242 message: format!("invalid branch '{}': {}", branch, message),
1243 }
1244 .into(),
1245 lance_core::Error::VersionNotFound { message } => {
1246 NamespaceError::TableVersionNotFound {
1247 message: format!(
1248 "source version for branch '{}' not found for table at '{}': {}",
1249 branch, table_uri, message
1250 ),
1251 }
1252 .into()
1253 }
1254 other => NamespaceError::Internal {
1255 message: format!(
1256 "branch operation failed for branch '{}' on table at '{}': {}",
1257 branch, table_uri, other
1258 ),
1259 }
1260 .into(),
1261 }
1262 }
1263
1264 fn map_mutation_error(
1276 err: lance_core::Error,
1277 operation: &str,
1278 table_uri: &str,
1279 ) -> lance_core::Error {
1280 let detail = err.to_string();
1281 let ns_err = match &err {
1282 lance_core::Error::InvalidInput { .. }
1283 | lance_core::Error::Unprocessable { .. }
1284 | lance_core::Error::InvalidRef { .. } => NamespaceError::InvalidInput {
1285 message: format!(
1286 "Invalid input for {} on table at '{}': {}",
1287 operation, table_uri, detail
1288 ),
1289 },
1290 lance_core::Error::NotFound { .. } | lance_core::Error::DatasetNotFound { .. } => {
1291 NamespaceError::TableNotFound {
1292 message: format!(
1293 "Table at '{}' not found while running {}: {}",
1294 table_uri, operation, detail
1295 ),
1296 }
1297 }
1298 lance_core::Error::SchemaMismatch { .. } | lance_core::Error::Schema { .. } => {
1299 NamespaceError::TableSchemaValidationError {
1300 message: format!(
1301 "Schema validation failed for {} on table at '{}': {}",
1302 operation, table_uri, detail
1303 ),
1304 }
1305 }
1306 lance_core::Error::CommitConflict { .. } => NamespaceError::Throttling {
1310 message: format!(
1311 "Too many concurrent writes for {} on table at '{}', please retry later: {}",
1312 operation, table_uri, detail
1313 ),
1314 },
1315 lance_core::Error::TooMuchWriteContention { .. }
1319 | lance_core::Error::RetryableCommitConflict { .. }
1320 | lance_core::Error::IncompatibleTransaction { .. }
1321 | lance_core::Error::VersionConflict { .. } => NamespaceError::ConcurrentModification {
1322 message: format!(
1323 "Concurrent modification detected for {} on table at '{}': {}",
1324 operation, table_uri, detail
1325 ),
1326 },
1327 lance_core::Error::NotSupported { .. } => NamespaceError::Unsupported {
1328 message: format!(
1329 "{} is not supported on table at '{}': {}",
1330 operation, table_uri, detail
1331 ),
1332 },
1333 _ => NamespaceError::Internal {
1334 message: format!(
1335 "Failed to run {} on table at '{}': {}",
1336 operation, table_uri, detail
1337 ),
1338 },
1339 };
1340 ns_err.into()
1341 }
1342
1343 async fn table_has_actual_manifests(&self, table_name: &str) -> Result<bool> {
1344 manifest::ManifestNamespace::path_has_actual_manifests(
1345 &self.object_store,
1346 &self.table_path(table_name),
1347 )
1348 .await
1349 }
1350
1351 async fn filter_declared_tables(
1352 &self,
1353 tables: Vec<String>,
1354 include_declared: bool,
1355 ) -> Result<Vec<String>> {
1356 if include_declared {
1357 return Ok(tables);
1358 }
1359
1360 let mut stream = futures::stream::iter(tables.into_iter().map(|table_name| async move {
1361 if self.table_has_actual_manifests(&table_name).await? {
1365 Ok::<Option<String>, Error>(Some(table_name))
1366 } else {
1367 Ok::<Option<String>, Error>(None)
1368 }
1369 }))
1370 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1371
1372 let mut filtered = Vec::new();
1373 while let Some(result) = stream.next().await {
1374 if let Some(table_name) = result? {
1375 filtered.push(table_name);
1376 }
1377 }
1378 Ok(filtered)
1379 }
1380
1381 fn ipc_reader_from_request_data(
1382 request_data: &Bytes,
1383 operation: &str,
1384 ) -> Result<(
1385 Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1386 usize,
1387 )> {
1388 if request_data.is_empty() {
1389 return Err(NamespaceError::InvalidInput {
1390 message: format!(
1391 "Request data (Arrow IPC stream) is required for {}",
1392 operation
1393 ),
1394 }
1395 .into());
1396 }
1397
1398 let cursor = Cursor::new(request_data.as_ref());
1399 let stream_reader =
1400 StreamReader::try_new(cursor, None).map_err(|e| NamespaceError::InvalidInput {
1401 message: format!("Invalid Arrow IPC stream: {}", e),
1402 })?;
1403 let arrow_schema = stream_reader.schema();
1404
1405 let mut num_rows = 0usize;
1406 let mut batches = Vec::new();
1407 for batch_result in stream_reader {
1408 let batch = batch_result.map_err(|e| NamespaceError::Internal {
1409 message: format!("Failed to read batch from IPC stream: {}", e),
1410 })?;
1411 num_rows += batch.num_rows();
1412 batches.push(batch);
1413 }
1414
1415 let reader: Box<dyn arrow::record_batch::RecordBatchReader + Send> = if batches.is_empty() {
1416 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
1417 Box::new(RecordBatchIterator::new(vec![Ok(batch)], arrow_schema))
1418 } else {
1419 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
1420 Box::new(RecordBatchIterator::new(batch_results, arrow_schema))
1421 };
1422
1423 Ok((reader, num_rows))
1424 }
1425
1426 async fn table_uri_has_actual_manifests(&self, table_uri: &str) -> Result<bool> {
1427 let table_path = self.object_store_path_from_uri(table_uri)?;
1428 manifest::ManifestNamespace::path_has_actual_manifests(&self.object_store, &table_path)
1429 .await
1430 }
1431
1432 fn object_store_path_from_uri(&self, uri: &str) -> Result<Path> {
1433 let registry = self
1434 .session
1435 .as_ref()
1436 .map(|session| session.store_registry())
1437 .unwrap_or_else(|| Arc::new(ObjectStoreRegistry::default()));
1438 ObjectStore::extract_path_from_uri(registry, uri)
1439 }
1440
1441 fn normalized_branch(branch: Option<&str>) -> Result<Option<&str>> {
1446 match branch.filter(|b| !b.is_empty() && *b != "main") {
1447 Some(branch) => {
1448 check_valid_branch(branch).map_err(|e| {
1449 lance_core::Error::from(NamespaceError::InvalidInput {
1450 message: format!("invalid branch name '{}': {}", branch, e),
1451 })
1452 })?;
1453 Ok(Some(branch))
1454 }
1455 None => Ok(None),
1456 }
1457 }
1458
1459 async fn open_validated_branch(&self, table_uri: &str, branch: &str) -> Result<Dataset> {
1460 let dataset = self
1461 .configured_builder(table_uri)
1462 .with_branch(branch, None)
1463 .load()
1464 .await
1465 .map_err(|e| {
1466 lance_core::Error::from(NamespaceError::TableNotFound {
1467 message: format!(
1468 "branch '{}' not found for table at '{}': {}",
1469 branch, table_uri, e
1470 ),
1471 })
1472 })?;
1473 dataset.branches().get(branch).await.map_err(|_| {
1474 lance_core::Error::from(NamespaceError::TableNotFound {
1475 message: format!("branch '{}' not found for table at '{}'", branch, table_uri),
1476 })
1477 })?;
1478 Ok(dataset)
1479 }
1480
1481 async fn resolve_branch_location(&self, table_uri: &str, branch: &str) -> Result<String> {
1482 Ok(self
1483 .open_validated_branch(table_uri, branch)
1484 .await?
1485 .branch_location()
1486 .uri)
1487 }
1488
1489 async fn resolve_branch_for_commit(
1496 &self,
1497 table_uri: &str,
1498 branch: &str,
1499 ) -> Result<(String, Path)> {
1500 let main = self
1501 .configured_builder(table_uri)
1502 .load()
1503 .await
1504 .map_err(|e| {
1505 lance_core::Error::from(NamespaceError::TableNotFound {
1506 message: format!("table at '{}' not found: {}", table_uri, e),
1507 })
1508 })?;
1509 let branch_location = main.branch_location().find_branch(Some(branch))?;
1510 match main.branches().get(branch).await {
1511 Ok(_) => Ok((branch_location.uri, branch_location.path)),
1512 Err(lance_core::Error::RefNotFound { .. }) => {
1513 if self
1514 .branch_has_committed_versions(&branch_location.path)
1515 .await?
1516 {
1517 return Err(NamespaceError::TableNotFound {
1518 message: format!(
1519 "branch '{}' not found for table at '{}'",
1520 branch, table_uri
1521 ),
1522 }
1523 .into());
1524 }
1525 Ok((branch_location.uri, branch_location.path))
1526 }
1527 Err(e) => Err(e),
1528 }
1529 }
1530
1531 async fn branch_has_committed_versions(&self, branch_path: &Path) -> Result<bool> {
1532 Ok(!self
1533 .list_versions_under(branch_path, false, Some(1))
1534 .await?
1535 .is_empty())
1536 }
1537
1538 fn validate_dir_only_properties(
1539 properties: Option<&HashMap<String, String>>,
1540 operation: &str,
1541 ) -> Result<()> {
1542 if properties.is_some_and(|properties| !properties.is_empty()) {
1546 return Err(NamespaceError::Unsupported {
1547 message: format!(
1548 "{} with non-empty table properties requires manifest_enabled=true",
1549 operation
1550 ),
1551 }
1552 .into());
1553 }
1554 Ok(())
1555 }
1556
1557 async fn write_reader_to_table(
1558 &self,
1559 table_uri: &str,
1560 reader: Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1561 mode: WriteMode,
1562 extra_storage_options: Option<HashMap<String, String>>,
1563 ) -> Result<Dataset> {
1564 let mut merged_storage_options = self.storage_options.clone().unwrap_or_default();
1567 if let Some(extra_storage_options) = extra_storage_options {
1568 merged_storage_options.extend(extra_storage_options);
1569 }
1570 let store_params = (!merged_storage_options.is_empty()).then(|| ObjectStoreParams {
1571 storage_options_accessor: Some(Arc::new(
1572 lance_io::object_store::StorageOptionsAccessor::with_static_options(
1573 merged_storage_options,
1574 ),
1575 )),
1576 ..Default::default()
1577 });
1578
1579 let write_params = WriteParams {
1580 mode,
1581 store_params,
1582 session: self.session.clone(),
1583 ..Default::default()
1584 };
1585
1586 let dataset = Dataset::write(reader, table_uri, Some(write_params))
1587 .await
1588 .map_err(|e| NamespaceError::Internal {
1589 message: format!("Failed to write table at '{}': {}", table_uri, e),
1590 })?;
1591
1592 Ok(dataset)
1593 }
1594
1595 fn manifest_version_from_filename(filename: &str) -> Option<u64> {
1599 ManifestNamingScheme::detect_scheme(filename)?.parse_version(filename)
1600 }
1601
1602 async fn list_table_versions_from_storage(
1603 &self,
1604 table_uri: &str,
1605 descending: bool,
1606 limit: Option<i32>,
1607 ) -> Result<Vec<TableVersion>> {
1608 let table_path = self.object_store_path_from_uri(table_uri)?;
1609 self.list_versions_under(&table_path, descending, limit)
1610 .await
1611 }
1612
1613 async fn list_versions_under(
1617 &self,
1618 table_path: &Path,
1619 descending: bool,
1620 limit: Option<i32>,
1621 ) -> Result<Vec<TableVersion>> {
1622 let versions_dir = table_path.clone().join(VERSIONS_DIR);
1623 let manifest_metas: Vec<_> = self
1624 .object_store
1625 .read_dir_all(&versions_dir, None)
1626 .try_collect()
1627 .await
1628 .map_err(|e| {
1629 lance_core::Error::from(NamespaceError::Internal {
1630 message: format!(
1631 "Failed to list manifest files under '{}': {}",
1632 versions_dir, e
1633 ),
1634 })
1635 })?;
1636
1637 let is_v2_naming = manifest_metas
1638 .first()
1639 .is_some_and(|meta| meta.location.filename().is_some_and(|f| f.len() == 29));
1640
1641 let mut table_versions: Vec<TableVersion> = manifest_metas
1642 .into_iter()
1643 .filter_map(|meta| {
1644 let filename = meta.location.filename()?;
1645 let actual_version = Self::manifest_version_from_filename(filename)?;
1646
1647 Some(TableVersion {
1648 version: actual_version as i64,
1649 manifest_path: meta.location.to_string(),
1650 manifest_size: Some(meta.size as i64),
1651 e_tag: meta.e_tag,
1652 timestamp_millis: Some(meta.last_modified.timestamp_millis()),
1653 metadata: None,
1654 })
1655 })
1656 .collect();
1657
1658 let list_is_ordered = self.object_store.list_is_lexically_ordered;
1659
1660 let needs_sort = if list_is_ordered {
1661 if is_v2_naming {
1662 !descending
1663 } else {
1664 descending
1665 }
1666 } else {
1667 true
1668 };
1669
1670 if needs_sort {
1671 if descending {
1672 table_versions.sort_by_key(|v| std::cmp::Reverse(v.version));
1673 } else {
1674 table_versions.sort_by_key(|v| v.version);
1675 }
1676 }
1677
1678 if let Some(limit) = limit {
1679 table_versions.truncate(limit as usize);
1680 }
1681
1682 Ok(table_versions)
1683 }
1684
1685 async fn describe_table_impl(
1689 &self,
1690 request: DescribeTableRequest,
1691 ) -> Result<DescribeTableResponse> {
1692 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
1693 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
1694 let skip_manifest_for_root = self.dir_listing_enabled
1695 && is_root_level
1696 && !self.dir_listing_to_manifest_migration_enabled;
1697 if let Some(manifest_ns) = self.manifest_ns_for_read()
1698 && !skip_manifest_for_root
1699 {
1700 match manifest_ns.describe_table(request.clone()).await {
1701 Ok(mut response) => {
1702 if let Some(ref table_uri) = response.table_uri {
1703 let vend = request.vend_credentials.unwrap_or(true);
1705 let identity = request.identity.as_deref();
1706 response.storage_options = self
1707 .get_storage_options_for_table(table_uri, vend, identity)
1708 .await?;
1709 }
1710 if self.table_version_tracking_enabled {
1712 response.managed_versioning = Some(true);
1713 }
1714 return Ok(response);
1715 }
1716 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
1717 return Err(e);
1720 }
1721 Err(_) if self.dir_listing_enabled && is_root_level => {
1722 }
1724 Err(e) => return Err(e),
1725 }
1726 }
1727 if is_child_table {
1728 return Err(self.child_namespace_requires_manifest_error());
1729 }
1730
1731 let table_name = Self::table_name_from_id(&request.id)?;
1732 let table_id = Self::format_table_id_from_request(&request.id);
1733 if !self.dir_listing_enabled {
1734 return Err(NamespaceError::TableNotFound { message: table_id }.into());
1735 }
1736
1737 let table_uri = self.table_full_uri(&table_name);
1738
1739 let status = self.check_table_status(&table_name).await;
1741
1742 if !status.exists {
1743 return Err(NamespaceError::TableNotFound {
1744 message: table_id.clone(),
1745 }
1746 .into());
1747 }
1748
1749 if status.is_deregistered {
1750 return Err(NamespaceError::TableNotFound {
1751 message: format!("Table is deregistered: {}", table_id),
1752 }
1753 .into());
1754 }
1755
1756 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
1757 let should_check_declared =
1758 load_detailed_metadata || request.check_declared.unwrap_or(false);
1759 let vend_credentials = request.vend_credentials.unwrap_or(true);
1761 let identity = request.identity.as_deref();
1762 let is_only_declared = if should_check_declared {
1763 if status.has_reserved_file {
1764 Some(!self.table_has_actual_manifests(&table_name).await?)
1765 } else {
1766 Some(false)
1767 }
1768 } else {
1769 None
1770 };
1771
1772 if !load_detailed_metadata {
1773 let storage_options = self
1774 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1775 .await?;
1776 return Ok(DescribeTableResponse {
1777 table: Some(table_name),
1778 namespace: request.id.as_ref().map(|id| {
1779 if id.len() > 1 {
1780 id[..id.len() - 1].to_vec()
1781 } else {
1782 vec![]
1783 }
1784 }),
1785 location: Some(table_uri.clone()),
1786 table_uri: Some(table_uri),
1787 storage_options,
1788 is_only_declared,
1789 managed_versioning: if self.table_version_tracking_enabled {
1790 Some(true)
1791 } else {
1792 None
1793 },
1794 ..Default::default()
1795 });
1796 }
1797
1798 if is_only_declared == Some(true) {
1799 let storage_options = self
1800 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1801 .await?;
1802 return Ok(DescribeTableResponse {
1803 table: Some(table_name),
1804 namespace: request.id.as_ref().map(|id| {
1805 if id.len() > 1 {
1806 id[..id.len() - 1].to_vec()
1807 } else {
1808 vec![]
1809 }
1810 }),
1811 location: Some(table_uri.clone()),
1812 table_uri: Some(table_uri),
1813 storage_options,
1814 is_only_declared,
1815 managed_versioning: if self.table_version_tracking_enabled {
1816 Some(true)
1817 } else {
1818 None
1819 },
1820 ..Default::default()
1821 });
1822 }
1823
1824 let mut builder = DatasetBuilder::from_uri(&table_uri);
1827 if let Some(opts) = &self.storage_options {
1828 builder = builder.with_storage_options(opts.clone());
1829 }
1830 if let Some(sess) = &self.session {
1831 builder = builder.with_session(sess.clone());
1832 }
1833 match builder.load().await {
1834 Ok(mut dataset) => {
1835 if let Some(requested_version) = request.version {
1837 dataset = dataset
1838 .checkout_version(requested_version as u64)
1839 .await
1840 .map_err(|e| {
1841 lance_core::Error::from(NamespaceError::TableVersionNotFound {
1842 message: format!(
1843 "Version {} not found for table '{}': {}",
1844 requested_version, table_name, e
1845 ),
1846 })
1847 })?;
1848 }
1849
1850 let version_info = dataset.version();
1851 let lance_schema = dataset.schema();
1852 let arrow_schema: arrow_schema::Schema = lance_schema.into();
1853 let json_schema = arrow_schema_to_json(&arrow_schema)?;
1854 let storage_options = self
1855 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1856 .await?;
1857
1858 let metadata: std::collections::HashMap<String, String> =
1860 version_info.metadata.into_iter().collect();
1861
1862 Ok(DescribeTableResponse {
1863 table: Some(table_name),
1864 namespace: request.id.as_ref().map(|id| {
1865 if id.len() > 1 {
1866 id[..id.len() - 1].to_vec()
1867 } else {
1868 vec![]
1869 }
1870 }),
1871 version: Some(version_info.version as i64),
1872 location: Some(table_uri.clone()),
1873 table_uri: Some(table_uri),
1874 schema: Some(Box::new(json_schema)),
1875 storage_options,
1876 metadata: Some(metadata),
1877 is_only_declared,
1878 managed_versioning: if self.table_version_tracking_enabled {
1879 Some(true)
1880 } else {
1881 None
1882 },
1883 ..Default::default()
1884 })
1885 }
1886 Err(err) => {
1887 if manifest::ManifestNamespace::is_not_found_load_error(&err)
1888 && is_only_declared == Some(true)
1889 {
1890 let storage_options = self
1891 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1892 .await?;
1893 Ok(DescribeTableResponse {
1894 table: Some(table_name),
1895 namespace: request.id.as_ref().map(|id| {
1896 if id.len() > 1 {
1897 id[..id.len() - 1].to_vec()
1898 } else {
1899 vec![]
1900 }
1901 }),
1902 location: Some(table_uri.clone()),
1903 table_uri: Some(table_uri),
1904 storage_options,
1905 is_only_declared,
1906 managed_versioning: if self.table_version_tracking_enabled {
1907 Some(true)
1908 } else {
1909 None
1910 },
1911 ..Default::default()
1912 })
1913 } else {
1914 Err(NamespaceError::Internal {
1915 message: format!(
1916 "Table directory exists but cannot load dataset {}: {:?}",
1917 table_name, err
1918 ),
1919 }
1920 .into())
1921 }
1922 }
1923 }
1924 }
1925
1926 fn configured_builder(&self, table_uri: &str) -> DatasetBuilder {
1929 let mut builder = DatasetBuilder::from_uri(table_uri);
1930 if let Some(opts) = &self.storage_options {
1931 builder = builder.with_storage_options(opts.clone());
1932 }
1933 if let Some(sess) = &self.session {
1934 builder = builder.with_session(sess.clone());
1935 }
1936 builder
1937 }
1938
1939 async fn load_dataset(
1940 &self,
1941 table_uri: &str,
1942 version: Option<i64>,
1943 operation: &str,
1944 ) -> Result<Dataset> {
1945 if let Some(version) = version
1946 && version < 0
1947 {
1948 return Err(NamespaceError::InvalidInput {
1949 message: format!(
1950 "Table version for {} must be non-negative, got {}",
1951 operation, version
1952 ),
1953 }
1954 .into());
1955 }
1956
1957 let builder = self.configured_builder(table_uri);
1958
1959 let dataset = builder.load().await.map_err(|e| {
1960 lance_core::Error::from(NamespaceError::TableNotFound {
1961 message: format!(
1962 "Failed to open table at '{}' for {}: {}",
1963 table_uri, operation, e
1964 ),
1965 })
1966 })?;
1967
1968 if let Some(version) = version {
1969 return dataset.checkout_version(version as u64).await.map_err(|e| {
1970 lance_core::Error::from(NamespaceError::TableVersionNotFound {
1971 message: format!(
1972 "Failed to checkout version {} for table at '{}' during {}: {}",
1973 version, table_uri, operation, e
1974 ),
1975 })
1976 });
1977 }
1978
1979 Ok(dataset)
1980 }
1981
1982 fn parse_index_type(index_type: &str) -> Result<IndexType> {
1983 match index_type.trim().to_ascii_uppercase().as_str() {
1984 "SCALAR" | "BTREE" => Ok(IndexType::BTree),
1985 "BITMAP" => Ok(IndexType::Bitmap),
1986 "LABEL_LIST" | "LABELLIST" => Ok(IndexType::LabelList),
1987 "INVERTED" | "FTS" => Ok(IndexType::Inverted),
1988 "NGRAM" => Ok(IndexType::NGram),
1989 "ZONEMAP" | "ZONE_MAP" => Ok(IndexType::ZoneMap),
1990 "BLOOMFILTER" | "BLOOM_FILTER" => Ok(IndexType::BloomFilter),
1991 "RTREE" | "R_TREE" => Ok(IndexType::RTree),
1992 "VECTOR" | "IVF_PQ" => Ok(IndexType::IvfPq),
1993 "IVF_FLAT" => Ok(IndexType::IvfFlat),
1994 "IVF_SQ" => Ok(IndexType::IvfSq),
1995 "IVF_RQ" => Ok(IndexType::IvfRq),
1996 "IVF_HNSW_FLAT" => Ok(IndexType::IvfHnswFlat),
1997 "IVF_HNSW_SQ" => Ok(IndexType::IvfHnswSq),
1998 "IVF_HNSW_PQ" => Ok(IndexType::IvfHnswPq),
1999 other => Err(NamespaceError::InvalidInput {
2000 message: format!("Unsupported index_type '{}'", other),
2001 }
2002 .into()),
2003 }
2004 }
2005
2006 fn parse_metric_type(distance_type: Option<&str>) -> Result<MetricType> {
2007 let distance_type = distance_type.unwrap_or("l2");
2008 MetricType::try_from(distance_type).map_err(|e| {
2009 lance_core::Error::from(NamespaceError::InvalidInput {
2010 message: format!(
2011 "Unsupported distance_type '{}' for vector index: {}",
2012 distance_type, e
2013 ),
2014 })
2015 })
2016 }
2017
2018 fn build_index_params(request: &CreateTableIndexRequest) -> Result<DirectoryIndexParams> {
2019 let index_type = Self::parse_index_type(&request.index_type)?;
2020 Ok(match index_type {
2021 IndexType::BTree => DirectoryIndexParams::Scalar {
2022 index_type,
2023 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
2024 },
2025 IndexType::Bitmap => DirectoryIndexParams::Scalar {
2026 index_type,
2027 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
2028 },
2029 IndexType::LabelList => DirectoryIndexParams::Scalar {
2030 index_type,
2031 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
2032 },
2033 IndexType::NGram => DirectoryIndexParams::Scalar {
2034 index_type,
2035 params: ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
2036 },
2037 IndexType::ZoneMap => DirectoryIndexParams::Scalar {
2038 index_type,
2039 params: ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
2040 },
2041 IndexType::BloomFilter => DirectoryIndexParams::Scalar {
2042 index_type,
2043 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter),
2044 },
2045 IndexType::RTree => DirectoryIndexParams::Scalar {
2046 index_type,
2047 params: ScalarIndexParams::for_builtin(BuiltinIndexType::RTree),
2048 },
2049 IndexType::Inverted => {
2050 let mut params = InvertedIndexParams::default();
2051 if let Some(with_position) = request.with_position {
2052 params = params.with_position(with_position);
2053 }
2054 if let Some(base_tokenizer) = &request.base_tokenizer {
2055 params = params.base_tokenizer(base_tokenizer.clone());
2056 }
2057 if let Some(language) = &request.language {
2058 params = params.language(language)?;
2059 }
2060 if let Some(max_token_length) = request.max_token_length {
2061 if max_token_length < 0 {
2062 return Err(NamespaceError::InvalidInput {
2063 message: format!(
2064 "FTS max_token_length must be non-negative, got {}",
2065 max_token_length
2066 ),
2067 }
2068 .into());
2069 }
2070 params = params.max_token_length(Some(max_token_length as usize));
2071 }
2072 if let Some(lower_case) = request.lower_case {
2073 params = params.lower_case(lower_case);
2074 }
2075 if let Some(stem) = request.stem {
2076 params = params.stem(stem);
2077 }
2078 if let Some(remove_stop_words) = request.remove_stop_words {
2079 params = params.remove_stop_words(remove_stop_words);
2080 }
2081 if let Some(ascii_folding) = request.ascii_folding {
2082 params = params.ascii_folding(ascii_folding);
2083 }
2084 DirectoryIndexParams::Inverted(params)
2085 }
2086 IndexType::IvfFlat => DirectoryIndexParams::Vector {
2087 index_type,
2088 params: VectorIndexParams::with_ivf_flat_params(
2089 Self::parse_metric_type(request.distance_type.as_deref())?,
2090 IvfBuildParams::default(),
2091 ),
2092 },
2093 IndexType::IvfPq => DirectoryIndexParams::Vector {
2094 index_type,
2095 params: VectorIndexParams::with_ivf_pq_params(
2096 Self::parse_metric_type(request.distance_type.as_deref())?,
2097 IvfBuildParams::default(),
2098 PQBuildParams::default(),
2099 ),
2100 },
2101 IndexType::IvfSq => DirectoryIndexParams::Vector {
2102 index_type,
2103 params: VectorIndexParams::with_ivf_sq_params(
2104 Self::parse_metric_type(request.distance_type.as_deref())?,
2105 IvfBuildParams::default(),
2106 SQBuildParams::default(),
2107 ),
2108 },
2109 IndexType::IvfRq => DirectoryIndexParams::Vector {
2110 index_type,
2111 params: VectorIndexParams::with_ivf_rq_params(
2112 Self::parse_metric_type(request.distance_type.as_deref())?,
2113 IvfBuildParams::default(),
2114 RQBuildParams::default(),
2115 ),
2116 },
2117 IndexType::IvfHnswFlat => DirectoryIndexParams::Vector {
2118 index_type,
2119 params: VectorIndexParams::ivf_hnsw(
2120 Self::parse_metric_type(request.distance_type.as_deref())?,
2121 IvfBuildParams::default(),
2122 HnswBuildParams::default(),
2123 ),
2124 },
2125 IndexType::IvfHnswSq => DirectoryIndexParams::Vector {
2126 index_type,
2127 params: VectorIndexParams::with_ivf_hnsw_sq_params(
2128 Self::parse_metric_type(request.distance_type.as_deref())?,
2129 IvfBuildParams::default(),
2130 HnswBuildParams::default(),
2131 SQBuildParams::default(),
2132 ),
2133 },
2134 IndexType::IvfHnswPq => DirectoryIndexParams::Vector {
2135 index_type,
2136 params: VectorIndexParams::with_ivf_hnsw_pq_params(
2137 Self::parse_metric_type(request.distance_type.as_deref())?,
2138 IvfBuildParams::default(),
2139 HnswBuildParams::default(),
2140 PQBuildParams::default(),
2141 ),
2142 },
2143 other => {
2144 return Err(NamespaceError::InvalidInput {
2145 message: format!("Unsupported index type for namespace API: {}", other),
2146 }
2147 .into());
2148 }
2149 })
2150 }
2151
2152 fn paginate_indices(
2153 indices: &mut Vec<IndexContent>,
2154 page_token: Option<String>,
2155 limit: Option<i32>,
2156 ) -> Option<String> {
2157 indices.sort_by(|a, b| a.index_name.cmp(&b.index_name));
2158
2159 if let Some(start_after) = page_token {
2160 if let Some(index) = indices
2161 .iter()
2162 .position(|index| index.index_name.as_str() > start_after.as_str())
2163 {
2164 indices.drain(0..index);
2165 } else {
2166 indices.clear();
2167 }
2168 }
2169
2170 let mut next_page_token = None;
2171 if let Some(limit) = limit
2172 && limit >= 0
2173 {
2174 let limit = limit as usize;
2175 if limit > 0 && indices.len() > limit {
2176 next_page_token = Some(indices[limit - 1].index_name.clone());
2177 }
2178 indices.truncate(limit);
2179 }
2180 if indices.is_empty() {
2181 None
2182 } else {
2183 next_page_token
2184 }
2185 }
2186
2187 fn transaction_operation_name(transaction: &Transaction) -> String {
2188 match &transaction.operation {
2189 Operation::CreateIndex {
2190 new_indices,
2191 removed_indices,
2192 } if new_indices.is_empty() && !removed_indices.is_empty() => "DropIndex".to_string(),
2193 _ => transaction.operation.to_string(),
2194 }
2195 }
2196
2197 fn transaction_response(
2198 version: u64,
2199 transaction: &Transaction,
2200 alteration: Option<TransactionAlteration>,
2201 ) -> DescribeTransactionResponse {
2202 let mut properties = transaction
2203 .transaction_properties
2204 .as_ref()
2205 .map(|properties| (**properties).clone())
2206 .unwrap_or_default();
2207
2208 let mut effective_status = "SUCCEEDED".to_string();
2211 if let Some(alteration) = alteration {
2212 for key in &alteration.removed_properties {
2213 properties.remove(key);
2214 }
2215 for (key, value) in alteration.properties {
2216 properties.insert(key, value);
2217 }
2218 if let Some(status) = alteration.status {
2219 effective_status = status;
2220 }
2221 }
2222
2223 properties.insert("uuid".to_string(), transaction.uuid.clone());
2224 properties.insert("version".to_string(), version.to_string());
2225 properties.insert(
2226 "read_version".to_string(),
2227 transaction.read_version.to_string(),
2228 );
2229 properties.insert(
2230 "operation".to_string(),
2231 Self::transaction_operation_name(transaction),
2232 );
2233 if let Some(tag) = &transaction.tag {
2234 properties.insert("tag".to_string(), tag.clone());
2235 }
2236
2237 DescribeTransactionResponse {
2238 status: effective_status,
2239 properties: Some(properties),
2240 }
2241 }
2242
2243 fn describe_table_index_stats_response(
2244 stats: &serde_json::Value,
2245 ) -> DescribeTableIndexStatsResponse {
2246 let get_i64 = |key: &str| {
2247 stats.get(key).and_then(|value| {
2248 value
2249 .as_i64()
2250 .or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
2251 })
2252 };
2253
2254 DescribeTableIndexStatsResponse {
2255 distance_type: stats
2256 .get("distance_type")
2257 .and_then(|value| value.as_str())
2258 .map(str::to_string),
2259 index_type: stats
2260 .get("index_type")
2261 .and_then(|value| value.as_str())
2262 .map(str::to_string),
2263 num_indexed_rows: get_i64("num_indexed_rows"),
2264 num_unindexed_rows: get_i64("num_unindexed_rows"),
2265 num_indices: get_i64("num_indices").and_then(|value| i32::try_from(value).ok()),
2266 }
2267 }
2268
2269 async fn find_transaction(&self, dataset: &Dataset, id: &str) -> Result<(u64, Transaction)> {
2274 if let Ok(version) = id.parse::<u64>() {
2275 let transaction = dataset
2276 .read_transaction_by_version(version)
2277 .await
2278 .map_err(|e| {
2279 lance_core::Error::from(NamespaceError::TransactionNotFound {
2280 message: format!(
2281 "Failed to read transaction for version {}: {}",
2282 version, e
2283 ),
2284 })
2285 })?
2286 .ok_or_else(|| {
2287 lance_core::Error::from(NamespaceError::TransactionNotFound {
2288 message: format!("version {}", version),
2289 })
2290 })?;
2291 return Ok((version, transaction));
2292 }
2293
2294 let versions = dataset.versions().await.map_err(|e| {
2295 lance_core::Error::from(NamespaceError::Internal {
2296 message: format!(
2297 "Failed to list table versions while resolving transaction '{}': {}",
2298 id, e
2299 ),
2300 })
2301 })?;
2302
2303 for version in versions.into_iter().rev() {
2304 if let Some(transaction) = dataset
2305 .read_transaction_by_version(version.version)
2306 .await
2307 .map_err(|e| {
2308 lance_core::Error::from(NamespaceError::Internal {
2309 message: format!(
2310 "Failed to read transaction for version {} while resolving '{}': {}",
2311 version.version, id, e
2312 ),
2313 })
2314 })?
2315 && transaction.uuid == id
2316 {
2317 return Ok((version.version, transaction));
2318 }
2319 }
2320
2321 Err(NamespaceError::TransactionNotFound {
2322 message: id.to_string(),
2323 }
2324 .into())
2325 }
2326
2327 const TRANSACTION_ALTERATIONS_DIR: &'static str = "_alter_transactions";
2331
2332 fn transaction_alteration_path(&self, table_uri: &str, txn_uuid: &str) -> Result<Path> {
2333 let table_path = self.object_store_path_from_uri(table_uri)?;
2334 Ok(table_path
2335 .join(Self::TRANSACTION_ALTERATIONS_DIR)
2336 .join(format!("{}.json", txn_uuid).as_str()))
2337 }
2338
2339 async fn load_transaction_alteration(
2340 &self,
2341 table_uri: &str,
2342 txn_uuid: &str,
2343 ) -> Result<Option<TransactionAlteration>> {
2344 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2345 match self.object_store.inner.get(&path).await {
2346 Ok(get_result) => {
2347 let bytes = get_result.bytes().await.map_err(|e| {
2348 lance_core::Error::from(NamespaceError::Internal {
2349 message: format!(
2350 "Failed to read alter_transaction sidecar for '{}': {}",
2351 txn_uuid, e
2352 ),
2353 })
2354 })?;
2355 let alteration = TransactionAlteration::from_json_slice(&bytes).map_err(|e| {
2356 lance_core::Error::from(NamespaceError::Internal {
2357 message: format!(
2358 "Failed to parse alter_transaction sidecar for '{}': {}",
2359 txn_uuid, e
2360 ),
2361 })
2362 })?;
2363 Ok(Some(alteration))
2364 }
2365 Err(ObjectStoreError::NotFound { .. }) => Ok(None),
2366 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2367 message: format!(
2368 "Failed to load alter_transaction sidecar for '{}': {}",
2369 txn_uuid, e
2370 ),
2371 })),
2372 }
2373 }
2374
2375 async fn save_transaction_alteration(
2376 &self,
2377 table_uri: &str,
2378 txn_uuid: &str,
2379 alteration: &TransactionAlteration,
2380 ) -> Result<()> {
2381 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2382 let bytes = alteration.to_json_bytes().map_err(|e| {
2383 lance_core::Error::from(NamespaceError::Internal {
2384 message: format!(
2385 "Failed to serialize alter_transaction sidecar for '{}': {}",
2386 txn_uuid, e
2387 ),
2388 })
2389 })?;
2390 self.object_store
2391 .inner
2392 .put(&path, bytes.into())
2393 .await
2394 .map_err(|e| {
2395 lance_core::Error::from(NamespaceError::Internal {
2396 message: format!(
2397 "Failed to persist alter_transaction sidecar for '{}': {}",
2398 txn_uuid, e
2399 ),
2400 })
2401 })?;
2402 Ok(())
2403 }
2404
2405 fn table_full_uri(&self, table_name: &str) -> String {
2406 format!("{}/{}.lance", self.root, table_name)
2407 }
2408
2409 fn table_path(&self, table_name: &str) -> Path {
2411 self.base_path
2412 .clone()
2413 .join(format!("{}.lance", table_name).as_str())
2414 }
2415
2416 fn table_reserved_file_path(&self, table_name: &str) -> Path {
2418 self.base_path
2419 .clone()
2420 .join(format!("{}.lance", table_name).as_str())
2421 .join(".lance-reserved")
2422 }
2423
2424 fn table_deregistered_file_path(&self, table_name: &str) -> Path {
2426 self.base_path
2427 .clone()
2428 .join(format!("{}.lance", table_name).as_str())
2429 .join(".lance-deregistered")
2430 }
2431
2432 pub(crate) async fn check_table_status(&self, table_name: &str) -> TableStatus {
2438 let table_path = self.table_path(table_name);
2439 match self.object_store.read_dir(table_path).await {
2440 Ok(entries) => {
2441 let exists = !entries.is_empty();
2442 let is_deregistered = entries.iter().any(|e| e.ends_with(".lance-deregistered"));
2443 let has_reserved_file = entries.iter().any(|e| e.ends_with(".lance-reserved"));
2444 TableStatus {
2445 exists,
2446 is_deregistered,
2447 has_reserved_file,
2448 }
2449 }
2450 Err(_) => TableStatus {
2451 exists: false,
2452 is_deregistered: false,
2453 has_reserved_file: false,
2454 },
2455 }
2456 }
2457
2458 async fn put_marker_file_atomic(
2459 &self,
2460 path: &Path,
2461 file_description: &str,
2462 ) -> std::result::Result<(), String> {
2463 let put_opts = PutOptions {
2464 mode: PutMode::Create,
2465 ..Default::default()
2466 };
2467
2468 match self
2469 .object_store
2470 .inner
2471 .put_opts(path, bytes::Bytes::new().into(), put_opts)
2472 .await
2473 {
2474 Ok(_) => Ok(()),
2475 Err(ObjectStoreError::AlreadyExists { .. })
2476 | Err(ObjectStoreError::Precondition { .. }) => {
2477 Err(format!("{} already exists", file_description))
2478 }
2479 Err(e) => Err(format!("Failed to create {}: {:?}", file_description, e)),
2480 }
2481 }
2482
2483 async fn get_storage_options_for_table(
2503 &self,
2504 table_uri: &str,
2505 vend_credentials: bool,
2506 identity: Option<&Identity>,
2507 ) -> Result<Option<HashMap<String, String>>> {
2508 if vend_credentials && let Some(ref vendor) = self.credential_vendor {
2509 let vended = vendor.vend_credentials(table_uri, identity).await?;
2510 return Ok(Some(vended.storage_options));
2511 }
2512 if self.vend_input_storage_options {
2515 let mut options = self.storage_options.clone().unwrap_or_default();
2516 if let Some(refresh_interval_millis) =
2518 self.vend_input_storage_options_refresh_interval_millis
2519 {
2520 let now_millis = std::time::SystemTime::now()
2521 .duration_since(std::time::UNIX_EPOCH)
2522 .unwrap()
2523 .as_millis() as u64;
2524 let expires_at_millis = now_millis + refresh_interval_millis;
2525 options.insert(
2526 "expires_at_millis".to_string(),
2527 expires_at_millis.to_string(),
2528 );
2529 }
2530 return Ok(Some(options));
2531 }
2532 Ok(None)
2535 }
2536
2537 pub async fn migrate(&self) -> Result<usize> {
2591 let Some(manifest_ns) = self.manifest_ns_for_write().await? else {
2593 return Ok(0); };
2595
2596 let manifest_locations = manifest_ns.list_manifest_table_locations().await?;
2598
2599 let dir_tables = self
2602 .filter_declared_tables(self.list_directory_tables().await?, false)
2603 .await?;
2604
2605 let mut migrated_count = 0;
2610 for table_name in dir_tables {
2611 let dir_name = format!("{}.lance", table_name);
2613 if !manifest_locations.contains(&dir_name) {
2614 manifest_ns.register_table(&table_name, dir_name).await?;
2615 migrated_count += 1;
2616 }
2617 }
2618
2619 Ok(migrated_count)
2620 }
2621
2622 async fn delete_physical_version_files(
2630 &self,
2631 table_entries: &[TableDeleteEntry],
2632 branch: Option<&str>,
2633 ) -> Result<i64> {
2634 let mut deleted_count = 0i64;
2635 for te in table_entries {
2636 let table_uri = self.resolve_table_location(&te.table_id).await?;
2637 let table_uri = match branch {
2638 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
2639 None => table_uri,
2640 };
2641 let table_path = self.object_store_path_from_uri(&table_uri)?;
2642 let versions_dir_path = table_path.clone().join(VERSIONS_DIR);
2643
2644 let manifest_metas: Vec<_> = self
2646 .object_store
2647 .read_dir_all(&versions_dir_path, None)
2648 .try_collect()
2649 .await
2650 .map_err(|e| {
2651 lance_core::Error::from(NamespaceError::Internal {
2652 message: format!(
2653 "Failed to list manifest files for table at '{}': {}",
2654 table_uri, e
2655 ),
2656 })
2657 })?;
2658 let location_by_version: HashMap<u64, Path> = manifest_metas
2659 .into_iter()
2660 .filter_map(|meta| {
2661 let version = Self::manifest_version_from_filename(meta.location.filename()?)?;
2662 Some((version, meta.location))
2663 })
2664 .collect();
2665
2666 for (&v, version_path) in &location_by_version {
2667 let vi = v as i64;
2668 if !te.ranges.iter().any(|&(s, e)| vi >= s && (e < 0 || vi < e)) {
2669 continue;
2670 }
2671 match self.object_store.inner.delete(version_path).await {
2672 Ok(_) => {
2673 deleted_count += 1;
2674 }
2675 Err(object_store::Error::NotFound { .. }) => {}
2676 Err(e) => {
2677 return Err(NamespaceError::Internal {
2678 message: format!(
2679 "Failed to delete version {} for table at '{}': {}",
2680 v, table_uri, e
2681 ),
2682 }
2683 .into());
2684 }
2685 }
2686 }
2687 }
2688 Ok(deleted_count)
2689 }
2690
2691 #[allow(clippy::too_many_arguments)]
2696 fn apply_query_params_to_scanner(
2697 scanner: &mut Scanner,
2698 filter: Option<&str>,
2699 columns: Option<&QueryTableRequestColumns>,
2700 vector_column: Option<&str>,
2701 vector: &QueryTableRequestVector,
2702 k: i32,
2703 offset: Option<i32>,
2704 prefilter: Option<bool>,
2705 bypass_vector_index: Option<bool>,
2706 nprobes: Option<i32>,
2707 ef: Option<i32>,
2708 refine_factor: Option<i32>,
2709 distance_type: Option<&str>,
2710 fast_search_flag: Option<bool>,
2711 with_row_id: Option<bool>,
2712 lower_bound: Option<f32>,
2713 upper_bound: Option<f32>,
2714 operation: &str,
2715 ) -> Result<()> {
2716 if let Some(pf) = prefilter {
2718 scanner.prefilter(pf);
2719 }
2720
2721 if let Some(filter) = filter {
2722 scanner.filter(filter).map_err(|e| {
2723 Error::invalid_input_source(
2724 format!("Invalid filter expression for {}: {}", operation, e).into(),
2725 )
2726 })?;
2727 }
2728
2729 if let Some(cols) = columns {
2730 if let Some(ref names) = cols.column_names {
2731 scanner.project(names.as_slice()).map_err(|e| {
2732 Error::invalid_input_source(
2733 format!("Invalid column projection for {}: {}", operation, e).into(),
2734 )
2735 })?;
2736 } else if let Some(ref aliases) = cols.column_aliases {
2737 let pairs: Vec<(&str, &str)> = aliases
2739 .iter()
2740 .map(|(alias, src)| (alias.as_str(), src.as_str()))
2741 .collect();
2742 scanner.project_with_transform(&pairs).map_err(|e| {
2743 Error::invalid_input_source(
2744 format!("Invalid column aliases for {}: {}", operation, e).into(),
2745 )
2746 })?;
2747 }
2748 }
2749
2750 let query_vec: Option<Vec<f32>> = vector
2752 .single_vector
2753 .as_ref()
2754 .filter(|v| !v.is_empty())
2755 .cloned()
2756 .or_else(|| {
2757 vector
2758 .multi_vector
2759 .as_ref()
2760 .and_then(|mv| mv.first())
2761 .filter(|v| !v.is_empty())
2762 .cloned()
2763 });
2764
2765 if let Some(q_vec) = query_vec {
2766 let col = vector_column.unwrap_or("vector");
2767 let q = Arc::new(Float32Array::from(q_vec));
2768 scanner
2769 .nearest(col, q.as_ref(), k.max(1) as usize)
2770 .map_err(|e| {
2771 Error::invalid_input_source(
2772 format!("Invalid vector query for {}: {}", operation, e).into(),
2773 )
2774 })?;
2775
2776 if let Some(n) = nprobes {
2778 scanner.nprobes(n.max(1) as usize);
2779 }
2780 if let Some(e) = ef {
2781 scanner.ef(e.max(1) as usize);
2782 }
2783 if let Some(rf) = refine_factor {
2784 scanner.refine(rf.max(0) as u32);
2785 }
2786 if let Some(true) = bypass_vector_index {
2788 scanner.use_index(false);
2789 }
2790 if let Some(true) = fast_search_flag {
2791 scanner.fast_search();
2792 }
2793 if lower_bound.is_some() || upper_bound.is_some() {
2794 scanner.distance_range(lower_bound, upper_bound);
2795 }
2796 if let Some(dt) = distance_type {
2797 let metric = Self::parse_metric_type(Some(dt))?;
2798 scanner.distance_metric(metric);
2799 }
2800 if let Some(off) = offset.filter(|&o| o > 0) {
2802 scanner.limit(None, Some(off as i64)).map_err(|e| {
2803 Error::invalid_input_source(
2804 format!("Invalid offset for {}: {}", operation, e).into(),
2805 )
2806 })?;
2807 }
2808 } else {
2809 let limit = if k > 0 { Some(k as i64) } else { None };
2811 scanner
2812 .limit(limit, offset.map(|o| o as i64))
2813 .map_err(|e| {
2814 Error::invalid_input_source(
2815 format!("Invalid limit/offset for {}: {}", operation, e).into(),
2816 )
2817 })?;
2818 }
2819
2820 if let Some(true) = with_row_id {
2821 scanner.with_row_id();
2822 }
2823
2824 Ok(())
2825 }
2826
2827 pub fn retrieve_ops_metrics(&self) -> HashMap<String, u64> {
2834 self.ops_metrics
2835 .as_ref()
2836 .map(|m| m.retrieve())
2837 .unwrap_or_default()
2838 }
2839
2840 pub fn reset_ops_metrics(&self) {
2844 if let Some(ref metrics) = self.ops_metrics {
2845 metrics.reset();
2846 }
2847 }
2848
2849 fn record_op(&self, operation: &str) {
2851 if let Some(ref metrics) = self.ops_metrics {
2852 metrics.increment(operation);
2853 }
2854 }
2855}
2856
2857#[async_trait]
2858impl LanceNamespace for DirectoryNamespace {
2859 async fn list_namespaces(
2860 &self,
2861 request: ListNamespacesRequest,
2862 ) -> Result<ListNamespacesResponse> {
2863 self.record_op("list_namespaces");
2864 if let Some(manifest_ns) = self.manifest_ns_for_read() {
2865 return manifest_ns.list_namespaces(request).await;
2866 }
2867
2868 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
2869 return Err(self.child_namespace_requires_manifest_error());
2870 }
2871 Self::validate_root_namespace_id(&request.id)?;
2872 Ok(ListNamespacesResponse::new(vec![]))
2873 }
2874
2875 async fn describe_namespace(
2876 &self,
2877 request: DescribeNamespaceRequest,
2878 ) -> Result<DescribeNamespaceResponse> {
2879 self.record_op("describe_namespace");
2880 if let Some(manifest_ns) = self.manifest_ns_for_read() {
2881 return manifest_ns.describe_namespace(request).await;
2882 }
2883
2884 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
2885 return Err(self.child_namespace_requires_manifest_error());
2886 }
2887 Self::validate_root_namespace_id(&request.id)?;
2888 #[allow(clippy::needless_update)]
2889 Ok(DescribeNamespaceResponse {
2890 properties: Some(HashMap::new()),
2891 ..Default::default()
2892 })
2893 }
2894
2895 async fn create_namespace(
2896 &self,
2897 request: CreateNamespaceRequest,
2898 ) -> Result<CreateNamespaceResponse> {
2899 self.record_op("create_namespace");
2900 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
2901 return manifest_ns.create_namespace(request).await;
2902 }
2903
2904 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
2905 return Err(NamespaceError::NamespaceAlreadyExists {
2906 message: "root namespace".to_string(),
2907 }
2908 .into());
2909 }
2910
2911 Err(NamespaceError::Unsupported {
2912 message: "Child namespaces are only supported when manifest mode is enabled"
2913 .to_string(),
2914 }
2915 .into())
2916 }
2917
2918 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
2919 self.record_op("drop_namespace");
2920 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
2921 return manifest_ns.drop_namespace(request).await;
2922 }
2923
2924 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
2925 return Err(NamespaceError::InvalidInput {
2926 message: "Root namespace cannot be dropped".to_string(),
2927 }
2928 .into());
2929 }
2930
2931 Err(NamespaceError::Unsupported {
2932 message: "Child namespaces are only supported when manifest mode is enabled"
2933 .to_string(),
2934 }
2935 .into())
2936 }
2937
2938 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
2939 self.record_op("namespace_exists");
2940 if let Some(manifest_ns) = self.manifest_ns_for_read() {
2941 return manifest_ns.namespace_exists(request).await;
2942 }
2943
2944 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
2945 return Ok(());
2946 }
2947
2948 Err(self.child_namespace_requires_manifest_error())
2949 }
2950
2951 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
2952 self.record_op("list_tables");
2953 let namespace_id = request.id.as_ref().ok_or_else(|| {
2955 lance_core::Error::from(NamespaceError::InvalidInput {
2956 message: "Namespace ID is required".to_string(),
2957 })
2958 })?;
2959
2960 if !namespace_id.is_empty() {
2962 if let Some(manifest_ns) = self.manifest_ns_for_read() {
2963 return manifest_ns.list_tables(request).await;
2964 }
2965 return Err(self.child_namespace_requires_manifest_error());
2966 }
2967
2968 if let Some(manifest_ns) = self.manifest_ns_for_read()
2970 && !self.dir_listing_enabled
2971 {
2972 return manifest_ns.list_tables(request).await;
2973 }
2974 if !self.dir_listing_enabled {
2975 return Ok(ListTablesResponse::new(vec![]));
2976 }
2977
2978 let mut tables = if self.manifest_ns_for_read().is_some()
2981 && self.dir_listing_enabled
2982 && self.dir_listing_to_manifest_migration_enabled
2983 {
2984 let manifest_locations = if let Some(manifest_ns) = self.manifest_ns_for_read() {
2986 manifest_ns.list_manifest_table_locations().await?
2987 } else {
2988 std::collections::HashSet::new()
2989 };
2990
2991 let mut manifest_request = request.clone();
2993 manifest_request.limit = None;
2994 manifest_request.page_token = None;
2995 let manifest_tables = if let Some(manifest_ns) = self.manifest_ns_for_read() {
2996 let manifest_response = manifest_ns.list_tables(manifest_request).await?;
2997 manifest_response.tables
2998 } else {
2999 vec![]
3000 };
3001
3002 let mut all_tables: Vec<String> = manifest_tables;
3005 let dir_tables = self.list_directory_tables().await?;
3006 for table_name in dir_tables {
3007 let full_location = format!("{}/{}.lance", self.root, table_name);
3010 let relative_location = format!("{}.lance", table_name);
3011 if !manifest_locations.contains(&full_location)
3012 && !manifest_locations.contains(&relative_location)
3013 {
3014 all_tables.push(table_name);
3015 }
3016 }
3017
3018 all_tables
3019 } else {
3020 self.list_directory_tables().await?
3021 };
3022
3023 tables = self
3024 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
3025 .await?;
3026
3027 let next_page_token =
3029 Self::apply_pagination(&mut tables, request.page_token, request.limit);
3030 let mut response = ListTablesResponse::new(tables);
3031 response.page_token = next_page_token;
3032 Ok(response)
3033 }
3034
3035 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
3036 self.record_op("describe_table");
3037 self.describe_table_impl(request).await
3038 }
3039
3040 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
3041 self.record_op("table_exists");
3042 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
3043 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
3044 let skip_manifest_for_root = self.dir_listing_enabled
3045 && is_root_level
3046 && !self.dir_listing_to_manifest_migration_enabled;
3047 if let Some(manifest_ns) = self.manifest_ns_for_read()
3048 && !skip_manifest_for_root
3049 {
3050 match manifest_ns.table_exists(request.clone()).await {
3051 Ok(()) => return Ok(()),
3052 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
3053 return Err(e);
3056 }
3057 Err(_) if self.dir_listing_enabled && is_root_level => {
3058 }
3060 Err(e) => return Err(e),
3061 }
3062 }
3063 if is_child_table {
3064 return Err(self.child_namespace_requires_manifest_error());
3065 }
3066
3067 let table_name = Self::table_name_from_id(&request.id)?;
3068 let table_id = Self::format_table_id_from_request(&request.id);
3069 if !self.dir_listing_enabled {
3070 return Err(NamespaceError::TableNotFound { message: table_id }.into());
3071 }
3072
3073 let status = self.check_table_status(&table_name).await;
3075
3076 if !status.exists {
3077 return Err(NamespaceError::TableNotFound {
3078 message: table_id.clone(),
3079 }
3080 .into());
3081 }
3082
3083 if status.is_deregistered {
3084 return Err(NamespaceError::TableNotFound {
3085 message: format!("Table is deregistered: {}", table_id),
3086 }
3087 .into());
3088 }
3089
3090 Ok(())
3091 }
3092
3093 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3094 self.record_op("drop_table");
3095 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3096 return manifest_ns.drop_table(request).await;
3097 }
3098
3099 let table_name = Self::table_name_from_id(&request.id)?;
3100 let table_uri = self.table_full_uri(&table_name);
3101 let table_path = self.table_path(&table_name);
3102
3103 self.object_store
3104 .remove_dir_all(table_path)
3105 .await
3106 .map_err(|e| {
3107 lance_core::Error::from(NamespaceError::Internal {
3108 message: format!("Failed to drop table {}: {:?}", table_name, e),
3109 })
3110 })?;
3111
3112 Ok(DropTableResponse {
3113 id: request.id,
3114 location: Some(table_uri),
3115 ..Default::default()
3116 })
3117 }
3118
3119 async fn create_table(
3120 &self,
3121 request: CreateTableRequest,
3122 request_data: Bytes,
3123 ) -> Result<CreateTableResponse> {
3124 self.record_op("create_table");
3125 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3126 return manifest_ns.create_table(request, request_data).await;
3127 }
3128
3129 Self::validate_dir_only_properties(request.properties.as_ref(), "create_table")?;
3130
3131 let table_name = Self::table_name_from_id(&request.id)?;
3132 let table_uri = self.table_full_uri(&table_name);
3133 let status = self.check_table_status(&table_name).await;
3134 let (reader, _num_rows) =
3135 Self::ipc_reader_from_request_data(&request_data, "create_table")?;
3136
3137 if status.exists && self.table_has_actual_manifests(&table_name).await? {
3138 return Err(NamespaceError::TableAlreadyExists {
3139 message: table_name,
3140 }
3141 .into());
3142 }
3143
3144 let write_result = self
3145 .write_reader_to_table(
3146 &table_uri,
3147 reader,
3148 WriteMode::Create,
3149 request.storage_options.clone(),
3150 )
3151 .await;
3152 if let Err(err) = write_result {
3153 if self.table_uri_has_actual_manifests(&table_uri).await? {
3154 return Err(NamespaceError::TableAlreadyExists {
3155 message: table_name,
3156 }
3157 .into());
3158 }
3159 return Err(err);
3160 }
3161 Ok(CreateTableResponse {
3162 version: Some(1),
3163 location: Some(table_uri),
3164 storage_options: self.storage_options.clone(),
3165 properties: request.properties,
3166 ..Default::default()
3167 })
3168 }
3169
3170 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3171 self.record_op("declare_table");
3172 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3173 let mut response = manifest_ns.declare_table(request.clone()).await?;
3174 if let Some(ref location) = response.location {
3175 let vend = request.vend_credentials.unwrap_or(true);
3177 let identity = request.identity.as_deref();
3178 response.storage_options = self
3179 .get_storage_options_for_table(location, vend, identity)
3180 .await?;
3181 }
3182 if self.table_version_tracking_enabled {
3184 response.managed_versioning = Some(true);
3185 }
3186 return Ok(response);
3187 }
3188
3189 Self::validate_dir_only_properties(request.properties.as_ref(), "declare_table")?;
3190
3191 let table_name = Self::table_name_from_id(&request.id)?;
3192 let table_uri = self.table_full_uri(&table_name);
3193
3194 if let Some(location) = &request.location {
3196 let location = location.trim_end_matches('/');
3197 if location != table_uri {
3198 return Err(NamespaceError::InvalidInput {
3199 message: format!(
3200 "Cannot declare table {} at location {}, must be at location {}",
3201 table_name, location, table_uri
3202 ),
3203 }
3204 .into());
3205 }
3206 }
3207
3208 let status = self.check_table_status(&table_name).await;
3212 if status.exists && !status.has_reserved_file {
3213 return Err(NamespaceError::TableAlreadyExists {
3215 message: table_name.to_string(),
3216 }
3217 .into());
3218 }
3219
3220 let reserved_file_path = self.table_reserved_file_path(&table_name);
3224
3225 self.put_marker_file_atomic(&reserved_file_path, &format!("table {}", table_name))
3226 .await
3227 .map_err(|e| {
3228 if e.contains("already exists") {
3229 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3230 message: table_name.to_string(),
3231 })
3232 } else {
3233 lance_core::Error::from(NamespaceError::Internal { message: e })
3234 }
3235 })?;
3236
3237 let vend_credentials = request.vend_credentials.unwrap_or(true);
3239 let identity = request.identity.as_deref();
3240 let storage_options = self
3241 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
3242 .await?;
3243
3244 Ok(DeclareTableResponse {
3245 location: Some(table_uri),
3246 storage_options,
3247 properties: request.properties,
3248 managed_versioning: if self.table_version_tracking_enabled {
3249 Some(true)
3250 } else {
3251 None
3252 },
3253 ..Default::default()
3254 })
3255 }
3256
3257 async fn register_table(
3258 &self,
3259 request: lance_namespace::models::RegisterTableRequest,
3260 ) -> Result<lance_namespace::models::RegisterTableResponse> {
3261 self.record_op("register_table");
3262 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3264 return LanceNamespace::register_table(manifest_ns.as_ref(), request).await;
3265 }
3266
3267 Err(NamespaceError::Unsupported {
3269 message: "register_table is only supported when manifest mode is enabled".to_string(),
3270 }
3271 .into())
3272 }
3273
3274 async fn deregister_table(
3275 &self,
3276 request: lance_namespace::models::DeregisterTableRequest,
3277 ) -> Result<lance_namespace::models::DeregisterTableResponse> {
3278 self.record_op("deregister_table");
3279 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3281 return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await;
3282 }
3283
3284 let table_name = Self::table_name_from_id(&request.id)?;
3286 let table_uri = self.table_full_uri(&table_name);
3287
3288 let status = self.check_table_status(&table_name).await;
3291
3292 if !status.exists {
3293 return Err(NamespaceError::TableNotFound {
3294 message: table_name.to_string(),
3295 }
3296 .into());
3297 }
3298
3299 if status.is_deregistered {
3300 return Err(NamespaceError::TableNotFound {
3301 message: format!("Table is already deregistered: {}", table_name),
3302 }
3303 .into());
3304 }
3305
3306 let deregistered_path = self.table_deregistered_file_path(&table_name);
3312 self.put_marker_file_atomic(
3313 &deregistered_path,
3314 &format!("deregistration marker for table {}", table_name),
3315 )
3316 .await
3317 .map_err(|e| {
3318 if e.contains("already exists") {
3319 lance_core::Error::from(NamespaceError::InvalidTableState {
3320 message: format!("Table is already deregistered: {}", table_name),
3321 })
3322 } else {
3323 lance_core::Error::from(NamespaceError::Internal { message: e })
3324 }
3325 })?;
3326
3327 Ok(lance_namespace::models::DeregisterTableResponse {
3328 id: request.id,
3329 location: Some(table_uri),
3330 ..Default::default()
3331 })
3332 }
3333
3334 async fn alter_table_add_columns(
3335 &self,
3336 request: AlterTableAddColumnsRequest,
3337 ) -> Result<AlterTableAddColumnsResponse> {
3338 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3339 return manifest_ns.alter_table_add_columns(request).await;
3340 }
3341
3342 let table_name = Self::table_name_from_id(&request.id)?;
3344 let table_uri = self.table_full_uri(&table_name);
3345
3346 let status = self.check_table_status(&table_name).await;
3348 if !status.exists {
3349 return Err(NamespaceError::TableNotFound {
3350 message: table_name,
3351 }
3352 .into());
3353 }
3354 if status.is_deregistered {
3355 return Err(NamespaceError::TableNotFound {
3356 message: format!("Table is deregistered: {}", table_name),
3357 }
3358 .into());
3359 }
3360
3361 let mut dataset = self
3362 .configured_builder(&table_uri)
3363 .load()
3364 .await
3365 .map_err(|e| {
3366 Error::io_source(box_error(std::io::Error::other(format!(
3367 "Failed to open dataset: {}",
3368 e
3369 ))))
3370 })?;
3371
3372 let sql_expressions = build_sql_expressions(&request.new_columns)?;
3373
3374 dataset
3375 .add_columns(
3376 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3377 None,
3378 None,
3379 )
3380 .await
3381 .map_err(|e| {
3382 Error::io_source(box_error(std::io::Error::other(format!(
3383 "Failed to add columns: {}",
3384 e
3385 ))))
3386 })?;
3387
3388 let version = dataset.version().version as i64;
3389 Ok(AlterTableAddColumnsResponse::new(version))
3390 }
3391
3392 async fn alter_table_alter_columns(
3393 &self,
3394 request: AlterTableAlterColumnsRequest,
3395 ) -> Result<AlterTableAlterColumnsResponse> {
3396 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3397 return manifest_ns.alter_table_alter_columns(request).await;
3398 }
3399
3400 let table_name = Self::table_name_from_id(&request.id)?;
3401 let table_uri = self.table_full_uri(&table_name);
3402
3403 let status = self.check_table_status(&table_name).await;
3405 if !status.exists {
3406 return Err(NamespaceError::TableNotFound {
3407 message: table_name,
3408 }
3409 .into());
3410 }
3411 if status.is_deregistered {
3412 return Err(NamespaceError::TableNotFound {
3413 message: format!("Table is deregistered: {}", table_name),
3414 }
3415 .into());
3416 }
3417
3418 let mut dataset = self
3419 .configured_builder(&table_uri)
3420 .load()
3421 .await
3422 .map_err(|e| {
3423 Error::io_source(box_error(std::io::Error::other(format!(
3424 "Failed to open dataset: {}",
3425 e
3426 ))))
3427 })?;
3428
3429 let alterations = build_column_alterations(&request.alterations)?;
3430
3431 dataset.alter_columns(&alterations).await.map_err(|e| {
3432 Error::io_source(box_error(std::io::Error::other(format!(
3433 "Failed to alter columns: {}",
3434 e
3435 ))))
3436 })?;
3437
3438 let version = dataset.version().version as i64;
3439 Ok(AlterTableAlterColumnsResponse::new(version))
3440 }
3441
3442 async fn alter_table_drop_columns(
3443 &self,
3444 request: AlterTableDropColumnsRequest,
3445 ) -> Result<AlterTableDropColumnsResponse> {
3446 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3447 return manifest_ns.alter_table_drop_columns(request).await;
3448 }
3449
3450 let table_name = Self::table_name_from_id(&request.id)?;
3451 let table_uri = self.table_full_uri(&table_name);
3452
3453 let status = self.check_table_status(&table_name).await;
3455 if !status.exists {
3456 return Err(NamespaceError::TableNotFound {
3457 message: table_name,
3458 }
3459 .into());
3460 }
3461 if status.is_deregistered {
3462 return Err(NamespaceError::TableNotFound {
3463 message: format!("Table is deregistered: {}", table_name),
3464 }
3465 .into());
3466 }
3467
3468 let mut dataset = self
3469 .configured_builder(&table_uri)
3470 .load()
3471 .await
3472 .map_err(|e| {
3473 Error::io_source(box_error(std::io::Error::other(format!(
3474 "Failed to open dataset: {}",
3475 e
3476 ))))
3477 })?;
3478
3479 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
3480 dataset.drop_columns(&columns).await.map_err(|e| {
3481 Error::io_source(box_error(std::io::Error::other(format!(
3482 "Failed to drop columns: {}",
3483 e
3484 ))))
3485 })?;
3486
3487 let version = dataset.version().version as i64;
3488 Ok(AlterTableDropColumnsResponse::new(version))
3489 }
3490
3491 async fn list_table_versions(
3492 &self,
3493 request: ListTableVersionsRequest,
3494 ) -> Result<ListTableVersionsResponse> {
3495 self.record_op("list_table_versions");
3496 let branch = Self::normalized_branch(request.branch.as_deref())?;
3497 let table_uri = self.resolve_table_location(&request.id).await?;
3498 let table_uri = match branch {
3499 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3500 None => table_uri,
3501 };
3502 let want_descending = request.descending == Some(true);
3503 let table_versions = self
3504 .list_table_versions_from_storage(&table_uri, want_descending, request.limit)
3505 .await?;
3506
3507 Ok(ListTableVersionsResponse {
3508 versions: table_versions,
3509 page_token: None,
3510 })
3511 }
3512
3513 async fn create_table_version(
3514 &self,
3515 request: CreateTableVersionRequest,
3516 ) -> Result<CreateTableVersionResponse> {
3517 self.record_op("create_table_version");
3518 let branch = Self::normalized_branch(request.branch.as_deref())?;
3519 let table_uri = self.resolve_table_location(&request.id).await?;
3520 let (table_uri, table_path) = match branch {
3521 Some(b) => self.resolve_branch_for_commit(&table_uri, b).await?,
3522 None => {
3523 let table_path = self.object_store_path_from_uri(&table_uri)?;
3524 (table_uri, table_path)
3525 }
3526 };
3527
3528 let staging_manifest_path = &request.manifest_path;
3529 let version = request.version as u64;
3530
3531 let naming_scheme = match request.naming_scheme.as_deref() {
3533 Some("V1") => ManifestNamingScheme::V1,
3534 _ => ManifestNamingScheme::V2,
3535 };
3536
3537 let final_path = naming_scheme.manifest_path(&table_path, version);
3539
3540 let staging_path = Path::parse(staging_manifest_path).map_err(|e| {
3541 lance_core::Error::from(NamespaceError::InvalidInput {
3542 message: format!(
3543 "Invalid staging manifest path '{}': {}",
3544 staging_manifest_path, e
3545 ),
3546 })
3547 })?;
3548
3549 let copy_result = match self
3550 .object_store
3551 .inner
3552 .copy_if_not_exists(&staging_path, &final_path)
3553 .await
3554 {
3555 Ok(()) => Ok(()),
3556 Err(ObjectStoreError::NotImplemented { .. })
3557 | Err(ObjectStoreError::NotSupported { .. }) => {
3558 let manifest_data = self
3559 .object_store
3560 .inner
3561 .get(&staging_path)
3562 .await
3563 .map_err(|e| {
3564 lance_core::Error::from(NamespaceError::Internal {
3565 message: format!(
3566 "Failed to read staging manifest at '{}': {}",
3567 staging_manifest_path, e
3568 ),
3569 })
3570 })?
3571 .bytes()
3572 .await
3573 .map_err(|e| {
3574 lance_core::Error::from(NamespaceError::Internal {
3575 message: format!(
3576 "Failed to read staging manifest bytes at '{}': {}",
3577 staging_manifest_path, e
3578 ),
3579 })
3580 })?;
3581 self.object_store
3582 .inner
3583 .put_opts(
3584 &final_path,
3585 manifest_data.into(),
3586 PutOptions {
3587 mode: PutMode::Create,
3588 ..Default::default()
3589 },
3590 )
3591 .await
3592 .map(|_| ())
3593 }
3594 Err(e) => Err(e),
3595 };
3596
3597 match copy_result {
3598 Ok(()) => {}
3599 Err(ObjectStoreError::AlreadyExists { .. })
3600 | Err(ObjectStoreError::Precondition { .. }) => {
3601 return Err(lance_core::Error::from(
3602 NamespaceError::ConcurrentModification {
3603 message: format!(
3604 "Version {} already exists for table at '{}'",
3605 version, table_uri
3606 ),
3607 },
3608 ));
3609 }
3610 Err(e) => {
3611 return Err(lance_core::Error::from(NamespaceError::Internal {
3612 message: format!(
3613 "Failed to create version {} for table at '{}': {}",
3614 version, table_uri, e
3615 ),
3616 }));
3617 }
3618 }
3619
3620 let final_meta = self
3621 .object_store
3622 .inner
3623 .head(&final_path)
3624 .await
3625 .map_err(|e| {
3626 lance_core::Error::from(NamespaceError::Internal {
3627 message: format!(
3628 "Failed to stat created version {} for table at '{}': {}",
3629 version, table_uri, e
3630 ),
3631 })
3632 })?;
3633 let manifest_size = final_meta.size as i64;
3634
3635 if let Err(e) = self.object_store.inner.delete(&staging_path).await {
3637 log::warn!(
3638 "Failed to delete staging manifest at '{}': {:?}",
3639 staging_path,
3640 e
3641 );
3642 }
3643
3644 Ok(CreateTableVersionResponse {
3645 transaction_id: None,
3646 version: Some(Box::new(TableVersion {
3647 version: version as i64,
3648 manifest_path: final_path.to_string(),
3649 manifest_size: Some(manifest_size),
3650 e_tag: final_meta.e_tag,
3651 timestamp_millis: None,
3652 metadata: None,
3653 })),
3654 })
3655 }
3656
3657 async fn describe_table_version(
3658 &self,
3659 request: DescribeTableVersionRequest,
3660 ) -> Result<DescribeTableVersionResponse> {
3661 self.record_op("describe_table_version");
3662 let branch = Self::normalized_branch(request.branch.as_deref())?;
3663 let table_uri = self.resolve_table_location(&request.id).await?;
3664 let table_uri = match branch {
3665 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3666 None => table_uri,
3667 };
3668 let versions = self
3669 .list_table_versions_from_storage(&table_uri, true, None)
3670 .await?;
3671 let table_version = if let Some(requested_version) = request.version {
3672 versions
3673 .into_iter()
3674 .find(|version| version.version == requested_version)
3675 .ok_or_else(|| {
3676 lance_core::Error::from(NamespaceError::TableVersionNotFound {
3677 message: format!(
3678 "version {} for table {}",
3679 requested_version,
3680 Self::format_table_id_from_request(&request.id)
3681 ),
3682 })
3683 })?
3684 } else {
3685 versions.into_iter().next().ok_or_else(|| {
3686 lance_core::Error::from(NamespaceError::TableVersionNotFound {
3687 message: format!(
3688 "latest version for table {}",
3689 Self::format_table_id_from_request(&request.id)
3690 ),
3691 })
3692 })?
3693 };
3694
3695 Ok(DescribeTableVersionResponse {
3696 version: Box::new(table_version),
3697 })
3698 }
3699
3700 async fn batch_delete_table_versions(
3701 &self,
3702 request: BatchDeleteTableVersionsRequest,
3703 ) -> Result<BatchDeleteTableVersionsResponse> {
3704 self.record_op("batch_delete_table_versions");
3705 let branch = Self::normalized_branch(request.branch.as_deref())?;
3706 let ranges: Vec<(i64, i64)> = request
3709 .ranges
3710 .iter()
3711 .map(|r| (r.start_version, r.end_version))
3712 .collect();
3713
3714 const MAX_VERSIONS_PER_REQUEST: i128 = 1_000_000;
3718 let requested: i128 = ranges
3719 .iter()
3720 .map(|(s, e)| {
3721 if *e < 0 {
3722 0
3723 } else {
3724 (*e as i128 - *s as i128).max(0)
3725 }
3726 })
3727 .sum();
3728 if requested > MAX_VERSIONS_PER_REQUEST {
3729 return Err(NamespaceError::InvalidInput {
3730 message: format!(
3731 "batch_delete requested {} versions; limit is {}",
3732 requested, MAX_VERSIONS_PER_REQUEST
3733 ),
3734 }
3735 .into());
3736 }
3737
3738 let table_entries = vec![TableDeleteEntry {
3739 table_id: request.id.clone(),
3740 ranges,
3741 }];
3742
3743 let total_deleted_count = self
3744 .delete_physical_version_files(&table_entries, branch)
3745 .await?;
3746
3747 Ok(BatchDeleteTableVersionsResponse {
3748 deleted_count: Some(total_deleted_count),
3749 transaction_id: None,
3750 })
3751 }
3752
3753 async fn create_table_index(
3754 &self,
3755 request: CreateTableIndexRequest,
3756 ) -> Result<CreateTableIndexResponse> {
3757 self.record_op("create_table_index");
3758 let table_uri = self.resolve_table_location(&request.id).await?;
3759 let mut dataset = self
3760 .load_dataset(&table_uri, None, "create_table_index")
3761 .await?;
3762 let index_request = Self::build_index_params(&request)?;
3763
3764 dataset
3765 .create_index(
3766 &[request.column.as_str()],
3767 index_request.index_type(),
3768 request.name.clone(),
3769 index_request.params(),
3770 false,
3771 )
3772 .await
3773 .map_err(|e| {
3774 let err_msg = format!("{}", e);
3775 let ns_err = if err_msg.contains("already exists") {
3776 NamespaceError::TableIndexAlreadyExists {
3777 message: format!(
3778 "Index '{}' already exists on table '{}': {:?}",
3779 request.name.as_deref().unwrap_or("<auto-generated>"),
3780 table_uri,
3781 e
3782 ),
3783 }
3784 } else if err_msg.contains("not found") || err_msg.contains("does not exist") {
3785 NamespaceError::TableColumnNotFound {
3786 message: format!(
3787 "Column '{}' not found for table '{}': {:?}",
3788 request.column, table_uri, e
3789 ),
3790 }
3791 } else {
3792 NamespaceError::Internal {
3793 message: format!(
3794 "Failed to create {} index '{}' on column '{}' for table '{}': {:?}",
3795 request.index_type,
3796 request.name.as_deref().unwrap_or("<auto-generated>"),
3797 request.column,
3798 table_uri,
3799 e
3800 ),
3801 }
3802 };
3803 lance_core::Error::from(ns_err)
3804 })?;
3805
3806 let transaction_id = dataset
3807 .read_transaction()
3808 .await
3809 .map_err(|e| {
3810 lance_core::Error::from(NamespaceError::Internal {
3811 message: format!(
3812 "Failed to read committed transaction after creating index on '{}': {}",
3813 table_uri, e
3814 ),
3815 })
3816 })?
3817 .map(|transaction| transaction.uuid);
3818
3819 Ok(CreateTableIndexResponse { transaction_id })
3820 }
3821
3822 async fn list_table_indices(
3823 &self,
3824 request: ListTableIndicesRequest,
3825 ) -> Result<ListTableIndicesResponse> {
3826 self.record_op("list_table_indices");
3827 let table_uri = self.resolve_table_location(&request.id).await?;
3828 let dataset = self
3829 .load_dataset(&table_uri, request.version, "list_table_indices")
3830 .await?;
3831 let total_rows = dataset.count_rows(None).await.map_err(|e| {
3832 lance_core::Error::from(NamespaceError::Internal {
3833 message: format!("Failed to count rows for table '{}': {:?}", table_uri, e),
3834 })
3835 })? as u64;
3836 let mut indices = dataset
3837 .describe_indices(None)
3838 .await
3839 .map_err(|e| {
3840 lance_core::Error::from(NamespaceError::Internal {
3841 message: format!("Failed to describe table indices for '{}': {:?}", table_uri, e),
3842 })
3843 })?
3844 .into_iter()
3845 .filter(|description| {
3846 description
3847 .metadata()
3848 .first()
3849 .map(|metadata| !is_system_index(metadata))
3850 .unwrap_or(false)
3851 })
3852 .map(|description| {
3853 let columns = description
3854 .field_ids()
3855 .iter()
3856 .map(|field_id| {
3857 dataset
3858 .schema()
3859 .field_path(i32::try_from(*field_id).map_err(|e| {
3860 lance_core::Error::from(NamespaceError::Internal {
3861 message: format!(
3862 "Field id {} does not fit in i32 for table '{}': {}",
3863 field_id, table_uri, e
3864 ),
3865 })
3866 })?)
3867 .map_err(|e| {
3868 lance_core::Error::from(NamespaceError::Internal {
3869 message: format!(
3870 "Failed to resolve field path for field_id {} in table '{}': {}",
3871 field_id, table_uri, e
3872 ),
3873 })
3874 })
3875 })
3876 .collect::<Result<Vec<_>>>()?;
3877
3878 let segments = description.segments();
3879 let created_at = segments
3880 .iter()
3881 .filter_map(|segment| segment.created_at)
3882 .min()
3883 .map(|ts| ts.to_rfc3339());
3884
3885 #[allow(clippy::needless_update)]
3888 let content = IndexContent {
3889 index_name: description.name().to_string(),
3890 index_uuid: description.metadata()[0].uuid.to_string(),
3891 columns,
3892 status: "SUCCEEDED".to_string(),
3893 index_type: Some(description.index_type().to_string()),
3894 type_url: Some(description.type_url().to_string()),
3895 num_indexed_rows: Some(description.rows_indexed() as i64),
3896 num_unindexed_rows: Some(
3897 total_rows.saturating_sub(description.rows_indexed()) as i64,
3898 ),
3899 size_bytes: description.total_size_bytes().map(|size| size as i64),
3900 num_segments: Some(segments.len() as i32),
3901 created_at,
3902 index_version: segments.first().map(|segment| segment.index_version),
3903 index_details: description.details().ok(),
3904 ..Default::default()
3905 };
3906 Ok(content)
3907 })
3908 .collect::<Result<Vec<_>>>()?;
3909
3910 let page_token = Self::paginate_indices(&mut indices, request.page_token, request.limit);
3911 Ok(ListTableIndicesResponse {
3912 indexes: indices,
3913 page_token,
3914 })
3915 }
3916
3917 async fn describe_table_index_stats(
3918 &self,
3919 request: DescribeTableIndexStatsRequest,
3920 ) -> Result<DescribeTableIndexStatsResponse> {
3921 self.record_op("describe_table_index_stats");
3922 let table_uri = self.resolve_table_location(&request.id).await?;
3923 let dataset = self
3924 .load_dataset(&table_uri, request.version, "describe_table_index_stats")
3925 .await?;
3926 let index_name = request.index_name.as_deref().ok_or_else(|| {
3927 lance_core::Error::from(NamespaceError::InvalidInput {
3928 message: "Index name is required for describe_table_index_stats".to_string(),
3929 })
3930 })?;
3931 let metadatas = dataset
3932 .load_indices_by_name(index_name)
3933 .await
3934 .map_err(|e| {
3935 lance_core::Error::from(NamespaceError::TableIndexNotFound {
3936 message: format!(
3937 "Failed to load index '{}' metadata for table '{}': {}",
3938 index_name, table_uri, e
3939 ),
3940 })
3941 })?;
3942 if metadatas.first().is_some_and(is_system_index) {
3943 return Err(NamespaceError::Unsupported {
3944 message: format!("System index '{}' is not exposed by this API", index_name),
3945 }
3946 .into());
3947 }
3948
3949 let stats = <Dataset as DatasetIndexExt>::index_statistics(&dataset, index_name)
3950 .await
3951 .map_err(|e| {
3952 lance_core::Error::from(NamespaceError::TableIndexNotFound {
3953 message: format!(
3954 "Failed to describe index statistics for '{}' on table '{}': {}",
3955 index_name, table_uri, e
3956 ),
3957 })
3958 })?;
3959 let stats: serde_json::Value = serde_json::from_str(&stats).map_err(|e| {
3960 lance_core::Error::from(NamespaceError::Internal {
3961 message: format!(
3962 "Failed to parse index statistics for '{}' on table '{}': {}",
3963 index_name, table_uri, e
3964 ),
3965 })
3966 })?;
3967
3968 Ok(Self::describe_table_index_stats_response(&stats))
3969 }
3970
3971 async fn describe_transaction(
3972 &self,
3973 request: DescribeTransactionRequest,
3974 ) -> Result<DescribeTransactionResponse> {
3975 self.record_op("describe_transaction");
3976 let mut request_id = request.id.ok_or_else(|| {
3977 lance_core::Error::from(NamespaceError::InvalidInput {
3978 message: "Transaction id must include table id and transaction identifier"
3979 .to_string(),
3980 })
3981 })?;
3982 if request_id.len() < 2 {
3983 return Err(NamespaceError::InvalidInput {
3984 message: format!(
3985 "Transaction request id must include table id and transaction identifier, got {:?}",
3986 request_id
3987 ),
3988 }
3989 .into());
3990 }
3991
3992 let id = request_id.pop().expect("request_id len checked above");
3993 let table_id = Some(request_id);
3994 let table_uri = self.resolve_table_location(&table_id).await?;
3995 let dataset = self
3996 .load_dataset(&table_uri, None, "describe_transaction")
3997 .await?;
3998 let (version, transaction) = self.find_transaction(&dataset, &id).await?;
3999
4000 let sidecar = self
4003 .load_transaction_alteration(&table_uri, &transaction.uuid)
4004 .await?;
4005
4006 Ok(Self::transaction_response(version, &transaction, sidecar))
4007 }
4008
4009 async fn alter_transaction(
4010 &self,
4011 request: AlterTransactionRequest,
4012 ) -> Result<AlterTransactionResponse> {
4013 self.record_op("alter_transaction");
4014
4015 let mut request_id = request.id.ok_or_else(|| {
4017 lance_core::Error::from(NamespaceError::InvalidInput {
4018 message: "Transaction id must include table id and transaction identifier"
4019 .to_string(),
4020 })
4021 })?;
4022 if request_id.len() < 2 {
4023 return Err(NamespaceError::InvalidInput {
4024 message: format!(
4025 "Transaction request id must include table id and transaction identifier, got {:?}",
4026 request_id
4027 ),
4028 }
4029 .into());
4030 }
4031
4032 let txn_id = request_id.pop().expect("request_id len checked above");
4033 let table_id = Some(request_id);
4034 let table_uri = self.resolve_table_location(&table_id).await?;
4035 let dataset = self
4036 .load_dataset(&table_uri, None, "alter_transaction")
4037 .await?;
4038 let (version, transaction) = self.find_transaction(&dataset, &txn_id).await?;
4039
4040 const RESERVED_KEYS: &[&str] = &["uuid", "version", "read_version", "operation", "tag"];
4044 let is_reserved = |key: &str| RESERVED_KEYS.contains(&key);
4045
4046 let mut sidecar = self
4049 .load_transaction_alteration(&table_uri, &transaction.uuid)
4050 .await?
4051 .unwrap_or_default();
4052
4053 for action in &request.actions {
4054 if let Some(ref set_status) = action.set_status_action
4055 && let Some(ref status) = set_status.status
4056 {
4057 let normalized = status.to_lowercase().replace('_', "");
4059 match normalized.as_str() {
4060 "queued" | "running" | "succeeded" | "failed" | "canceled" => {
4061 sidecar.status = Some(status.clone());
4062 }
4063 _ => {
4064 return Err(NamespaceError::InvalidInput {
4065 message: format!(
4066 "Invalid transaction status '{}'. Valid values are: Queued, Running, Succeeded, Failed, Canceled",
4067 status
4068 ),
4069 }
4070 .into());
4071 }
4072 }
4073 }
4074
4075 if let Some(ref set_property) = action.set_property_action
4076 && let (Some(key), Some(value)) = (&set_property.key, &set_property.value)
4077 {
4078 if is_reserved(key) {
4079 return Err(NamespaceError::InvalidInput {
4080 message: format!("Property '{}' is reserved and cannot be modified", key),
4081 }
4082 .into());
4083 }
4084 let mode = set_property
4085 .mode
4086 .as_deref()
4087 .unwrap_or("Overwrite")
4088 .to_lowercase();
4089 match mode.as_str() {
4090 "overwrite" => {
4091 sidecar.properties.insert(key.clone(), value.clone());
4092 }
4093 "fail" => {
4094 let exists = sidecar.properties.contains_key(key)
4097 || transaction
4098 .transaction_properties
4099 .as_ref()
4100 .is_some_and(|props| props.contains_key(key));
4101 if exists {
4102 return Err(NamespaceError::ConcurrentModification {
4103 message: format!(
4104 "Property '{}' already exists and mode is 'Fail'",
4105 key
4106 ),
4107 }
4108 .into());
4109 }
4110 sidecar.properties.insert(key.clone(), value.clone());
4111 }
4112 "skip" => {
4113 let exists = sidecar.properties.contains_key(key)
4114 || transaction
4115 .transaction_properties
4116 .as_ref()
4117 .is_some_and(|props| props.contains_key(key));
4118 if !exists {
4119 sidecar.properties.insert(key.clone(), value.clone());
4120 }
4121 }
4122 _ => {
4123 return Err(NamespaceError::InvalidInput {
4124 message: format!(
4125 "Invalid set_property mode '{}'. Valid values are: Overwrite, Fail, Skip",
4126 mode
4127 ),
4128 }
4129 .into());
4130 }
4131 }
4132 }
4133
4134 if let Some(ref unset_property) = action.unset_property_action
4135 && let Some(ref key) = unset_property.key
4136 {
4137 if is_reserved(key) {
4138 return Err(NamespaceError::InvalidInput {
4139 message: format!("Property '{}' is reserved and cannot be modified", key),
4140 }
4141 .into());
4142 }
4143 let mode = unset_property
4144 .mode
4145 .as_deref()
4146 .unwrap_or("Skip")
4147 .to_lowercase();
4148 let exists_in_transaction = transaction
4149 .transaction_properties
4150 .as_ref()
4151 .is_some_and(|props| props.contains_key(key));
4152 match mode.as_str() {
4153 "skip" => {
4154 sidecar.properties.remove(key);
4155 if exists_in_transaction {
4156 sidecar.removed_properties.insert(key.clone());
4159 }
4160 }
4161 "fail" => {
4162 if !sidecar.properties.contains_key(key) && !exists_in_transaction {
4163 return Err(NamespaceError::InvalidInput {
4164 message: format!(
4165 "Property '{}' does not exist and mode is 'Fail'",
4166 key
4167 ),
4168 }
4169 .into());
4170 }
4171 sidecar.properties.remove(key);
4172 if exists_in_transaction {
4173 sidecar.removed_properties.insert(key.clone());
4174 }
4175 }
4176 _ => {
4177 return Err(NamespaceError::InvalidInput {
4178 message: format!(
4179 "Invalid unset_property mode '{}'. Valid values are: Skip, Fail",
4180 mode
4181 ),
4182 }
4183 .into());
4184 }
4185 }
4186 }
4187 }
4188
4189 self.save_transaction_alteration(&table_uri, &transaction.uuid, &sidecar)
4193 .await?;
4194
4195 let final_status = sidecar
4198 .status
4199 .clone()
4200 .unwrap_or_else(|| "SUCCEEDED".to_string());
4201 let response = Self::transaction_response(version, &transaction, Some(sidecar));
4202 Ok(AlterTransactionResponse {
4203 status: final_status,
4204 properties: response.properties,
4205 })
4206 }
4207
4208 async fn create_table_scalar_index(
4209 &self,
4210 request: CreateTableIndexRequest,
4211 ) -> Result<CreateTableScalarIndexResponse> {
4212 self.record_op("create_table_scalar_index");
4213 let index_type = Self::parse_index_type(&request.index_type)?;
4214 if !index_type.is_scalar() {
4215 return Err(NamespaceError::InvalidInput {
4216 message: format!(
4217 "create_table_scalar_index only supports scalar index types, got {}",
4218 request.index_type
4219 ),
4220 }
4221 .into());
4222 }
4223
4224 let response = self.create_table_index(request).await?;
4225 Ok(CreateTableScalarIndexResponse {
4226 transaction_id: response.transaction_id,
4227 })
4228 }
4229
4230 async fn drop_table_index(
4231 &self,
4232 request: DropTableIndexRequest,
4233 ) -> Result<DropTableIndexResponse> {
4234 self.record_op("drop_table_index");
4235 let table_uri = self.resolve_table_location(&request.id).await?;
4236 let index_name = request.index_name.as_deref().ok_or_else(|| {
4237 lance_core::Error::from(NamespaceError::InvalidInput {
4238 message: "Index name is required for drop_table_index".to_string(),
4239 })
4240 })?;
4241 let mut dataset = self
4242 .load_dataset(&table_uri, None, "drop_table_index")
4243 .await?;
4244 let metadatas = dataset
4245 .load_indices_by_name(index_name)
4246 .await
4247 .map_err(|e| {
4248 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4249 message: format!(
4250 "Failed to load index '{}' before dropping it from table '{}': {}",
4251 index_name, table_uri, e
4252 ),
4253 })
4254 })?;
4255 if metadatas.first().is_some_and(is_system_index) {
4256 return Err(NamespaceError::Unsupported {
4257 message: format!(
4258 "System index '{}' cannot be dropped via this API",
4259 index_name
4260 ),
4261 }
4262 .into());
4263 }
4264
4265 dataset.drop_index(index_name).await.map_err(|e| {
4266 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4267 message: format!(
4268 "Failed to drop index '{}' from table '{}': {}",
4269 index_name, table_uri, e
4270 ),
4271 })
4272 })?;
4273
4274 let transaction_id = dataset
4275 .read_transaction()
4276 .await
4277 .map_err(|e| {
4278 lance_core::Error::from(NamespaceError::Internal {
4279 message: format!(
4280 "Failed to read committed transaction after dropping index '{}' from '{}': {}",
4281 index_name, table_uri, e
4282 ),
4283 })
4284 })?
4285 .map(|transaction| transaction.uuid);
4286
4287 Ok(DropTableIndexResponse { transaction_id })
4288 }
4289
4290 async fn list_all_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
4291 let mut tables = self.list_directory_tables().await?;
4294 tables = self
4295 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
4296 .await?;
4297 Self::apply_pagination(&mut tables, request.page_token, request.limit);
4298 Ok(ListTablesResponse::new(tables))
4299 }
4300
4301 async fn restore_table(&self, request: RestoreTableRequest) -> Result<RestoreTableResponse> {
4302 let version = request.version;
4303 if version < 0 {
4304 return Err(Error::invalid_input_source(
4305 format!(
4306 "Table version for restore_table must be non-negative, got {}",
4307 version
4308 )
4309 .into(),
4310 ));
4311 }
4312
4313 let branch = Self::normalized_branch(request.branch.as_deref())?;
4314 let table_uri = self.resolve_table_location(&request.id).await?;
4315 let mut dataset = match branch {
4316 Some(branch) => self.open_validated_branch(&table_uri, branch).await?,
4317 None => self.load_dataset(&table_uri, None, "restore_table").await?,
4318 };
4319
4320 dataset = dataset
4321 .checkout_version(version as u64)
4322 .await
4323 .map_err(|e| {
4324 Error::namespace_source(
4325 format!(
4326 "Failed to checkout version {} for restore at '{}': {}",
4327 version, table_uri, e
4328 )
4329 .into(),
4330 )
4331 })?;
4332
4333 dataset.restore().await.map_err(|e| {
4334 Error::namespace_source(
4335 format!(
4336 "Failed to restore table at '{}' to version {}: {}",
4337 table_uri, version, e
4338 )
4339 .into(),
4340 )
4341 })?;
4342
4343 let transaction_id = dataset
4344 .read_transaction()
4345 .await
4346 .map_err(|e| {
4347 Error::namespace_source(
4348 format!(
4349 "Failed to read transaction after restoring '{}': {}",
4350 table_uri, e
4351 )
4352 .into(),
4353 )
4354 })?
4355 .map(|t| t.uuid);
4356
4357 Ok(RestoreTableResponse { transaction_id })
4358 }
4359
4360 async fn update_table_schema_metadata(
4361 &self,
4362 request: UpdateTableSchemaMetadataRequest,
4363 ) -> Result<UpdateTableSchemaMetadataResponse> {
4364 let table_uri = self.resolve_table_location(&request.id).await?;
4365 let mut dataset = self
4366 .load_dataset(&table_uri, None, "update_table_schema_metadata")
4367 .await?;
4368
4369 let new_metadata = request.metadata.unwrap_or_default();
4370 let updated_metadata = dataset
4371 .update_schema_metadata(new_metadata.iter().map(|(k, v)| (k.as_str(), v.as_str())))
4372 .await
4373 .map_err(|e| {
4374 Error::namespace_source(
4375 format!(
4376 "Failed to update schema metadata for table at '{}': {}",
4377 table_uri, e
4378 )
4379 .into(),
4380 )
4381 })?;
4382
4383 let transaction_id = dataset
4384 .read_transaction()
4385 .await
4386 .map_err(|e| {
4387 Error::namespace_source(
4388 format!(
4389 "Failed to read transaction after updating metadata for '{}': {}",
4390 table_uri, e
4391 )
4392 .into(),
4393 )
4394 })?
4395 .map(|t| t.uuid);
4396
4397 Ok(UpdateTableSchemaMetadataResponse {
4398 metadata: Some(updated_metadata),
4399 transaction_id,
4400 })
4401 }
4402
4403 async fn get_table_stats(
4404 &self,
4405 request: GetTableStatsRequest,
4406 ) -> Result<GetTableStatsResponse> {
4407 let table_uri = self.resolve_table_location(&request.id).await?;
4408 let dataset = Arc::new(
4409 self.load_dataset(&table_uri, None, "get_table_stats")
4410 .await?,
4411 );
4412
4413 let data_stats = dataset.calculate_data_stats().await.map_err(|e| {
4415 Error::namespace_source(
4416 format!(
4417 "Failed to calculate data statistics for table at '{}': {}",
4418 table_uri, e
4419 )
4420 .into(),
4421 )
4422 })?;
4423 let total_bytes: i64 = data_stats
4424 .fields
4425 .iter()
4426 .map(|f| f.bytes_on_disk as i64)
4427 .sum();
4428
4429 let fragment_row_futures: Vec<_> = dataset
4431 .get_fragments()
4432 .into_iter()
4433 .map(|f| async move { f.physical_rows().await })
4434 .collect();
4435 let fragment_row_results = futures::future::join_all(fragment_row_futures).await;
4436 let mut fragment_row_counts: Vec<i64> = fragment_row_results
4437 .into_iter()
4438 .filter_map(|r| r.ok())
4439 .map(|r| r as i64)
4440 .collect();
4441
4442 let num_fragments = fragment_row_counts.len() as i64;
4443 let num_rows: i64 = fragment_row_counts.iter().sum();
4444
4445 const SMALL_FRAGMENT_THRESHOLD: i64 = 1024 * 1024;
4448 let num_small_fragments = fragment_row_counts
4449 .iter()
4450 .filter(|&&r| r < SMALL_FRAGMENT_THRESHOLD)
4451 .count() as i64;
4452
4453 fragment_row_counts.sort_unstable();
4455 let lengths = if fragment_row_counts.is_empty() {
4456 FragmentSummary::new(0, 0, 0, 0, 0, 0, 0)
4457 } else {
4458 let len = fragment_row_counts.len();
4459 let min = fragment_row_counts[0];
4460 let max = fragment_row_counts[len - 1];
4461 let mean = num_rows / num_fragments;
4462 let pct = |p: f64| fragment_row_counts[((len - 1) as f64 * p) as usize];
4463 FragmentSummary::new(min, max, mean, pct(0.25), pct(0.50), pct(0.75), pct(0.99))
4464 };
4465
4466 let indices = dataset.load_indices().await.map_err(|e| {
4468 Error::namespace_source(
4469 format!("Failed to load indices for table at '{}': {}", table_uri, e).into(),
4470 )
4471 })?;
4472 let num_indices = indices.iter().filter(|m| !is_system_index(m)).count() as i64;
4473
4474 let fragment_stats = FragmentStats::new(num_fragments, num_small_fragments, lengths);
4475 Ok(GetTableStatsResponse::new(
4476 total_bytes,
4477 num_rows,
4478 num_indices,
4479 fragment_stats,
4480 ))
4481 }
4482
4483 async fn explain_table_query_plan(
4484 &self,
4485 request: ExplainTableQueryPlanRequest,
4486 ) -> Result<String> {
4487 let table_uri = self.resolve_table_location(&request.id).await?;
4488 let dataset = self
4489 .load_dataset(
4490 &table_uri,
4491 request.query.version,
4492 "explain_table_query_plan",
4493 )
4494 .await?;
4495 let verbose = request.verbose.unwrap_or(false);
4496
4497 let mut scanner = dataset.scan();
4498 Self::apply_query_params_to_scanner(
4499 &mut scanner,
4500 request.query.filter.as_deref(),
4501 request.query.columns.as_deref(),
4502 request.query.vector_column.as_deref(),
4503 &request.query.vector,
4504 request.query.k,
4505 request.query.offset,
4506 request.query.prefilter,
4507 request.query.bypass_vector_index,
4508 request.query.nprobes,
4509 request.query.ef,
4510 request.query.refine_factor,
4511 request.query.distance_type.as_deref(),
4512 request.query.fast_search,
4513 request.query.with_row_id,
4514 request.query.lower_bound,
4515 request.query.upper_bound,
4516 "explain_table_query_plan",
4517 )?;
4518
4519 scanner.explain_plan(verbose).await.map_err(|e| {
4520 Error::namespace_source(
4521 format!(
4522 "Failed to explain query plan for table at '{}': {}",
4523 table_uri, e
4524 )
4525 .into(),
4526 )
4527 })
4528 }
4529
4530 async fn analyze_table_query_plan(
4531 &self,
4532 request: AnalyzeTableQueryPlanRequest,
4533 ) -> Result<String> {
4534 let table_uri = self.resolve_table_location(&request.id).await?;
4535 let dataset = self
4536 .load_dataset(&table_uri, request.version, "analyze_table_query_plan")
4537 .await?;
4538
4539 let mut scanner = dataset.scan();
4540 Self::apply_query_params_to_scanner(
4541 &mut scanner,
4542 request.filter.as_deref(),
4543 request.columns.as_deref(),
4544 request.vector_column.as_deref(),
4545 &request.vector,
4546 request.k,
4547 request.offset,
4548 request.prefilter,
4549 request.bypass_vector_index,
4550 request.nprobes,
4551 request.ef,
4552 request.refine_factor,
4553 request.distance_type.as_deref(),
4554 request.fast_search,
4555 request.with_row_id,
4556 request.lower_bound,
4557 request.upper_bound,
4558 "analyze_table_query_plan",
4559 )?;
4560
4561 scanner.analyze_plan().await.map_err(|e| {
4562 Error::namespace_source(
4563 format!(
4564 "Failed to analyze query plan for table at '{}': {}",
4565 table_uri, e
4566 )
4567 .into(),
4568 )
4569 })
4570 }
4571
4572 async fn count_table_rows(&self, request: CountTableRowsRequest) -> Result<i64> {
4573 self.record_op("count_table_rows");
4574 let table_uri = self.resolve_table_location(&request.id).await?;
4575 let dataset = self
4576 .load_dataset(&table_uri, request.version, "count_table_rows")
4577 .await?;
4578
4579 let count =
4580 dataset
4581 .count_rows(request.predicate)
4582 .await
4583 .map_err(|e| NamespaceError::Internal {
4584 message: format!("Failed to count rows for table at '{}': {:?}", table_uri, e),
4585 })?;
4586
4587 Ok(count as i64)
4588 }
4589
4590 async fn insert_into_table(
4591 &self,
4592 request: InsertIntoTableRequest,
4593 request_data: Bytes,
4594 ) -> Result<InsertIntoTableResponse> {
4595 self.record_op("insert_into_table");
4596 let table_uri = self.resolve_table_location(&request.id).await?;
4597 let (reader, _num_rows) =
4598 Self::ipc_reader_from_request_data(&request_data, "insert_into_table")?;
4599
4600 let mode = match request.mode.as_deref() {
4601 Some(m) if m.eq_ignore_ascii_case("overwrite") => WriteMode::Overwrite,
4602 Some(m) if m.eq_ignore_ascii_case("append") => WriteMode::Append,
4603 None => WriteMode::Append,
4604 Some(m) => {
4605 return Err(lance_namespace::error::NamespaceError::InvalidInput {
4606 message: format!(
4607 "Unsupported write mode '{}'. Supported modes are: 'append', 'overwrite'",
4608 m
4609 ),
4610 }
4611 .into());
4612 }
4613 };
4614
4615 if !self.table_uri_has_actual_manifests(&table_uri).await? {
4616 self.write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
4617 .await?;
4618 } else {
4619 self.write_reader_to_table(&table_uri, reader, mode, None)
4620 .await?;
4621 }
4622
4623 Ok(InsertIntoTableResponse {
4624 transaction_id: None,
4625 })
4626 }
4627
4628 async fn merge_insert_into_table(
4629 &self,
4630 request: MergeInsertIntoTableRequest,
4631 request_data: Bytes,
4632 ) -> Result<MergeInsertIntoTableResponse> {
4633 self.record_op("merge_insert_into_table");
4634 let table_uri = self.resolve_table_location(&request.id).await?;
4635 let on = request.on.as_ref().ok_or_else(|| {
4636 lance_core::Error::from(NamespaceError::InvalidInput {
4637 message: "'on' field is required for merge_insert_into_table".to_string(),
4638 })
4639 })?;
4640
4641 let table_has_manifests = self.table_uri_has_actual_manifests(&table_uri).await?;
4642 let (reader, num_rows) =
4643 Self::ipc_reader_from_request_data(&request_data, "merge_insert_into_table")?;
4644
4645 if !table_has_manifests {
4646 let dataset = self
4647 .write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
4648 .await?;
4649 let version = dataset.version().version as i64;
4650 return Ok(MergeInsertIntoTableResponse {
4651 transaction_id: None,
4652 num_updated_rows: Some(0),
4653 num_inserted_rows: Some(num_rows as i64),
4654 num_deleted_rows: Some(0),
4655 version: Some(version),
4656 });
4657 }
4658
4659 let dataset = Arc::new(
4660 self.load_dataset(&table_uri, None, "merge_insert_into_table")
4661 .await?,
4662 );
4663
4664 let mut merge_builder = MergeInsertBuilder::try_new(dataset.clone(), vec![on.clone()])
4665 .map_err(|e| {
4666 lance_core::Error::from(NamespaceError::InvalidInput {
4667 message: format!("Failed to create merge_insert_into_table builder: {}", e),
4668 })
4669 })?;
4670
4671 if let Some(filter) = request.when_matched_update_all_filt.as_deref() {
4672 let behavior = WhenMatched::update_if(dataset.as_ref(), filter).map_err(|e| {
4673 lance_core::Error::from(NamespaceError::InvalidInput {
4674 message: format!(
4675 "Invalid when_matched_update_all_filt for merge_insert_into_table: {}",
4676 e
4677 ),
4678 })
4679 })?;
4680 merge_builder.when_matched(behavior);
4681 } else if request.when_matched_update_all.unwrap_or(false) {
4682 merge_builder.when_matched(WhenMatched::UpdateAll);
4683 }
4684
4685 if matches!(request.when_not_matched_insert_all, Some(false)) {
4686 merge_builder.when_not_matched(WhenNotMatched::DoNothing);
4687 } else {
4688 merge_builder.when_not_matched(WhenNotMatched::InsertAll);
4689 }
4690
4691 if let Some(filter) = request.when_not_matched_by_source_delete_filt.as_deref() {
4692 let behavior = WhenNotMatchedBySource::delete_if(dataset.as_ref(), filter).map_err(|e| {
4693 lance_core::Error::from(NamespaceError::InvalidInput {
4694 message: format!(
4695 "Invalid when_not_matched_by_source_delete_filt for merge_insert_into_table: {}",
4696 e
4697 ),
4698 })
4699 })?;
4700 merge_builder.when_not_matched_by_source(behavior);
4701 } else if request.when_not_matched_by_source_delete.unwrap_or(false) {
4702 merge_builder.when_not_matched_by_source(WhenNotMatchedBySource::Delete);
4703 }
4704
4705 if let Some(use_index) = request.use_index {
4706 merge_builder.use_index(use_index);
4707 }
4708
4709 let (dataset, stats) = merge_builder
4710 .try_build()
4711 .map_err(|e| {
4712 lance_core::Error::from(NamespaceError::InvalidInput {
4713 message: format!("Failed to build merge_insert_into_table job: {}", e),
4714 })
4715 })?
4716 .execute_reader(reader)
4717 .await
4718 .map_err(|e| Self::map_mutation_error(e, "merge_insert_into_table", &table_uri))?;
4719
4720 Ok(MergeInsertIntoTableResponse {
4721 transaction_id: None,
4722 num_updated_rows: Some(stats.num_updated_rows as i64),
4723 num_inserted_rows: Some(stats.num_inserted_rows as i64),
4724 num_deleted_rows: Some(stats.num_deleted_rows as i64),
4725 version: Some(dataset.version().version as i64),
4726 })
4727 }
4728
4729 async fn update_table(&self, request: UpdateTableRequest) -> Result<UpdateTableResponse> {
4730 self.record_op("update_table");
4731
4732 if request.updates.is_empty() {
4733 return Err(NamespaceError::InvalidInput {
4734 message: "update_table requires at least one [column, expression] pair".to_string(),
4735 }
4736 .into());
4737 }
4738
4739 let mut seen_columns: HashMap<String, ()> = HashMap::with_capacity(request.updates.len());
4742 for (idx, pair) in request.updates.iter().enumerate() {
4743 if pair.len() != 2 {
4744 return Err(NamespaceError::InvalidInput {
4745 message: format!(
4746 "update_table updates[{}] must be a [column, expression] pair, got {} elements",
4747 idx,
4748 pair.len()
4749 ),
4750 }
4751 .into());
4752 }
4753 let column = &pair[0];
4754 if column.trim().is_empty() {
4755 return Err(NamespaceError::InvalidInput {
4756 message: format!("update_table updates[{}] has an empty column name", idx),
4757 }
4758 .into());
4759 }
4760 if seen_columns.insert(column.clone(), ()).is_some() {
4761 return Err(NamespaceError::InvalidInput {
4762 message: format!(
4763 "update_table cannot update column '{}' more than once",
4764 column
4765 ),
4766 }
4767 .into());
4768 }
4769 }
4770
4771 let table_uri = self.resolve_table_location(&request.id).await?;
4772 let dataset = Arc::new(self.load_dataset(&table_uri, None, "update_table").await?);
4773
4774 let mut builder = UpdateBuilder::new(dataset);
4775 for pair in &request.updates {
4776 builder = builder.set(&pair[0], &pair[1]).map_err(|e| {
4778 lance_core::Error::from(NamespaceError::InvalidInput {
4779 message: format!("Invalid update expression for column '{}': {}", pair[0], e),
4780 })
4781 })?;
4782 }
4783 if let Some(predicate) = request.predicate.as_deref()
4784 && !predicate.trim().is_empty()
4785 {
4786 builder = builder.update_where(predicate).map_err(|e| {
4787 lance_core::Error::from(NamespaceError::InvalidInput {
4788 message: format!("Invalid update_table predicate '{}': {}", predicate, e),
4789 })
4790 })?;
4791 }
4792
4793 let job = builder.build().map_err(|e| {
4794 lance_core::Error::from(NamespaceError::InvalidInput {
4795 message: format!("Failed to build update_table job: {}", e),
4796 })
4797 })?;
4798
4799 let result = job
4800 .execute()
4801 .await
4802 .map_err(|e| Self::map_mutation_error(e, "update_table", &table_uri))?;
4803
4804 let version = result.new_dataset.version().version as i64;
4805 Ok(UpdateTableResponse {
4806 transaction_id: None,
4807 updated_rows: result.rows_updated as i64,
4808 version,
4809 properties: None,
4810 })
4811 }
4812
4813 async fn delete_from_table(
4814 &self,
4815 request: DeleteFromTableRequest,
4816 ) -> Result<DeleteFromTableResponse> {
4817 self.record_op("delete_from_table");
4818
4819 if request.predicate.trim().is_empty() {
4820 return Err(NamespaceError::InvalidInput {
4821 message: "delete_from_table requires a non-empty predicate".to_string(),
4822 }
4823 .into());
4824 }
4825
4826 let table_uri = self.resolve_table_location(&request.id).await?;
4827 let mut dataset = self
4828 .load_dataset(&table_uri, None, "delete_from_table")
4829 .await?;
4830
4831 let result = dataset
4832 .delete(&request.predicate)
4833 .await
4834 .map_err(|e| Self::map_mutation_error(e, "delete_from_table", &table_uri))?;
4835
4836 Ok(DeleteFromTableResponse {
4837 transaction_id: None,
4838 version: Some(result.new_dataset.version().version as i64),
4839 })
4840 }
4841
4842 async fn query_table(&self, request: QueryTableRequest) -> Result<Bytes> {
4843 use arrow::ipc::writer::FileWriter;
4844
4845 self.record_op("query_table");
4846 let table_uri = self.resolve_table_location(&request.id).await?;
4847 let dataset = self
4848 .load_dataset(&table_uri, request.version, "query_table")
4849 .await?;
4850
4851 let mut scanner = dataset.scan();
4853
4854 let has_vector_query = request
4857 .vector
4858 .single_vector
4859 .as_ref()
4860 .map(|sv| !sv.is_empty())
4861 .unwrap_or(false)
4862 || request
4863 .vector
4864 .multi_vector
4865 .as_ref()
4866 .map(|mv| !mv.is_empty())
4867 .unwrap_or(false);
4868
4869 if let Some(prefilter) = request.prefilter {
4871 scanner.prefilter(prefilter);
4872 }
4873
4874 if has_vector_query {
4876 let vector_column = request.vector_column.as_deref().unwrap_or("vector");
4877
4878 let query_vector: Vec<f32> = request
4880 .vector
4881 .single_vector
4882 .clone()
4883 .or_else(|| {
4884 request
4885 .vector
4886 .multi_vector
4887 .as_ref()
4888 .and_then(|mv| mv.first().cloned())
4889 })
4890 .unwrap_or_default();
4891
4892 if !query_vector.is_empty() {
4893 let k = if request.k > 0 {
4894 request.k as usize
4895 } else {
4896 10
4897 };
4898 let query_array = Float32Array::from(query_vector);
4899 scanner
4900 .nearest(vector_column, &query_array, k)
4901 .map_err(|e| NamespaceError::InvalidInput {
4902 message: format!("Invalid vector search: {:?}", e),
4903 })?;
4904
4905 if let Some(ref distance_type) = request.distance_type {
4907 let metric = match distance_type.to_lowercase().as_str() {
4908 "l2" | "euclidean" => MetricType::L2,
4909 "cosine" => MetricType::Cosine,
4910 "dot" | "inner_product" => MetricType::Dot,
4911 "hamming" => MetricType::Hamming,
4912 _ => {
4913 return Err(NamespaceError::InvalidInput {
4914 message: format!("Unknown distance type: {}", distance_type),
4915 }
4916 .into());
4917 }
4918 };
4919 scanner.distance_metric(metric);
4920 }
4921
4922 if let Some(nprobes) = request.nprobes {
4924 scanner.minimum_nprobes(nprobes as usize);
4925 }
4926
4927 if let Some(ef) = request.ef {
4929 scanner.ef(ef as usize);
4930 }
4931
4932 if let Some(refine_factor) = request.refine_factor {
4934 scanner.refine(refine_factor as u32);
4935 }
4936
4937 if request.lower_bound.is_some() || request.upper_bound.is_some() {
4939 scanner.distance_range(request.lower_bound, request.upper_bound);
4940 }
4941
4942 if let Some(bypass) = request.bypass_vector_index {
4944 scanner.use_index(!bypass);
4945 }
4946
4947 if request.fast_search == Some(true) {
4949 scanner.fast_search();
4950 }
4951 }
4952 }
4953
4954 if let Some(ref fts_query) = request.full_text_query {
4956 if let Some(ref string_query) = fts_query.string_query {
4958 let mut fts = FullTextSearchQuery::new(string_query.query.clone());
4959
4960 if let Some(ref columns) = string_query.columns
4962 && !columns.is_empty()
4963 {
4964 fts = fts
4965 .with_columns(columns)
4966 .map_err(|e| NamespaceError::InvalidInput {
4967 message: format!("Invalid FTS columns: {:?}", e),
4968 })?;
4969 }
4970
4971 scanner
4972 .full_text_search(fts)
4973 .map_err(|e| NamespaceError::InvalidInput {
4974 message: format!("Invalid full text search: {:?}", e),
4975 })?;
4976 }
4977 }
4980
4981 if let Some(ref columns) = request.columns {
4983 if let Some(ref column_names) = columns.column_names
4984 && !column_names.is_empty()
4985 {
4986 scanner
4987 .project(column_names)
4988 .map_err(|e| NamespaceError::InvalidInput {
4989 message: format!("Invalid column projection: {:?}", e),
4990 })?;
4991 } else if let Some(ref column_aliases) = columns.column_aliases
4992 && !column_aliases.is_empty()
4993 {
4994 let transform_pairs: Vec<(String, String)> = column_aliases
4996 .iter()
4997 .map(|(alias, sql)| (alias.clone(), sql.clone()))
4998 .collect();
4999 scanner
5000 .project_with_transform(
5001 &transform_pairs
5002 .iter()
5003 .map(|(a, s)| (a.as_str(), s.as_str()))
5004 .collect::<Vec<_>>(),
5005 )
5006 .map_err(|e| NamespaceError::InvalidInput {
5007 message: format!("Invalid column alias expression: {:?}", e),
5008 })?;
5009 }
5010 }
5011
5012 if let Some(ref filter) = request.filter
5014 && !filter.is_empty()
5015 {
5016 scanner
5017 .filter(filter)
5018 .map_err(|e| NamespaceError::InvalidInput {
5019 message: format!("Invalid filter expression: {:?}", e),
5020 })?;
5021 }
5022
5023 if request.with_row_id == Some(true) {
5025 scanner.with_row_id();
5026 }
5027
5028 if !has_vector_query && request.k > 0 {
5032 let offset = request.offset.map(|o| o as i64);
5033 scanner.limit(Some(request.k as i64), offset).map_err(|e| {
5034 NamespaceError::InvalidInput {
5035 message: format!("Invalid limit/offset: {:?}", e),
5036 }
5037 })?;
5038 } else if has_vector_query && request.offset.is_some() {
5039 let offset = request.offset.map(|o| o as i64);
5041 scanner
5042 .limit(None, offset)
5043 .map_err(|e| NamespaceError::InvalidInput {
5044 message: format!("Invalid offset: {:?}", e),
5045 })?;
5046 }
5047
5048 let batch = scanner
5050 .try_into_batch()
5051 .await
5052 .map_err(|e| NamespaceError::Internal {
5053 message: format!("Failed to execute query: {:?}", e),
5054 })?;
5055
5056 let schema = batch.schema();
5058 let mut buffer = Vec::new();
5059 {
5060 let mut writer = FileWriter::try_new(&mut buffer, &schema).map_err(|e| {
5061 NamespaceError::Internal {
5062 message: format!("Failed to create IPC writer: {:?}", e),
5063 }
5064 })?;
5065 writer.write(&batch).map_err(|e| NamespaceError::Internal {
5066 message: format!("Failed to write batch to IPC: {:?}", e),
5067 })?;
5068 writer.finish().map_err(|e| NamespaceError::Internal {
5069 message: format!("Failed to finish IPC writer: {:?}", e),
5070 })?;
5071 }
5072
5073 Ok(Bytes::from(buffer))
5074 }
5075
5076 async fn list_table_tags(
5077 &self,
5078 request: ListTableTagsRequest,
5079 ) -> Result<ListTableTagsResponse> {
5080 self.record_op("list_table_tags");
5081 let table_uri = self.resolve_table_location(&request.id).await?;
5082 let dataset = self
5083 .load_dataset(&table_uri, None, "list_table_tags")
5084 .await?;
5085
5086 let raw_tags = dataset.tags().list().await.map_err(|e| {
5087 lance_core::Error::from(NamespaceError::Internal {
5088 message: format!("Failed to list tags for table at '{}': {}", table_uri, e),
5089 })
5090 })?;
5091
5092 let tags = raw_tags
5093 .into_iter()
5094 .map(|(name, contents)| {
5095 let mut tag_model =
5096 ModelTagContents::new(contents.version as i64, contents.manifest_size as i64);
5097 tag_model.branch = contents.branch;
5098 (name, tag_model)
5099 })
5100 .collect();
5101
5102 Ok(ListTableTagsResponse {
5103 tags,
5104 page_token: None,
5105 })
5106 }
5107
5108 async fn get_table_tag_version(
5109 &self,
5110 request: GetTableTagVersionRequest,
5111 ) -> Result<GetTableTagVersionResponse> {
5112 self.record_op("get_table_tag_version");
5113 if request.tag.is_empty() {
5114 return Err(NamespaceError::InvalidInput {
5115 message: "tag name must not be empty for get_table_tag_version".to_string(),
5116 }
5117 .into());
5118 }
5119
5120 let table_uri = self.resolve_table_location(&request.id).await?;
5121 let dataset = self
5122 .load_dataset(&table_uri, None, "get_table_tag_version")
5123 .await?;
5124
5125 let contents = dataset
5126 .tags()
5127 .get(&request.tag)
5128 .await
5129 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5130
5131 Ok(GetTableTagVersionResponse {
5132 version: contents.version as i64,
5133 branch: contents.branch,
5134 })
5135 }
5136
5137 async fn create_table_tag(
5138 &self,
5139 request: CreateTableTagRequest,
5140 ) -> Result<CreateTableTagResponse> {
5141 self.record_op("create_table_tag");
5142 if request.tag.is_empty() {
5143 return Err(NamespaceError::InvalidInput {
5144 message: "tag name must not be empty for create_table_tag".to_string(),
5145 }
5146 .into());
5147 }
5148 if request.version <= 0 {
5149 return Err(NamespaceError::InvalidInput {
5150 message: format!(
5151 "tag version must be a positive integer, got {} for create_table_tag",
5152 request.version
5153 ),
5154 }
5155 .into());
5156 }
5157
5158 let table_uri = self.resolve_table_location(&request.id).await?;
5159 let dataset = self
5160 .load_dataset(&table_uri, None, "create_table_tag")
5161 .await?;
5162
5163 dataset
5164 .tags()
5165 .create(&request.tag, request.version as u64)
5166 .await
5167 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5168
5169 Ok(CreateTableTagResponse {
5170 transaction_id: None,
5171 })
5172 }
5173
5174 async fn delete_table_tag(
5175 &self,
5176 request: DeleteTableTagRequest,
5177 ) -> Result<DeleteTableTagResponse> {
5178 self.record_op("delete_table_tag");
5179 if request.tag.is_empty() {
5180 return Err(NamespaceError::InvalidInput {
5181 message: "tag name must not be empty for delete_table_tag".to_string(),
5182 }
5183 .into());
5184 }
5185
5186 let table_uri = self.resolve_table_location(&request.id).await?;
5187 let dataset = self
5188 .load_dataset(&table_uri, None, "delete_table_tag")
5189 .await?;
5190
5191 dataset
5192 .tags()
5193 .delete(&request.tag)
5194 .await
5195 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5196
5197 Ok(DeleteTableTagResponse {
5198 transaction_id: None,
5199 })
5200 }
5201
5202 async fn update_table_tag(
5203 &self,
5204 request: UpdateTableTagRequest,
5205 ) -> Result<UpdateTableTagResponse> {
5206 self.record_op("update_table_tag");
5207 if request.tag.is_empty() {
5208 return Err(NamespaceError::InvalidInput {
5209 message: "tag name must not be empty for update_table_tag".to_string(),
5210 }
5211 .into());
5212 }
5213 if request.version <= 0 {
5214 return Err(NamespaceError::InvalidInput {
5215 message: format!(
5216 "tag version must be a positive integer, got {} for update_table_tag",
5217 request.version
5218 ),
5219 }
5220 .into());
5221 }
5222
5223 let table_uri = self.resolve_table_location(&request.id).await?;
5224 let dataset = self
5225 .load_dataset(&table_uri, None, "update_table_tag")
5226 .await?;
5227
5228 dataset
5229 .tags()
5230 .update(&request.tag, request.version as u64)
5231 .await
5232 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5233
5234 Ok(UpdateTableTagResponse {
5235 transaction_id: None,
5236 })
5237 }
5238
5239 async fn create_table_branch(
5240 &self,
5241 request: CreateTableBranchRequest,
5242 ) -> Result<CreateTableBranchResponse> {
5243 self.record_op("create_table_branch");
5244 if request.name.is_empty() {
5245 return Err(NamespaceError::InvalidInput {
5246 message: "branch name must not be empty for create_table_branch".to_string(),
5247 }
5248 .into());
5249 }
5250 let from_version = match request.from_version {
5251 Some(v) if v <= 0 => {
5252 return Err(NamespaceError::InvalidInput {
5253 message: format!(
5254 "from_version must be a positive integer, got {} for create_table_branch",
5255 v
5256 ),
5257 }
5258 .into());
5259 }
5260 Some(v) => Some(v as u64),
5261 None => None,
5262 };
5263
5264 let table_uri = self.resolve_table_location(&request.id).await?;
5265 let mut dataset = self
5266 .load_dataset(&table_uri, None, "create_table_branch")
5267 .await?;
5268
5269 if dataset.branches().get(&request.name).await.is_ok() {
5274 return Err(NamespaceError::TableBranchAlreadyExists {
5275 message: format!("branch '{}' for table at '{}'", request.name, table_uri),
5276 }
5277 .into());
5278 }
5279
5280 dataset
5281 .create_branch(
5282 &request.name,
5283 (request.from_branch.as_deref(), from_version),
5284 None,
5285 )
5286 .await
5287 .map_err(|e| {
5288 if matches!(e, lance_core::Error::DatasetNotFound { .. }) {
5291 NamespaceError::InvalidInput {
5292 message: format!(
5293 "from_branch/from_version for branch '{}' refers to a source that does not exist: {}",
5294 request.name, e
5295 ),
5296 }
5297 .into()
5298 } else {
5299 Self::map_branch_error(e, &request.name, &table_uri)
5300 }
5301 })?;
5302
5303 Ok(CreateTableBranchResponse {
5304 transaction_id: None,
5305 })
5306 }
5307
5308 async fn list_table_branches(
5309 &self,
5310 request: ListTableBranchesRequest,
5311 ) -> Result<ListTableBranchesResponse> {
5312 self.record_op("list_table_branches");
5313 let table_uri = self.resolve_table_location(&request.id).await?;
5314 let dataset = self
5315 .load_dataset(&table_uri, None, "list_table_branches")
5316 .await?;
5317
5318 let raw_branches = dataset.list_branches().await.map_err(|e| {
5319 lance_core::Error::from(NamespaceError::Internal {
5320 message: format!(
5321 "Failed to list branches for table at '{}': {}",
5322 table_uri, e
5323 ),
5324 })
5325 })?;
5326
5327 let branches = raw_branches
5328 .into_iter()
5329 .map(|(name, contents)| {
5330 let mut branch_model = ModelBranchContents::new(
5333 contents.parent_version as i64,
5334 contents.create_at as i64,
5335 contents.manifest_size as i64,
5336 );
5337 branch_model.parent_branch = contents.parent_branch;
5338 branch_model.metadata = if contents.metadata.is_empty() {
5339 None
5340 } else {
5341 Some(contents.metadata)
5342 };
5343 (name, branch_model)
5344 })
5345 .collect();
5346
5347 Ok(ListTableBranchesResponse {
5348 branches,
5349 page_token: None,
5350 })
5351 }
5352
5353 async fn delete_table_branch(
5354 &self,
5355 request: DeleteTableBranchRequest,
5356 ) -> Result<DeleteTableBranchResponse> {
5357 self.record_op("delete_table_branch");
5358 if request.name.is_empty() {
5359 return Err(NamespaceError::InvalidInput {
5360 message: "branch name must not be empty for delete_table_branch".to_string(),
5361 }
5362 .into());
5363 }
5364
5365 let table_uri = self.resolve_table_location(&request.id).await?;
5366 let mut dataset = self
5367 .load_dataset(&table_uri, None, "delete_table_branch")
5368 .await?;
5369
5370 dataset
5371 .delete_branch(&request.name)
5372 .await
5373 .map_err(|e| match e {
5374 lance_core::Error::RefConflict { message } => NamespaceError::InvalidInput {
5375 message: format!(
5376 "branch '{}' for table at '{}': {}",
5377 request.name, table_uri, message
5378 ),
5379 }
5380 .into(),
5381 other => Self::map_branch_error(other, &request.name, &table_uri),
5382 })?;
5383
5384 Ok(DeleteTableBranchResponse {
5385 transaction_id: None,
5386 })
5387 }
5388
5389 fn namespace_id(&self) -> String {
5390 format!("DirectoryNamespace {{ root: {:?} }}", self.root)
5391 }
5392}
5393
5394#[cfg(test)]
5395mod tests {
5396 use super::*;
5397 use arrow_ipc::reader::{FileReader, StreamReader};
5398 use lance::dataset::Dataset;
5399 use lance::index::DatasetIndexExt;
5400 use lance_core::utils::tempfile::{TempStdDir, TempStrDir};
5401 use lance_core::utils::testing::CountingObjectStore;
5402 use lance_io::object_store::{providers::local::FileStoreProvider, uri_to_url};
5403 use lance_namespace::error::ErrorCode;
5404 use lance_namespace::models::{
5405 CreateTableRequest, JsonArrowDataType, JsonArrowField, JsonArrowSchema, ListTablesRequest,
5406 QueryTableRequestColumns,
5407 };
5408 use lance_namespace::schema::convert_json_arrow_schema;
5409 use std::io::Cursor;
5410 use std::sync::{
5411 Arc,
5412 atomic::{AtomicUsize, Ordering},
5413 };
5414 use url::Url;
5415
5416 fn assert_plan_contains_all(plan: &str, expected_fragments: &[&str], context: &str) {
5417 for expected_fragment in expected_fragments {
5418 assert!(
5419 plan.contains(expected_fragment),
5420 "{}. Missing fragment: '{}'. Plan:\n{}",
5421 context,
5422 expected_fragment,
5423 plan
5424 );
5425 }
5426 }
5427
5428 fn mutation_error_code(err: lance_core::Error) -> ErrorCode {
5429 match err {
5430 lance_core::Error::Namespace { source, .. } => source
5431 .downcast_ref::<NamespaceError>()
5432 .expect("mutation error should wrap a NamespaceError")
5433 .code(),
5434 other => panic!("expected Namespace error, got: {other:?}"),
5435 }
5436 }
5437
5438 #[test]
5443 fn test_map_mutation_error_commit_conflict_alignment() {
5444 let boxed = || -> Box<dyn std::error::Error + Send + Sync + 'static> {
5445 Box::<dyn std::error::Error + Send + Sync>::from("inner conflict")
5446 };
5447
5448 let throttling_cases = vec![lance_core::Error::commit_conflict_source(1, boxed())];
5449 for err in throttling_cases {
5450 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
5451 err,
5452 "update",
5453 "memory://t",
5454 ));
5455 assert_eq!(code, ErrorCode::Throttling);
5456 }
5457
5458 let concurrent_cases = vec![
5459 lance_core::Error::too_much_write_contention("contention"),
5460 lance_core::Error::retryable_commit_conflict_source(1, boxed()),
5461 lance_core::Error::incompatible_transaction_source(boxed()),
5462 lance_core::Error::version_conflict("conflict", 0, 3),
5463 ];
5464 for err in concurrent_cases {
5465 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
5466 err,
5467 "update",
5468 "memory://t",
5469 ));
5470 assert_eq!(code, ErrorCode::ConcurrentModification);
5471 }
5472 }
5473
5474 async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) {
5476 let temp_dir = TempStdDir::default();
5477
5478 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
5479 .build()
5480 .await
5481 .unwrap();
5482 (namespace, temp_dir)
5483 }
5484
5485 #[derive(Debug)]
5486 #[allow(dead_code)]
5487 struct CountingFileStoreProvider {
5488 listing_count: Arc<AtomicUsize>,
5489 }
5490
5491 #[async_trait]
5492 impl lance_io::object_store::ObjectStoreProvider for CountingFileStoreProvider {
5493 async fn new_store(
5494 &self,
5495 base_path: Url,
5496 params: &ObjectStoreParams,
5497 ) -> Result<ObjectStore> {
5498 let provider = FileStoreProvider;
5499 let mut store = provider.new_store(base_path, params).await?;
5500 store.inner = Arc::new(CountingObjectStore::new(
5501 store.inner.clone(),
5502 self.listing_count.clone(),
5503 ));
5504 Ok(store)
5505 }
5506
5507 fn extract_path(&self, url: &Url) -> Result<Path> {
5508 let provider = FileStoreProvider;
5509 provider.extract_path(url)
5510 }
5511
5512 fn calculate_object_store_prefix(
5513 &self,
5514 url: &Url,
5515 storage_options: Option<&HashMap<String, String>>,
5516 ) -> Result<String> {
5517 let provider = FileStoreProvider;
5518 provider.calculate_object_store_prefix(url, storage_options)
5519 }
5520 }
5521
5522 #[allow(dead_code)]
5523 fn file_object_store_uri(path: &str) -> String {
5524 let file_url = uri_to_url(path).unwrap();
5525 let mut url = Url::parse("file-object-store:///").unwrap();
5526 url.set_path(file_url.path());
5527 url.to_string()
5528 }
5529
5530 #[allow(dead_code)]
5531 fn build_listing_counting_session(listing_count: Arc<AtomicUsize>) -> Arc<Session> {
5532 let registry = Arc::new(ObjectStoreRegistry::default());
5533 registry.insert(
5534 "file-object-store",
5535 Arc::new(CountingFileStoreProvider { listing_count }),
5536 );
5537 Arc::new(Session::new(0, 0, registry))
5538 }
5539
5540 fn create_test_ipc_data(schema: &JsonArrowSchema) -> Vec<u8> {
5542 use arrow::ipc::writer::StreamWriter;
5543
5544 let arrow_schema = convert_json_arrow_schema(schema).unwrap();
5545 let arrow_schema = Arc::new(arrow_schema);
5546 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
5547 let mut buffer = Vec::new();
5548 {
5549 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
5550 writer.write(&batch).unwrap();
5551 writer.finish().unwrap();
5552 }
5553 buffer
5554 }
5555
5556 fn create_ipc_data_from_batches(
5557 schema: Arc<arrow_schema::Schema>,
5558 batches: Vec<arrow::record_batch::RecordBatch>,
5559 ) -> Vec<u8> {
5560 use arrow::ipc::writer::StreamWriter;
5561
5562 let mut buffer = Vec::new();
5563 {
5564 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
5565 for batch in &batches {
5566 writer.write(batch).unwrap();
5567 }
5568 writer.finish().unwrap();
5569 }
5570 buffer
5571 }
5572
5573 fn create_non_empty_test_ipc_data() -> Vec<u8> {
5574 use arrow::array::{Int32Array, StringArray};
5575 use arrow::record_batch::RecordBatch;
5576
5577 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
5578 let batch = RecordBatch::try_new(
5579 schema.clone(),
5580 vec![
5581 Arc::new(Int32Array::from(vec![1, 2])),
5582 Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])),
5583 ],
5584 )
5585 .unwrap();
5586 create_ipc_data_from_batches(schema, vec![batch])
5587 }
5588
5589 fn create_single_row_test_ipc_data() -> Vec<u8> {
5590 use arrow::array::{Int32Array, StringArray};
5591 use arrow::record_batch::RecordBatch;
5592
5593 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
5594 let batch = RecordBatch::try_new(
5595 schema.clone(),
5596 vec![
5597 Arc::new(Int32Array::from(vec![10])),
5598 Arc::new(StringArray::from(vec![Some("carol")])),
5599 ],
5600 )
5601 .unwrap();
5602 create_ipc_data_from_batches(schema, vec![batch])
5603 }
5604
5605 fn create_test_schema() -> JsonArrowSchema {
5607 let int_type = JsonArrowDataType::new("int32".to_string());
5608 let string_type = JsonArrowDataType::new("utf8".to_string());
5609
5610 let id_field = JsonArrowField {
5611 name: "id".to_string(),
5612 r#type: Box::new(int_type),
5613 nullable: false,
5614 metadata: None,
5615 };
5616
5617 let name_field = JsonArrowField {
5618 name: "name".to_string(),
5619 r#type: Box::new(string_type),
5620 nullable: true,
5621 metadata: None,
5622 };
5623
5624 JsonArrowSchema {
5625 fields: vec![id_field, name_field],
5626 metadata: None,
5627 }
5628 }
5629
5630 fn create_scalar_table_ipc_data() -> Vec<u8> {
5631 use arrow::array::{Int32Array, StringArray};
5632 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
5633
5634 let schema = Arc::new(ArrowSchema::new(vec![
5635 Field::new("id", DataType::Int32, false),
5636 Field::new("name", DataType::Utf8, true),
5637 ]));
5638 let batch = arrow::record_batch::RecordBatch::try_new(
5639 schema.clone(),
5640 vec![
5641 Arc::new(Int32Array::from(vec![1, 2, 3])),
5642 Arc::new(StringArray::from(vec!["alice", "bob", "cory"])),
5643 ],
5644 )
5645 .unwrap();
5646 create_ipc_data_from_batches(schema, vec![batch])
5647 }
5648
5649 async fn create_legacy_manifest_without_primary_key_metadata(root: &str) {
5650 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
5651 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
5652
5653 let schema = Arc::new(ArrowSchema::new(vec![
5654 Field::new("object_id", DataType::Utf8, false),
5655 Field::new("object_type", DataType::Utf8, false),
5656 Field::new("location", DataType::Utf8, true),
5657 Field::new("metadata", DataType::Utf8, true),
5658 Field::new(
5659 "base_objects",
5660 DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))),
5661 true,
5662 ),
5663 ]));
5664 let batch = RecordBatch::new_empty(schema.clone());
5665 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
5666 Dataset::write(Box::new(reader), &format!("{}/__manifest", root), None)
5667 .await
5668 .unwrap();
5669 }
5670
5671 async fn manifest_has_primary_key_metadata(root: &str) -> bool {
5672 let dataset = Dataset::open(&format!("{}/__manifest", root))
5673 .await
5674 .unwrap();
5675 dataset
5676 .schema()
5677 .field("object_id")
5678 .map(|field| {
5679 field
5680 .metadata
5681 .contains_key(lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
5682 })
5683 .unwrap_or(false)
5684 }
5685
5686 fn create_vector_table_ipc_data() -> Vec<u8> {
5687 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
5688 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
5689
5690 let schema = Arc::new(ArrowSchema::new(vec![
5691 Field::new("id", DataType::Int32, false),
5692 Field::new(
5693 "vector",
5694 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
5695 true,
5696 ),
5697 ]));
5698 let vector_field = Arc::new(Field::new("item", DataType::Float32, true));
5699 let vectors = FixedSizeListArray::try_new(
5700 vector_field,
5701 2,
5702 Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6])),
5703 None,
5704 )
5705 .unwrap();
5706 let batch = arrow::record_batch::RecordBatch::try_new(
5707 schema.clone(),
5708 vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(vectors)],
5709 )
5710 .unwrap();
5711 create_ipc_data_from_batches(schema, vec![batch])
5712 }
5713
5714 async fn create_scalar_table(namespace: &DirectoryNamespace, table_name: &str) {
5715 let mut create_table_request = CreateTableRequest::new();
5716 create_table_request.id = Some(vec![table_name.to_string()]);
5717 namespace
5718 .create_table(
5719 create_table_request,
5720 Bytes::from(create_scalar_table_ipc_data()),
5721 )
5722 .await
5723 .unwrap();
5724 }
5725
5726 async fn create_vector_table(namespace: &DirectoryNamespace, table_name: &str) {
5727 let mut create_table_request = CreateTableRequest::new();
5728 create_table_request.id = Some(vec![table_name.to_string()]);
5729 namespace
5730 .create_table(
5731 create_table_request,
5732 Bytes::from(create_vector_table_ipc_data()),
5733 )
5734 .await
5735 .unwrap();
5736 }
5737
5738 async fn open_dataset(namespace: &DirectoryNamespace, table_name: &str) -> Dataset {
5739 let mut describe_request = DescribeTableRequest::new();
5740 describe_request.id = Some(vec![table_name.to_string()]);
5741 let table_uri = namespace
5742 .describe_table(describe_request)
5743 .await
5744 .unwrap()
5745 .location
5746 .expect("table location should exist");
5747 Dataset::open(&table_uri).await.unwrap()
5748 }
5749
5750 async fn create_scalar_index(
5751 namespace: &DirectoryNamespace,
5752 table_name: &str,
5753 index_name: &str,
5754 ) -> Option<String> {
5755 use lance_namespace::models::CreateTableIndexRequest;
5756
5757 let mut create_index_request =
5758 CreateTableIndexRequest::new("id".to_string(), "BTREE".to_string());
5759 create_index_request.id = Some(vec![table_name.to_string()]);
5760 create_index_request.name = Some(index_name.to_string());
5761 namespace
5762 .create_table_scalar_index(create_index_request)
5763 .await
5764 .unwrap()
5765 .transaction_id
5766 }
5767
5768 async fn create_branch_with_commits(
5773 namespace: &DirectoryNamespace,
5774 table_name: &str,
5775 branch_name: &str,
5776 extra_versions: usize,
5777 ) -> String {
5778 let mut main = open_dataset(namespace, table_name).await;
5779 let fork_version = main.version().version;
5780 let branch = main
5781 .create_branch(branch_name, fork_version, None)
5782 .await
5783 .unwrap();
5784 let branch_uri = branch.uri().to_string();
5785 for i in 0..extra_versions {
5786 append_scalar_version(&branch_uri, (i as i32 + 1) * 100).await;
5787 }
5788 branch_uri
5789 }
5790
5791 async fn append_scalar_version(uri: &str, seed: i32) {
5794 use arrow::array::{Int32Array, StringArray};
5795 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
5796 let schema = Arc::new(ArrowSchema::new(vec![
5797 Field::new("id", DataType::Int32, false),
5798 Field::new("name", DataType::Utf8, true),
5799 ]));
5800 let batch = arrow::record_batch::RecordBatch::try_new(
5801 schema.clone(),
5802 vec![
5803 Arc::new(Int32Array::from(vec![seed, seed + 1])),
5804 Arc::new(StringArray::from(vec![Some("x"), Some("y")])),
5805 ],
5806 )
5807 .unwrap();
5808 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
5809 Dataset::write(
5810 reader,
5811 uri,
5812 Some(WriteParams {
5813 mode: WriteMode::Append,
5814 ..Default::default()
5815 }),
5816 )
5817 .await
5818 .unwrap();
5819 }
5820
5821 async fn list_versions(
5823 namespace: &DirectoryNamespace,
5824 table_name: &str,
5825 branch: Option<&str>,
5826 ) -> Result<Vec<TableVersion>> {
5827 let req = ListTableVersionsRequest {
5828 id: Some(vec![table_name.to_string()]),
5829 branch: branch.map(|b| b.to_string()),
5830 ..Default::default()
5831 };
5832 namespace.list_table_versions(req).await.map(|r| r.versions)
5833 }
5834
5835 #[tokio::test]
5836 async fn test_list_table_versions_on_branch() {
5837 let (namespace, _temp_dir) = create_test_namespace().await;
5838 create_scalar_table(&namespace, "users").await;
5839 create_branch_with_commits(&namespace, "users", "exp", 2).await;
5840
5841 let branch_versions = list_versions(&namespace, "users", Some("exp"))
5844 .await
5845 .unwrap();
5846 assert!(branch_versions.len() >= 2);
5847 assert!(
5848 branch_versions
5849 .iter()
5850 .all(|v| v.manifest_path.contains("tree/exp")),
5851 "branch versions must resolve to branch manifests: {:?}",
5852 branch_versions
5853 );
5854
5855 let main_versions = list_versions(&namespace, "users", None).await.unwrap();
5857 let main_explicit = list_versions(&namespace, "users", Some("main"))
5858 .await
5859 .unwrap();
5860 assert_eq!(main_versions.len(), main_explicit.len());
5861 assert!(
5862 main_versions
5863 .iter()
5864 .all(|v| !v.manifest_path.contains("tree/"))
5865 );
5866
5867 let missing = list_versions(&namespace, "users", Some("does-not-exist")).await;
5869 assert!(missing.is_err());
5870 assert!(missing.unwrap_err().to_string().contains("not found"));
5871 }
5872
5873 #[tokio::test]
5874 async fn test_describe_table_version_on_branch() {
5875 let (namespace, _temp_dir) = create_test_namespace().await;
5876 create_scalar_table(&namespace, "users").await;
5877 create_branch_with_commits(&namespace, "users", "exp", 2).await;
5878
5879 let branch_versions = list_versions(&namespace, "users", Some("exp"))
5880 .await
5881 .unwrap();
5882 let latest = branch_versions.iter().map(|v| v.version).max().unwrap();
5883
5884 let req = DescribeTableVersionRequest {
5886 id: Some(vec!["users".to_string()]),
5887 branch: Some("exp".to_string()),
5888 ..Default::default()
5889 };
5890 let resp = namespace.describe_table_version(req).await.unwrap();
5891 assert_eq!(resp.version.version, latest);
5892 assert!(resp.version.manifest_path.contains("tree/exp"));
5893
5894 let req = DescribeTableVersionRequest {
5896 id: Some(vec!["users".to_string()]),
5897 version: Some(latest),
5898 branch: Some("exp".to_string()),
5899 ..Default::default()
5900 };
5901 assert!(namespace.describe_table_version(req).await.is_ok());
5902
5903 let req = DescribeTableVersionRequest {
5905 id: Some(vec!["users".to_string()]),
5906 version: Some(999_999),
5907 branch: Some("exp".to_string()),
5908 ..Default::default()
5909 };
5910 assert!(namespace.describe_table_version(req).await.is_err());
5911
5912 let req = DescribeTableVersionRequest {
5914 id: Some(vec!["users".to_string()]),
5915 branch: Some("nope".to_string()),
5916 ..Default::default()
5917 };
5918 let err = namespace.describe_table_version(req).await;
5919 assert!(err.is_err() && err.unwrap_err().to_string().contains("not found"));
5920 }
5921
5922 #[tokio::test]
5923 async fn test_restore_table_on_branch() {
5924 use lance_namespace::models::RestoreTableRequest;
5925
5926 let (namespace, _temp_dir) = create_test_namespace().await;
5927 create_scalar_table(&namespace, "users").await;
5928 create_branch_with_commits(&namespace, "users", "exp", 2).await;
5929
5930 let before = list_versions(&namespace, "users", Some("exp"))
5931 .await
5932 .unwrap();
5933 let branch_latest = before.iter().map(|v| v.version).max().unwrap();
5934 let earliest = before.iter().map(|v| v.version).min().unwrap();
5935 let main_before = list_versions(&namespace, "users", None)
5936 .await
5937 .unwrap()
5938 .len();
5939
5940 let req = RestoreTableRequest {
5943 id: Some(vec!["users".to_string()]),
5944 version: earliest,
5945 branch: Some("exp".to_string()),
5946 ..Default::default()
5947 };
5948 let resp = namespace.restore_table(req).await.unwrap();
5949 assert!(resp.transaction_id.is_some());
5950
5951 let after = list_versions(&namespace, "users", Some("exp"))
5952 .await
5953 .unwrap();
5954 let new_latest = after.iter().map(|v| v.version).max().unwrap();
5955 assert!(
5956 new_latest > branch_latest,
5957 "restore should add a branch version"
5958 );
5959
5960 let main_after = list_versions(&namespace, "users", None)
5961 .await
5962 .unwrap()
5963 .len();
5964 assert_eq!(main_after, main_before, "main must be unaffected");
5965 }
5966
5967 #[tokio::test]
5968 async fn test_batch_delete_table_versions_on_branch() {
5969 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
5970
5971 let (namespace, _temp_dir) = create_test_namespace().await;
5972 create_scalar_table(&namespace, "users").await;
5973 create_branch_with_commits(&namespace, "users", "exp", 2).await;
5974
5975 let before = list_versions(&namespace, "users", Some("exp"))
5976 .await
5977 .unwrap();
5978 let main_before = list_versions(&namespace, "users", None).await.unwrap();
5979
5980 let req = BatchDeleteTableVersionsRequest {
5985 id: Some(vec!["users".to_string()]),
5986 branch: Some("exp".to_string()),
5987 ranges: vec![VersionRange::new(0, -1)],
5988 ..Default::default()
5989 };
5990 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
5991 assert_eq!(
5992 resp.deleted_count,
5993 Some(before.len() as i64),
5994 "every branch manifest should be physically deleted"
5995 );
5996
5997 assert!(
5999 list_versions(&namespace, "users", Some("exp"))
6000 .await
6001 .is_err()
6002 );
6003 let main_after = list_versions(&namespace, "users", None).await.unwrap();
6004 assert_eq!(
6005 main_after.len(),
6006 main_before.len(),
6007 "main must be untouched"
6008 );
6009 }
6010
6011 #[tokio::test]
6012 async fn test_create_table_version_on_branch() {
6013 use futures::TryStreamExt;
6014 use lance_namespace::models::CreateTableVersionRequest;
6015
6016 let (namespace, _temp_dir) = create_test_namespace().await;
6017 create_scalar_table(&namespace, "users").await;
6018 let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 1).await;
6019
6020 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
6022 let versions_dir = branch_ds.versions_dir();
6023 let store = branch_ds.object_store(None).await.unwrap();
6024 let existing = store
6025 .inner
6026 .list(Some(&versions_dir))
6027 .try_collect::<Vec<_>>()
6028 .await
6029 .unwrap()
6030 .into_iter()
6031 .find(|m| {
6032 m.location
6033 .filename()
6034 .map(|f| f.ends_with(".manifest"))
6035 .unwrap_or(false)
6036 })
6037 .expect("a branch manifest");
6038 let bytes = store
6039 .inner
6040 .get(&existing.location)
6041 .await
6042 .unwrap()
6043 .bytes()
6044 .await
6045 .unwrap();
6046 let staging = versions_dir.join("staging_manifest");
6047 store.inner.put(&staging, bytes.into()).await.unwrap();
6048
6049 let main_before = list_versions(&namespace, "users", None)
6050 .await
6051 .unwrap()
6052 .len();
6053 let new_version = list_versions(&namespace, "users", Some("exp"))
6054 .await
6055 .unwrap()
6056 .iter()
6057 .map(|v| v.version)
6058 .max()
6059 .unwrap()
6060 + 1;
6061
6062 let req = CreateTableVersionRequest {
6063 id: Some(vec!["users".to_string()]),
6064 version: new_version,
6065 manifest_path: staging.to_string(),
6066 naming_scheme: Some("V2".to_string()),
6067 branch: Some("exp".to_string()),
6068 ..Default::default()
6069 };
6070 let resp = namespace.create_table_version(req).await.unwrap();
6071 let info = resp.version.expect("version info");
6072 assert!(
6074 info.manifest_path.contains("tree/exp"),
6075 "got {}",
6076 info.manifest_path
6077 );
6078
6079 let after = list_versions(&namespace, "users", Some("exp"))
6081 .await
6082 .unwrap();
6083 assert!(after.iter().any(|v| v.version == new_version));
6084 let main_after = list_versions(&namespace, "users", None)
6085 .await
6086 .unwrap()
6087 .len();
6088 assert_eq!(main_after, main_before, "main must be unaffected");
6089 }
6090
6091 #[tokio::test]
6096 async fn test_external_manifest_store_resolves_branch_from_base_path() {
6097 use futures::TryStreamExt;
6098 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
6099 use lance_table::io::commit::external_manifest::ExternalManifestStore;
6100
6101 let (namespace, _temp_dir) = create_test_namespace().await;
6102 create_scalar_table(&namespace, "users").await; let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 2).await;
6104
6105 let namespace = Arc::new(namespace);
6106 let table_id = vec!["users".to_string()];
6107 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
6108 let branch_base = branch_ds.branch_location().path;
6109 let root_base = branch_ds.branch_location().find_main().unwrap().path;
6110 let store = LanceNamespaceExternalManifestStore::new(
6111 namespace.clone(),
6112 table_id.clone(),
6113 root_base.clone(),
6114 );
6115
6116 let (branch_latest, branch_path) = store
6119 .get_latest_version(branch_base.as_ref())
6120 .await
6121 .unwrap()
6122 .expect("branch has versions");
6123 let (_main_latest, main_path) = store
6124 .get_latest_version(root_base.as_ref())
6125 .await
6126 .unwrap()
6127 .expect("main has versions");
6128 assert!(
6129 branch_path.contains("tree/exp"),
6130 "branch latest must resolve to the branch tree: {}",
6131 branch_path
6132 );
6133 assert!(
6134 !main_path.contains("tree/exp"),
6135 "main latest must not resolve to a branch tree: {}",
6136 main_path
6137 );
6138
6139 let described = store
6141 .get(branch_base.as_ref(), branch_latest)
6142 .await
6143 .unwrap();
6144 assert!(
6145 described.contains("tree/exp"),
6146 "describe on the branch must resolve to the branch tree: {}",
6147 described
6148 );
6149
6150 assert!(store.get_latest_version("somewhere/else").await.is_err());
6152
6153 let versions_dir = branch_ds.versions_dir();
6156 let obj = branch_ds.object_store(None).await.unwrap();
6157 let existing = obj
6158 .inner
6159 .list(Some(&versions_dir))
6160 .try_collect::<Vec<_>>()
6161 .await
6162 .unwrap()
6163 .into_iter()
6164 .find(|m| {
6165 m.location
6166 .filename()
6167 .map(|f| f.ends_with(".manifest"))
6168 .unwrap_or(false)
6169 })
6170 .expect("a branch manifest");
6171 let bytes = obj
6172 .inner
6173 .get(&existing.location)
6174 .await
6175 .unwrap()
6176 .bytes()
6177 .await
6178 .unwrap();
6179 let size = bytes.len() as u64;
6180 let staging = versions_dir.clone().join("staging_manifest");
6181 obj.inner.put(&staging, bytes.into()).await.unwrap();
6182
6183 let committed = store
6184 .put(
6185 &branch_base,
6186 branch_latest + 1,
6187 &staging,
6188 size,
6189 None,
6190 obj.inner.as_ref(),
6191 ManifestNamingScheme::V2,
6192 )
6193 .await
6194 .unwrap();
6195 assert!(
6196 committed.path.to_string().contains("tree/exp"),
6197 "a commit through a branch-qualified base must land on the branch tree: {}",
6198 committed.path
6199 );
6200 }
6201
6202 #[tokio::test]
6210 async fn test_write_into_namespace_on_branch_appends_to_branch() {
6211 use lance::dataset::builder::DatasetBuilder;
6212 use lance_namespace::models::CreateTableBranchRequest;
6213
6214 let temp = TempStdDir::default();
6215 let namespace = Arc::new(
6216 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
6217 .manifest_enabled(true)
6218 .table_version_tracking_enabled(true)
6219 .ops_metrics_enabled(true)
6220 .build()
6221 .await
6222 .unwrap(),
6223 );
6224 let ns: Arc<dyn LanceNamespace> = namespace.clone();
6225 let table_id = vec!["t".to_string()];
6226 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
6228 id: Some(table_id.clone()),
6229 name: "exp".to_string(),
6230 ..Default::default()
6231 })
6232 .await
6233 .unwrap();
6234
6235 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
6236 ns.list_table_versions(ListTableVersionsRequest {
6237 id: Some(table_id),
6238 ..Default::default()
6239 })
6240 .await
6241 .unwrap()
6242 .versions
6243 .len()
6244 };
6245 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
6246 let commits_before = namespace
6247 .retrieve_ops_metrics()
6248 .get("create_table_version")
6249 .copied()
6250 .unwrap_or(0);
6251
6252 let branch_ds = Dataset::write_into_namespace_on_branch(
6253 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
6254 ns.clone(),
6255 table_id.clone(),
6256 "exp",
6257 Some(WriteParams {
6258 mode: WriteMode::Append,
6259 ..Default::default()
6260 }),
6261 )
6262 .await
6263 .unwrap();
6264 assert_eq!(branch_ds.manifest.branch.as_deref(), Some("exp"));
6265 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
6266
6267 let commits_after = namespace
6270 .retrieve_ops_metrics()
6271 .get("create_table_version")
6272 .copied()
6273 .unwrap_or(0);
6274 assert_eq!(
6275 commits_after,
6276 commits_before + 1,
6277 "the branch append must register its version via create_table_version"
6278 );
6279 let exp_versions = ns
6280 .list_table_versions(ListTableVersionsRequest {
6281 id: Some(table_id.clone()),
6282 branch: Some("exp".to_string()),
6283 ..Default::default()
6284 })
6285 .await
6286 .unwrap()
6287 .versions;
6288 assert!(
6289 exp_versions
6290 .iter()
6291 .all(|v| v.manifest_path.contains("tree/exp")),
6292 "branch versions must resolve to the branch tree: {:?}",
6293 exp_versions
6294 );
6295 assert_eq!(
6296 main_chain_len(ns.clone(), table_id.clone()).await,
6297 main_before,
6298 "main's catalog must be untouched by the branch append"
6299 );
6300
6301 Dataset::write_into_namespace(
6304 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
6305 ns.clone(),
6306 table_id.clone(),
6307 Some(WriteParams {
6308 mode: WriteMode::Append,
6309 ..Default::default()
6310 }),
6311 )
6312 .await
6313 .unwrap();
6314 assert_eq!(
6315 main_chain_len(ns.clone(), table_id.clone()).await,
6316 main_before + 1,
6317 "a managed main append must register its version in the catalog"
6318 );
6319 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6320 .await
6321 .unwrap()
6322 .load()
6323 .await
6324 .unwrap();
6325 assert_eq!(
6326 scan_id_column(&fresh).await,
6327 vec![1, 2, 100],
6328 "a fresh managed open must resolve the appended version, not a stale latest"
6329 );
6330 }
6331
6332 #[tokio::test]
6334 async fn test_write_into_namespace_on_branch_rejects_create() {
6335 use arrow::array::{Int32Array, StringArray};
6336 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6337
6338 let (namespace, _temp_dir) = create_test_namespace().await;
6339 let namespace = Arc::new(namespace);
6340
6341 let schema = Arc::new(ArrowSchema::new(vec![
6342 Field::new("id", DataType::Int32, false),
6343 Field::new("name", DataType::Utf8, true),
6344 ]));
6345 let batch = arrow::record_batch::RecordBatch::try_new(
6346 schema.clone(),
6347 vec![
6348 Arc::new(Int32Array::from(vec![1])),
6349 Arc::new(StringArray::from(vec![Some("a")])),
6350 ],
6351 )
6352 .unwrap();
6353 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
6354
6355 let result = Dataset::write_into_namespace_on_branch(
6356 reader,
6357 namespace.clone(),
6358 vec!["new_table".to_string()],
6359 "exp",
6360 Some(WriteParams {
6361 mode: WriteMode::Create,
6362 ..Default::default()
6363 }),
6364 )
6365 .await;
6366 assert!(result.is_err(), "create on a branch must be rejected");
6367 assert!(
6368 result.unwrap_err().to_string().contains("branch"),
6369 "error should mention the branch restriction"
6370 );
6371 }
6372
6373 #[tokio::test]
6374 async fn test_branch_name_validation_rejects_traversal() {
6375 let (namespace, _temp_dir) = create_test_namespace().await;
6376 create_scalar_table(&namespace, "users").await;
6377
6378 let err = list_versions(&namespace, "users", Some("../evil")).await;
6381 assert!(err.is_err());
6382 assert!(err.unwrap_err().to_string().contains("invalid branch name"));
6383 }
6384
6385 #[tokio::test]
6386 async fn test_branch_ops_reject_zombie_branch() {
6387 use futures::TryStreamExt;
6388 use lance_namespace::models::{
6389 BatchDeleteTableVersionsRequest, CreateTableVersionRequest, RestoreTableRequest,
6390 VersionRange,
6391 };
6392
6393 let (namespace, _temp_dir) = create_test_namespace().await;
6394 create_scalar_table(&namespace, "users").await;
6395
6396 let dataset = open_dataset(&namespace, "users").await;
6397 let store = dataset.object_store(None).await.unwrap();
6398 let manifest = store
6399 .inner
6400 .list(Some(&dataset.versions_dir()))
6401 .try_collect::<Vec<_>>()
6402 .await
6403 .unwrap()
6404 .into_iter()
6405 .find(|m| {
6406 m.location
6407 .filename()
6408 .map(|f| f.ends_with(".manifest"))
6409 .unwrap_or(false)
6410 })
6411 .expect("a manifest");
6412 let bytes = store
6413 .inner
6414 .get(&manifest.location)
6415 .await
6416 .unwrap()
6417 .bytes()
6418 .await
6419 .unwrap();
6420 let zombie = dataset
6421 .branch_location()
6422 .find_branch(Some("ghost"))
6423 .unwrap()
6424 .path
6425 .join(VERSIONS_DIR)
6426 .join(manifest.location.filename().unwrap());
6427 store.inner.put(&zombie, bytes.into()).await.unwrap();
6428
6429 assert!(dataset.branches().get("ghost").await.is_err());
6430
6431 fn rejected<T: std::fmt::Debug>(label: &str, r: Result<T>) {
6432 match r {
6433 Ok(v) => panic!("{label} must reject the zombie branch, got Ok({v:?})"),
6434 Err(e) => assert!(e.to_string().contains("not found"), "{label}: {e}"),
6435 }
6436 }
6437
6438 rejected(
6439 "list",
6440 list_versions(&namespace, "users", Some("ghost")).await,
6441 );
6442 rejected(
6443 "describe",
6444 namespace
6445 .describe_table_version(DescribeTableVersionRequest {
6446 id: Some(vec!["users".to_string()]),
6447 branch: Some("ghost".to_string()),
6448 ..Default::default()
6449 })
6450 .await,
6451 );
6452 rejected(
6453 "create",
6454 namespace
6455 .create_table_version(CreateTableVersionRequest {
6456 id: Some(vec!["users".to_string()]),
6457 version: 2,
6458 manifest_path: zombie.to_string(),
6459 branch: Some("ghost".to_string()),
6460 ..Default::default()
6461 })
6462 .await,
6463 );
6464 rejected(
6465 "restore",
6466 namespace
6467 .restore_table(RestoreTableRequest {
6468 id: Some(vec!["users".to_string()]),
6469 version: 1,
6470 branch: Some("ghost".to_string()),
6471 ..Default::default()
6472 })
6473 .await,
6474 );
6475 rejected(
6476 "batch_delete",
6477 namespace
6478 .batch_delete_table_versions(BatchDeleteTableVersionsRequest {
6479 id: Some(vec!["users".to_string()]),
6480 branch: Some("ghost".to_string()),
6481 ranges: vec![VersionRange::new(1, 1)],
6482 ..Default::default()
6483 })
6484 .await,
6485 );
6486 }
6487
6488 #[tokio::test]
6493 async fn test_batch_delete_table_versions_main_v2() {
6494 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
6495
6496 let (namespace, _temp_dir) = create_test_namespace().await;
6497 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
6499 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
6503 assert!(before.len() >= 3);
6504 assert!(
6507 before
6508 .iter()
6509 .all(|v| v.manifest_path.rsplit('/').next().unwrap().len() == 29),
6510 "expected V2-named manifests: {:?}",
6511 before
6512 );
6513 let min_v = before.iter().map(|v| v.version).min().unwrap();
6514 let max_v = before.iter().map(|v| v.version).max().unwrap();
6515
6516 let req = BatchDeleteTableVersionsRequest {
6519 id: Some(vec!["users".to_string()]),
6520 ranges: vec![VersionRange::new(min_v, max_v)],
6521 ..Default::default()
6522 };
6523 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
6524 assert_eq!(
6525 resp.deleted_count,
6526 Some((before.len() - 1) as i64),
6527 "V2 manifests must actually be deleted (was 0 before the fix)"
6528 );
6529
6530 let after = list_versions(&namespace, "users", None).await.unwrap();
6531 assert_eq!(after.len(), 1);
6532 assert_eq!(after[0].version, max_v);
6533 }
6534
6535 #[tokio::test]
6537 async fn test_batch_delete_end_is_exclusive() {
6538 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
6539
6540 let (namespace, _temp_dir) = create_test_namespace().await;
6541 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
6543 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
6547 let min_v = before.iter().map(|v| v.version).min().unwrap();
6548
6549 let req = BatchDeleteTableVersionsRequest {
6550 id: Some(vec!["users".to_string()]),
6551 ranges: vec![VersionRange::new(min_v, min_v + 1)],
6552 ..Default::default()
6553 };
6554 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
6555 assert_eq!(
6556 resp.deleted_count,
6557 Some(1),
6558 "only min_v is in [min_v, min_v+1)"
6559 );
6560
6561 let after = list_versions(&namespace, "users", None).await.unwrap();
6562 assert!(
6563 !after.iter().any(|v| v.version == min_v),
6564 "min_v must be deleted"
6565 );
6566 assert_eq!(after.len(), before.len() - 1, "exactly one version removed");
6567 }
6568
6569 #[tokio::test]
6570 async fn test_batch_delete_rejects_unbounded_range() {
6571 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
6572
6573 let (namespace, _temp_dir) = create_test_namespace().await;
6574 create_scalar_table(&namespace, "users").await;
6575
6576 let req = BatchDeleteTableVersionsRequest {
6579 id: Some(vec!["users".to_string()]),
6580 ranges: vec![VersionRange::new(0, i64::MAX)],
6581 ..Default::default()
6582 };
6583 let err = namespace.batch_delete_table_versions(req).await;
6584 assert!(err.is_err());
6585 assert!(
6586 err.unwrap_err().to_string().contains("limit"),
6587 "expected a range-too-large error"
6588 );
6589 }
6590
6591 async fn create_managed_namespace(path: &str) -> Arc<dyn LanceNamespace> {
6593 Arc::new(
6594 DirectoryNamespaceBuilder::new(path)
6595 .manifest_enabled(true)
6596 .table_version_tracking_enabled(true)
6597 .build()
6598 .await
6599 .unwrap(),
6600 )
6601 }
6602
6603 fn single_int_schema() -> Arc<arrow::datatypes::Schema> {
6604 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6605 Arc::new(ArrowSchema::new(vec![Field::new(
6606 "id",
6607 DataType::Int32,
6608 false,
6609 )]))
6610 }
6611
6612 fn single_int_batch(seed: i32) -> arrow::record_batch::RecordBatch {
6613 use arrow::array::Int32Array;
6614 arrow::record_batch::RecordBatch::try_new(
6615 single_int_schema(),
6616 vec![Arc::new(Int32Array::from(vec![seed]))],
6617 )
6618 .unwrap()
6619 }
6620
6621 async fn create_managed_table(ns: &Arc<dyn LanceNamespace>, table_id: &[String]) -> Dataset {
6624 let mut ds = Dataset::write_into_namespace(
6625 RecordBatchIterator::new(vec![Ok(single_int_batch(1))], single_int_schema()),
6626 ns.clone(),
6627 table_id.to_vec(),
6628 Some(WriteParams {
6629 mode: WriteMode::Create,
6630 ..Default::default()
6631 }),
6632 )
6633 .await
6634 .unwrap();
6635 ds.append(
6636 RecordBatchIterator::new(vec![Ok(single_int_batch(2))], single_int_schema()),
6637 None,
6638 )
6639 .await
6640 .unwrap();
6641 ds
6642 }
6643
6644 async fn scan_id_column(ds: &Dataset) -> Vec<i32> {
6646 use arrow::array::Int32Array;
6647 use futures::TryStreamExt;
6648 let batches: Vec<arrow::record_batch::RecordBatch> = ds
6649 .scan()
6650 .try_into_stream()
6651 .await
6652 .unwrap()
6653 .try_collect()
6654 .await
6655 .unwrap();
6656 let mut ids: Vec<i32> = batches
6657 .iter()
6658 .flat_map(|b| {
6659 b.column(0)
6660 .as_any()
6661 .downcast_ref::<Int32Array>()
6662 .unwrap()
6663 .values()
6664 .to_vec()
6665 })
6666 .collect();
6667 ids.sort();
6668 ids
6669 }
6670
6671 #[tokio::test]
6676 async fn test_managed_branch_open_and_commit() {
6677 use futures::TryStreamExt;
6678 use lance::dataset::builder::DatasetBuilder;
6679 use lance_namespace::models::CreateTableBranchRequest;
6680
6681 let temp = TempStdDir::default();
6682 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
6683 let table_id = vec!["t".to_string()];
6684 create_managed_table(&ns, &table_id).await;
6685 let main_before = ns
6686 .list_table_versions(ListTableVersionsRequest {
6687 id: Some(table_id.clone()),
6688 ..Default::default()
6689 })
6690 .await
6691 .unwrap()
6692 .versions
6693 .len();
6694
6695 ns.create_table_branch(CreateTableBranchRequest {
6698 id: Some(table_id.clone()),
6699 name: "exp".to_string(),
6700 ..Default::default()
6701 })
6702 .await
6703 .unwrap();
6704
6705 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6708 .await
6709 .unwrap()
6710 .with_branch("exp", None)
6711 .load()
6712 .await
6713 .unwrap();
6714 assert_eq!(
6715 branch_ds.manifest.branch.as_deref(),
6716 Some("exp"),
6717 "with_branch on a managed table must open the branch chain"
6718 );
6719 let branch_base = branch_ds.branch_location().path;
6720 assert!(
6721 branch_base.as_ref().ends_with("tree/exp"),
6722 "the branch dataset must be rooted at the branch chain: {}",
6723 branch_base
6724 );
6725 let branch_v_before = branch_ds.version().version;
6726
6727 branch_ds
6729 .append(
6730 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
6731 None,
6732 )
6733 .await
6734 .unwrap();
6735 assert_eq!(
6736 branch_ds.manifest.branch.as_deref(),
6737 Some("exp"),
6738 "the commit must stay on the branch"
6739 );
6740 assert!(
6741 branch_ds.version().version > branch_v_before,
6742 "the branch version must advance after the commit"
6743 );
6744 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
6745
6746 let store = branch_ds.object_store(None).await.unwrap();
6750 let branch_data = branch_base.clone().join("data");
6751 let branch_files = store
6752 .inner
6753 .list(Some(&branch_data))
6754 .try_collect::<Vec<_>>()
6755 .await
6756 .unwrap();
6757 assert!(
6758 !branch_files.is_empty(),
6759 "the branch commit must place data files under the branch chain"
6760 );
6761
6762 let table_uri = ns
6764 .describe_table(DescribeTableRequest {
6765 id: Some(table_id.clone()),
6766 ..Default::default()
6767 })
6768 .await
6769 .unwrap()
6770 .location
6771 .unwrap();
6772 let fs_branch_ds = DatasetBuilder::from_uri(&table_uri)
6773 .with_branch("exp", None)
6774 .load()
6775 .await
6776 .unwrap();
6777 assert_eq!(fs_branch_ds.manifest.branch.as_deref(), Some("exp"));
6778 assert_eq!(scan_id_column(&fs_branch_ds).await, vec![1, 2, 3]);
6779
6780 let main_after = ns
6783 .list_table_versions(ListTableVersionsRequest {
6784 id: Some(table_id.clone()),
6785 ..Default::default()
6786 })
6787 .await
6788 .unwrap()
6789 .versions
6790 .len();
6791 assert_eq!(
6792 main_after, main_before,
6793 "committing on the branch must not change main's chain"
6794 );
6795 let main_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6796 .await
6797 .unwrap()
6798 .load()
6799 .await
6800 .unwrap();
6801 assert_eq!(main_ds.manifest.branch, None);
6802 assert_eq!(scan_id_column(&main_ds).await, vec![1, 2]);
6803 }
6804
6805 #[tokio::test]
6810 async fn test_managed_branch_tags() {
6811 use lance::dataset::builder::DatasetBuilder;
6812 use lance::dataset::refs::Ref;
6813 use lance_namespace::models::CreateTableBranchRequest;
6814
6815 let temp = TempStdDir::default();
6816 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
6817 let table_id = vec!["t".to_string()];
6818 let main_ds = create_managed_table(&ns, &table_id).await;
6819 ns.create_table_branch(CreateTableBranchRequest {
6820 id: Some(table_id.clone()),
6821 name: "exp".to_string(),
6822 ..Default::default()
6823 })
6824 .await
6825 .unwrap();
6826 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6827 .await
6828 .unwrap()
6829 .with_branch("exp", None)
6830 .load()
6831 .await
6832 .unwrap();
6833 branch_ds
6834 .append(
6835 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
6836 None,
6837 )
6838 .await
6839 .unwrap();
6840 let branch_version = branch_ds.version().version;
6841
6842 main_ds
6845 .tags()
6846 .create("exp-tag", ("exp", Some(branch_version)))
6847 .await
6848 .unwrap();
6849 let tag = main_ds.tags().get("exp-tag").await.unwrap();
6850 assert_eq!(tag.branch.as_deref(), Some("exp"));
6851 assert_eq!(tag.version, branch_version);
6852
6853 branch_ds
6855 .tags()
6856 .create("exp-tag2", branch_version)
6857 .await
6858 .unwrap();
6859 let tag2 = branch_ds.tags().get("exp-tag2").await.unwrap();
6860 assert_eq!(tag2.branch.as_deref(), Some("exp"));
6861
6862 let tag_open = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6865 .await
6866 .unwrap()
6867 .with_tag("exp-tag")
6868 .load()
6869 .await
6870 .unwrap();
6871 assert_eq!(tag_open.manifest.branch.as_deref(), Some("exp"));
6872 assert_eq!(tag_open.version().version, branch_version);
6873 assert_eq!(scan_id_column(&tag_open).await, vec![1, 2, 3]);
6874
6875 let tag_checkout = main_ds
6877 .checkout_version(Ref::Tag("exp-tag".to_string()))
6878 .await
6879 .unwrap();
6880 assert_eq!(tag_checkout.manifest.branch.as_deref(), Some("exp"));
6881 assert_eq!(scan_id_column(&tag_checkout).await, vec![1, 2, 3]);
6882
6883 let err = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6885 .await
6886 .unwrap()
6887 .with_tag("no-such-tag")
6888 .load()
6889 .await;
6890 assert!(err.is_err(), "a missing tag must error");
6891 }
6892
6893 #[tokio::test]
6898 async fn test_managed_cross_branch_checkout() {
6899 use lance::dataset::builder::DatasetBuilder;
6900 use lance::dataset::refs::Ref;
6901 use lance_namespace::models::CreateTableBranchRequest;
6902
6903 let temp = TempStdDir::default();
6904 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
6905 let table_id = vec!["t".to_string()];
6906 let mut main_ds = create_managed_table(&ns, &table_id).await;
6907 ns.create_table_branch(CreateTableBranchRequest {
6908 id: Some(table_id.clone()),
6909 name: "exp".to_string(),
6910 ..Default::default()
6911 })
6912 .await
6913 .unwrap();
6914
6915 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
6917 .await
6918 .unwrap()
6919 .with_branch("exp", None)
6920 .load()
6921 .await
6922 .unwrap();
6923 branch_ds
6924 .append(
6925 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
6926 None,
6927 )
6928 .await
6929 .unwrap();
6930 let overlap_version = branch_ds.version().version;
6931 while main_ds.version().version < overlap_version {
6932 main_ds
6933 .append(
6934 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
6935 None,
6936 )
6937 .await
6938 .unwrap();
6939 }
6940
6941 let on_branch = main_ds
6944 .checkout_version(Ref::Version(Some("exp".to_string()), Some(overlap_version)))
6945 .await
6946 .unwrap();
6947 assert_eq!(on_branch.manifest.branch.as_deref(), Some("exp"));
6948 assert_eq!(scan_id_column(&on_branch).await, vec![1, 2, 3]);
6949
6950 let mut on_branch_latest = main_ds.checkout_branch("exp").await.unwrap();
6952 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
6953 assert_eq!(on_branch_latest.version().version, overlap_version);
6954
6955 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
6958 ns.list_table_versions(ListTableVersionsRequest {
6959 id: Some(table_id),
6960 ..Default::default()
6961 })
6962 .await
6963 .unwrap()
6964 .versions
6965 .len()
6966 };
6967 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
6968 on_branch_latest
6969 .append(
6970 RecordBatchIterator::new(vec![Ok(single_int_batch(4))], single_int_schema()),
6971 None,
6972 )
6973 .await
6974 .unwrap();
6975 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
6976 assert_eq!(scan_id_column(&on_branch_latest).await, vec![1, 2, 3, 4]);
6977 assert_eq!(
6978 main_chain_len(ns.clone(), table_id.clone()).await,
6979 main_before,
6980 "a commit on the checked-out branch must not advance main's chain"
6981 );
6982
6983 let on_main = branch_ds
6985 .checkout_version(Ref::Version(None, Some(1)))
6986 .await
6987 .unwrap();
6988 assert_eq!(on_main.manifest.branch, None);
6989 assert_eq!(scan_id_column(&on_main).await, vec![1]);
6990
6991 ns.create_table_branch(CreateTableBranchRequest {
6993 id: Some(table_id.clone()),
6994 name: "exp2".to_string(),
6995 ..Default::default()
6996 })
6997 .await
6998 .unwrap();
6999 let on_branch2 = branch_ds.checkout_branch("exp2").await.unwrap();
7000 assert_eq!(on_branch2.manifest.branch.as_deref(), Some("exp2"));
7001
7002 let err = main_ds
7004 .checkout_version(Ref::Version(Some("exp".to_string()), Some(999)))
7005 .await;
7006 assert!(err.is_err(), "a version missing from the branch must error");
7007 }
7008
7009 #[tokio::test]
7015 async fn test_commit_builder_honors_explicit_handler_for_dataset_dest() {
7016 use lance::dataset::write::{CommitBuilder, InsertBuilder};
7017 use lance::dataset::{WriteDestination, builder::DatasetBuilder};
7018 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
7019 use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
7020
7021 let temp = TempStdDir::default();
7022 let namespace = Arc::new(
7023 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
7024 .manifest_enabled(true)
7025 .table_version_tracking_enabled(true)
7026 .ops_metrics_enabled(true)
7027 .build()
7028 .await
7029 .unwrap(),
7030 );
7031 let ns: Arc<dyn LanceNamespace> = namespace.clone();
7032 let table_id = vec!["t".to_string()];
7033 create_managed_table(&ns, &table_id).await; let table_uri = ns
7037 .describe_table(DescribeTableRequest {
7038 id: Some(table_id.clone()),
7039 ..Default::default()
7040 })
7041 .await
7042 .unwrap()
7043 .location
7044 .unwrap();
7045 let plain_ds = Arc::new(Dataset::open(&table_uri).await.unwrap());
7046
7047 let transaction = InsertBuilder::new(WriteDestination::Dataset(plain_ds.clone()))
7048 .with_params(&WriteParams {
7049 mode: WriteMode::Append,
7050 ..Default::default()
7051 })
7052 .execute_uncommitted(vec![single_int_batch(3)])
7053 .await
7054 .unwrap();
7055
7056 let handler = Arc::new(ExternalManifestCommitHandler {
7057 external_manifest_store: Arc::new(
7058 LanceNamespaceExternalManifestStore::for_table_uri(
7059 ns.clone(),
7060 table_id.clone(),
7061 &table_uri,
7062 )
7063 .unwrap(),
7064 ),
7065 });
7066 let commits_before = namespace
7067 .retrieve_ops_metrics()
7068 .get("create_table_version")
7069 .copied()
7070 .unwrap_or(0);
7071 let committed = CommitBuilder::new(WriteDestination::Dataset(plain_ds))
7072 .with_commit_handler(handler)
7073 .execute(transaction)
7074 .await
7075 .unwrap();
7076 assert_eq!(scan_id_column(&committed).await, vec![1, 2, 3]);
7077
7078 let commits_after = namespace
7079 .retrieve_ops_metrics()
7080 .get("create_table_version")
7081 .copied()
7082 .unwrap_or(0);
7083 assert_eq!(
7084 commits_after,
7085 commits_before + 1,
7086 "the explicit handler must route the commit through create_table_version"
7087 );
7088 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
7089 .await
7090 .unwrap()
7091 .load()
7092 .await
7093 .unwrap();
7094 assert_eq!(
7095 scan_id_column(&fresh).await,
7096 vec![1, 2, 3],
7097 "a fresh managed open must resolve the committed version"
7098 );
7099 }
7100
7101 #[tokio::test]
7103 async fn test_managed_branch_from_non_latest_fork() {
7104 use lance::dataset::builder::DatasetBuilder;
7105 use lance_namespace::models::CreateTableBranchRequest;
7106
7107 let temp = TempStdDir::default();
7108 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
7109 let table_id = vec!["t".to_string()];
7110 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
7113 id: Some(table_id.clone()),
7114 name: "old".to_string(),
7115 from_version: Some(1),
7116 ..Default::default()
7117 })
7118 .await
7119 .unwrap();
7120
7121 let old_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
7122 .await
7123 .unwrap()
7124 .with_branch("old", None)
7125 .load()
7126 .await
7127 .unwrap();
7128 assert_eq!(old_ds.manifest.branch.as_deref(), Some("old"));
7129 assert_eq!(
7130 scan_id_column(&old_ds).await,
7131 vec![1],
7132 "the fork must contain only the fork-point data"
7133 );
7134 }
7135
7136 #[test]
7139 fn test_manifest_version_from_filename() {
7140 assert_eq!(
7142 DirectoryNamespace::manifest_version_from_filename("5.manifest"),
7143 Some(5)
7144 );
7145 assert_eq!(
7146 DirectoryNamespace::manifest_version_from_filename("0.manifest"),
7147 Some(0)
7148 );
7149 let v2_five = format!("{:020}.manifest", u64::MAX - 5);
7151 assert_eq!(
7152 DirectoryNamespace::manifest_version_from_filename(&v2_five),
7153 Some(5)
7154 );
7155 let v2_zero = format!("{:020}.manifest", u64::MAX);
7156 assert_eq!(
7157 DirectoryNamespace::manifest_version_from_filename(&v2_zero),
7158 Some(0)
7159 );
7160 assert_eq!(
7162 DirectoryNamespace::manifest_version_from_filename("data.lance"),
7163 None
7164 );
7165 assert_eq!(
7166 DirectoryNamespace::manifest_version_from_filename("d5.manifest"),
7167 None
7168 );
7169 }
7170
7171 #[tokio::test]
7172 async fn test_create_table() {
7173 let (namespace, _temp_dir) = create_test_namespace().await;
7174
7175 let schema = create_test_schema();
7177 let ipc_data = create_test_ipc_data(&schema);
7178
7179 let mut request = CreateTableRequest::new();
7180 request.id = Some(vec!["test_table".to_string()]);
7181
7182 let response = namespace
7183 .create_table(request, bytes::Bytes::from(ipc_data))
7184 .await
7185 .unwrap();
7186
7187 assert!(response.location.is_some());
7188 assert!(response.location.unwrap().ends_with("test_table.lance"));
7189 assert_eq!(response.version, Some(1));
7190 }
7191
7192 #[tokio::test]
7193 async fn test_create_table_without_data() {
7194 let (namespace, _temp_dir) = create_test_namespace().await;
7195
7196 let mut request = CreateTableRequest::new();
7197 request.id = Some(vec!["test_table".to_string()]);
7198
7199 let result = namespace.create_table(request, bytes::Bytes::new()).await;
7200 assert!(result.is_err());
7201 assert!(
7202 result
7203 .unwrap_err()
7204 .to_string()
7205 .contains("Arrow IPC stream) is required")
7206 );
7207 }
7208
7209 #[tokio::test]
7210 async fn test_create_table_with_invalid_id() {
7211 let (namespace, _temp_dir) = create_test_namespace().await;
7212
7213 let schema = create_test_schema();
7215 let ipc_data = create_test_ipc_data(&schema);
7216
7217 let mut request = CreateTableRequest::new();
7219 request.id = Some(vec![]);
7220
7221 let result = namespace
7222 .create_table(request, bytes::Bytes::from(ipc_data.clone()))
7223 .await;
7224 assert!(result.is_err());
7225
7226 let mut create_ns_req = CreateNamespaceRequest::new();
7229 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
7230 namespace.create_namespace(create_ns_req).await.unwrap();
7231
7232 let mut request = CreateTableRequest::new();
7234 request.id = Some(vec!["test_namespace".to_string(), "table".to_string()]);
7235
7236 let result = namespace
7237 .create_table(request, bytes::Bytes::from(ipc_data))
7238 .await;
7239 assert!(
7241 result.is_ok(),
7242 "Multi-level table IDs should work with manifest enabled"
7243 );
7244 }
7245
7246 #[tokio::test]
7247 async fn test_list_tables() {
7248 let (namespace, _temp_dir) = create_test_namespace().await;
7249
7250 let mut request = ListTablesRequest::new();
7252 request.id = Some(vec![]);
7253 let response = namespace.list_tables(request).await.unwrap();
7254 assert_eq!(response.tables.len(), 0);
7255
7256 let schema = create_test_schema();
7258 let ipc_data = create_test_ipc_data(&schema);
7259
7260 let mut create_request = CreateTableRequest::new();
7262 create_request.id = Some(vec!["table1".to_string()]);
7263 namespace
7264 .create_table(create_request, bytes::Bytes::from(ipc_data.clone()))
7265 .await
7266 .unwrap();
7267
7268 let mut create_request = CreateTableRequest::new();
7270 create_request.id = Some(vec!["table2".to_string()]);
7271 namespace
7272 .create_table(create_request, bytes::Bytes::from(ipc_data))
7273 .await
7274 .unwrap();
7275
7276 let mut request = ListTablesRequest::new();
7278 request.id = Some(vec![]);
7279 let response = namespace.list_tables(request).await.unwrap();
7280 let tables = response.tables;
7281 assert_eq!(tables.len(), 2);
7282 assert!(tables.contains(&"table1".to_string()));
7283 assert!(tables.contains(&"table2".to_string()));
7284 }
7285
7286 #[tokio::test]
7287 async fn test_list_tables_pagination() {
7288 let (namespace, _temp_dir) = create_test_namespace().await;
7289
7290 let schema = create_test_schema();
7291 let ipc_data = create_test_ipc_data(&schema);
7292
7293 for name in ["alpha", "bravo", "charlie"] {
7294 let mut req = CreateTableRequest::new();
7295 req.id = Some(vec![name.to_string()]);
7296 namespace
7297 .create_table(req, bytes::Bytes::from(ipc_data.clone()))
7298 .await
7299 .unwrap();
7300 }
7301
7302 let first_page = namespace
7304 .list_tables(ListTablesRequest {
7305 id: Some(vec![]),
7306 limit: Some(2),
7307 ..Default::default()
7308 })
7309 .await
7310 .unwrap();
7311
7312 assert_eq!(first_page.tables, vec!["alpha", "bravo"]);
7313 assert_eq!(first_page.page_token.as_deref(), Some("bravo"));
7314
7315 let second_page = namespace
7317 .list_tables(ListTablesRequest {
7318 id: Some(vec![]),
7319 limit: Some(2),
7320 page_token: first_page.page_token.clone(),
7321 ..Default::default()
7322 })
7323 .await
7324 .unwrap();
7325
7326 assert_eq!(second_page.tables, vec!["charlie"]);
7327 assert!(second_page.page_token.is_none());
7328 }
7329
7330 #[tokio::test]
7331 async fn test_list_tables_pagination_limit_zero() {
7332 let (namespace, _temp_dir) = create_test_namespace().await;
7333
7334 let schema = create_test_schema();
7335 let ipc_data = create_test_ipc_data(&schema);
7336
7337 let mut req = CreateTableRequest::new();
7338 req.id = Some(vec!["alpha".to_string()]);
7339 namespace
7340 .create_table(req, bytes::Bytes::from(ipc_data))
7341 .await
7342 .unwrap();
7343
7344 let response = namespace
7345 .list_tables(ListTablesRequest {
7346 id: Some(vec![]),
7347 limit: Some(0),
7348 ..Default::default()
7349 })
7350 .await
7351 .unwrap();
7352
7353 assert!(response.tables.is_empty());
7354 assert!(response.page_token.is_none());
7355 }
7356
7357 #[tokio::test]
7358 async fn test_list_tables_with_namespace_id() {
7359 let (namespace, _temp_dir) = create_test_namespace().await;
7360
7361 let mut create_ns_req = CreateNamespaceRequest::new();
7363 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
7364 namespace.create_namespace(create_ns_req).await.unwrap();
7365
7366 let mut request = ListTablesRequest::new();
7368 request.id = Some(vec!["test_namespace".to_string()]);
7369
7370 let result = namespace.list_tables(request).await;
7371 assert!(
7373 result.is_ok(),
7374 "list_tables should work with child namespace when manifest is enabled"
7375 );
7376 let response = result.unwrap();
7377 assert_eq!(
7378 response.tables.len(),
7379 0,
7380 "Namespace should have no tables yet"
7381 );
7382 }
7383
7384 #[tokio::test]
7385 async fn test_create_scalar_index() {
7386 let (namespace, _temp_dir) = create_test_namespace().await;
7387 create_scalar_table(&namespace, "users").await;
7388
7389 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
7390 let dataset = open_dataset(&namespace, "users").await;
7391 let expected_transaction_id = dataset
7392 .read_transaction()
7393 .await
7394 .unwrap()
7395 .map(|transaction| transaction.uuid);
7396 assert_eq!(transaction_id, expected_transaction_id);
7397 let indices = dataset.load_indices().await.unwrap();
7398 assert!(indices.iter().any(|index| index.name == "users_id_idx"));
7399 }
7400
7401 #[tokio::test]
7402 async fn test_create_vector_index() {
7403 use lance_namespace::models::CreateTableIndexRequest;
7404
7405 let (namespace, _temp_dir) = create_test_namespace().await;
7406 create_vector_table(&namespace, "vectors").await;
7407
7408 let mut create_index_request =
7409 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
7410 create_index_request.id = Some(vec!["vectors".to_string()]);
7411 create_index_request.name = Some("vector_idx".to_string());
7412 create_index_request.distance_type = Some("l2".to_string());
7413 let transaction_id = namespace
7414 .create_table_index(create_index_request)
7415 .await
7416 .unwrap()
7417 .transaction_id;
7418
7419 let dataset = open_dataset(&namespace, "vectors").await;
7420 let expected_transaction_id = dataset
7421 .read_transaction()
7422 .await
7423 .unwrap()
7424 .map(|transaction| transaction.uuid);
7425 assert_eq!(transaction_id, expected_transaction_id);
7426 let indices = dataset.load_indices().await.unwrap();
7427 assert!(indices.iter().any(|index| index.name == "vector_idx"));
7428 }
7429
7430 #[tokio::test]
7431 async fn test_list_table_indices() {
7432 use lance_namespace::models::{CreateTableIndexRequest, ListTableIndicesRequest};
7433
7434 let (namespace, _temp_dir) = create_test_namespace().await;
7435 create_scalar_table(&namespace, "users").await;
7436 create_scalar_index(&namespace, "users", "a_idx").await;
7437 create_scalar_index(&namespace, "users", "b_idx").await;
7438 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
7439
7440 let response = namespace
7441 .list_table_indices(ListTableIndicesRequest {
7442 id: Some(vec!["users".to_string()]),
7443 ..Default::default()
7444 })
7445 .await
7446 .unwrap();
7447
7448 assert_eq!(response.indexes.len(), 3);
7449 assert_eq!(response.indexes[0].index_name, "a_idx");
7450 assert_eq!(response.indexes[1].index_name, "b_idx");
7451 assert_eq!(response.indexes[2].index_name, "users_id_idx");
7452 assert!(response.page_token.is_none());
7453 let users_id_idx = response
7454 .indexes
7455 .iter()
7456 .find(|index| index.index_name == "users_id_idx")
7457 .unwrap();
7458 assert_eq!(users_id_idx.columns, vec!["id"]);
7459 assert_eq!(users_id_idx.status, "SUCCEEDED");
7460
7461 assert_eq!(users_id_idx.index_type.as_deref(), Some("BTree"));
7463 assert!(
7464 users_id_idx
7465 .type_url
7466 .as_deref()
7467 .is_some_and(|s| !s.is_empty())
7468 );
7469 assert_eq!(users_id_idx.num_indexed_rows, Some(3));
7470 assert_eq!(users_id_idx.num_unindexed_rows, Some(0));
7471 assert_eq!(users_id_idx.num_segments, Some(1));
7472 assert!(users_id_idx.size_bytes.is_some_and(|size| size > 0));
7473 assert!(users_id_idx.created_at.is_some());
7474 assert!(users_id_idx.index_version.is_some());
7475 assert!(users_id_idx.index_details.is_some());
7476
7477 let dataset = open_dataset(&namespace, "users").await;
7478 let expected_transaction_id = dataset
7479 .read_transaction()
7480 .await
7481 .unwrap()
7482 .map(|transaction| transaction.uuid);
7483 assert_eq!(transaction_id, expected_transaction_id);
7484 let indices = dataset.load_indices().await.unwrap();
7485 assert_eq!(
7486 indices
7487 .iter()
7488 .filter(|index| index.name == "users_id_idx")
7489 .count(),
7490 1
7491 );
7492
7493 let first_page = namespace
7494 .list_table_indices(ListTableIndicesRequest {
7495 id: Some(vec!["users".to_string()]),
7496 limit: Some(2),
7497 ..Default::default()
7498 })
7499 .await
7500 .unwrap();
7501
7502 assert_eq!(first_page.indexes.len(), 2);
7503 assert_eq!(first_page.indexes[0].index_name, "a_idx");
7504 assert_eq!(first_page.indexes[1].index_name, "b_idx");
7505 assert_eq!(first_page.page_token.as_deref(), Some("b_idx"));
7506
7507 let second_page = namespace
7508 .list_table_indices(ListTableIndicesRequest {
7509 id: Some(vec!["users".to_string()]),
7510 page_token: first_page.page_token.clone(),
7511 limit: Some(2),
7512 ..Default::default()
7513 })
7514 .await
7515 .unwrap();
7516
7517 assert_eq!(second_page.indexes.len(), 1);
7518 assert_eq!(second_page.indexes[0].index_name, "users_id_idx");
7519 assert!(second_page.page_token.is_none());
7520
7521 create_vector_table(&namespace, "vectors").await;
7523 let mut create_index_request =
7524 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
7525 create_index_request.id = Some(vec!["vectors".to_string()]);
7526 create_index_request.name = Some("vector_idx".to_string());
7527 create_index_request.distance_type = Some("l2".to_string());
7528 namespace
7529 .create_table_index(create_index_request)
7530 .await
7531 .unwrap();
7532
7533 let vector_response = namespace
7534 .list_table_indices(ListTableIndicesRequest {
7535 id: Some(vec!["vectors".to_string()]),
7536 ..Default::default()
7537 })
7538 .await
7539 .unwrap();
7540
7541 assert_eq!(vector_response.indexes.len(), 1);
7542 let vector_idx = &vector_response.indexes[0];
7543 assert_eq!(vector_idx.index_name, "vector_idx");
7544 assert_eq!(vector_idx.columns, vec!["vector"]);
7545 assert_eq!(vector_idx.index_type.as_deref(), Some("IVF_FLAT"));
7546 assert!(
7547 vector_idx
7548 .type_url
7549 .as_deref()
7550 .is_some_and(|s| !s.is_empty())
7551 );
7552 assert!(vector_idx.num_indexed_rows.is_some());
7553 assert!(vector_idx.num_unindexed_rows.is_some());
7554 assert_eq!(vector_idx.num_segments, Some(1));
7555 assert!(vector_idx.created_at.is_some());
7556 assert!(vector_idx.index_version.is_some());
7557 assert!(vector_idx.index_details.is_some());
7558 }
7559
7560 #[tokio::test]
7561 async fn test_describe_table_index_stats() {
7562 use lance_namespace::models::DescribeTableIndexStatsRequest;
7563
7564 let (namespace, _temp_dir) = create_test_namespace().await;
7565 create_scalar_table(&namespace, "users").await;
7566 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
7567
7568 let response = namespace
7569 .describe_table_index_stats(DescribeTableIndexStatsRequest {
7570 id: Some(vec!["users".to_string()]),
7571 index_name: Some("users_id_idx".to_string()),
7572 ..Default::default()
7573 })
7574 .await
7575 .unwrap();
7576 assert_eq!(response.index_type, Some("BTree".to_string()));
7577 assert_eq!(response.num_indices, Some(1));
7578 assert_eq!(response.num_indexed_rows, Some(3));
7579 assert_eq!(response.num_unindexed_rows, Some(0));
7580
7581 let dataset = open_dataset(&namespace, "users").await;
7582 let expected_transaction_id = dataset
7583 .read_transaction()
7584 .await
7585 .unwrap()
7586 .map(|transaction| transaction.uuid);
7587 assert_eq!(transaction_id, expected_transaction_id);
7588 let stats: serde_json::Value =
7589 serde_json::from_str(&dataset.index_statistics("users_id_idx").await.unwrap()).unwrap();
7590 assert_eq!(stats["index_type"], "BTree");
7591 assert_eq!(stats["num_indices"], 1);
7592 assert_eq!(stats["num_indexed_rows"], 3);
7593 assert_eq!(stats["num_unindexed_rows"], 0);
7594 }
7595
7596 #[tokio::test]
7597 async fn test_describe_transaction() {
7598 use lance_namespace::models::DescribeTransactionRequest;
7599
7600 let (namespace, _temp_dir) = create_test_namespace().await;
7601 create_scalar_table(&namespace, "users").await;
7602 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
7603 let dataset = open_dataset(&namespace, "users").await;
7604 let latest_transaction = dataset.read_transaction().await.unwrap();
7605 assert_eq!(
7606 transaction_id,
7607 latest_transaction
7608 .as_ref()
7609 .map(|transaction| transaction.uuid.clone())
7610 );
7611
7612 if let Some(transaction_id) = transaction_id {
7613 let response = namespace
7614 .describe_transaction(DescribeTransactionRequest {
7615 id: Some(vec!["users".to_string(), transaction_id.clone()]),
7616 ..Default::default()
7617 })
7618 .await
7619 .unwrap();
7620 assert_eq!(response.status, "SUCCEEDED");
7621 assert_eq!(
7622 response
7623 .properties
7624 .as_ref()
7625 .and_then(|props| props.get("operation")),
7626 Some(&"CreateIndex".to_string())
7627 );
7628 assert_eq!(
7629 response
7630 .properties
7631 .as_ref()
7632 .and_then(|props| props.get("uuid")),
7633 Some(&transaction_id)
7634 );
7635 } else {
7636 assert!(latest_transaction.is_none());
7637 }
7638 }
7639
7640 #[tokio::test]
7641 async fn test_drop_table_index() {
7642 use lance_namespace::models::{DropTableIndexRequest, ListTableIndicesRequest};
7643
7644 let (namespace, _temp_dir) = create_test_namespace().await;
7645 create_scalar_table(&namespace, "users").await;
7646 let create_transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
7647
7648 let drop_transaction_id = namespace
7649 .drop_table_index(DropTableIndexRequest {
7650 id: Some(vec!["users".to_string()]),
7651 index_name: Some("users_id_idx".to_string()),
7652 ..Default::default()
7653 })
7654 .await
7655 .unwrap()
7656 .transaction_id;
7657
7658 let dataset = open_dataset(&namespace, "users").await;
7659 let previous_dataset = dataset
7660 .checkout_version(dataset.version().version - 1)
7661 .await
7662 .unwrap();
7663 let previous_transaction_id = previous_dataset
7664 .read_transaction()
7665 .await
7666 .unwrap()
7667 .map(|transaction| transaction.uuid);
7668 assert_eq!(create_transaction_id, previous_transaction_id);
7669 let expected_drop_transaction_id = dataset
7670 .read_transaction()
7671 .await
7672 .unwrap()
7673 .map(|transaction| transaction.uuid);
7674 assert_eq!(drop_transaction_id, expected_drop_transaction_id);
7675 let indices = dataset.load_indices().await.unwrap();
7676 assert!(!indices.iter().any(|index| index.name == "users_id_idx"));
7677
7678 let list_response = namespace
7679 .list_table_indices(ListTableIndicesRequest {
7680 id: Some(vec!["users".to_string()]),
7681 ..Default::default()
7682 })
7683 .await
7684 .unwrap();
7685 assert!(list_response.indexes.is_empty());
7686 }
7687
7688 #[tokio::test]
7689 async fn test_describe_table() {
7690 let (namespace, _temp_dir) = create_test_namespace().await;
7691
7692 let schema = create_test_schema();
7694 let ipc_data = create_test_ipc_data(&schema);
7695
7696 let mut create_request = CreateTableRequest::new();
7697 create_request.id = Some(vec!["test_table".to_string()]);
7698 namespace
7699 .create_table(create_request, bytes::Bytes::from(ipc_data))
7700 .await
7701 .unwrap();
7702
7703 let mut request = DescribeTableRequest::new();
7705 request.id = Some(vec!["test_table".to_string()]);
7706 let response = namespace.describe_table(request).await.unwrap();
7707
7708 assert!(response.location.is_some());
7709 assert!(response.location.unwrap().ends_with("test_table.lance"));
7710 }
7711
7712 #[tokio::test]
7713 async fn test_describe_nonexistent_table() {
7714 let (namespace, _temp_dir) = create_test_namespace().await;
7715
7716 let mut request = DescribeTableRequest::new();
7717 request.id = Some(vec!["nonexistent".to_string()]);
7718
7719 let result = namespace.describe_table(request).await;
7720 assert!(result.is_err());
7721 assert!(result.unwrap_err().to_string().contains("Table not found"));
7722 }
7723
7724 #[tokio::test]
7725 async fn test_table_exists() {
7726 let (namespace, _temp_dir) = create_test_namespace().await;
7727
7728 let schema = create_test_schema();
7730 let ipc_data = create_test_ipc_data(&schema);
7731
7732 let mut create_request = CreateTableRequest::new();
7733 create_request.id = Some(vec!["existing_table".to_string()]);
7734 namespace
7735 .create_table(create_request, bytes::Bytes::from(ipc_data))
7736 .await
7737 .unwrap();
7738
7739 let mut request = TableExistsRequest::new();
7741 request.id = Some(vec!["existing_table".to_string()]);
7742 let result = namespace.table_exists(request).await;
7743 assert!(result.is_ok());
7744
7745 let mut request = TableExistsRequest::new();
7747 request.id = Some(vec!["nonexistent".to_string()]);
7748 let result = namespace.table_exists(request).await;
7749 assert!(result.is_err());
7750 assert!(result.unwrap_err().to_string().contains("Table not found"));
7751 }
7752
7753 #[tokio::test]
7754 async fn test_drop_table() {
7755 let (namespace, _temp_dir) = create_test_namespace().await;
7756
7757 let schema = create_test_schema();
7759 let ipc_data = create_test_ipc_data(&schema);
7760
7761 let mut create_request = CreateTableRequest::new();
7762 create_request.id = Some(vec!["table_to_drop".to_string()]);
7763 namespace
7764 .create_table(create_request, bytes::Bytes::from(ipc_data))
7765 .await
7766 .unwrap();
7767
7768 let mut exists_request = TableExistsRequest::new();
7770 exists_request.id = Some(vec!["table_to_drop".to_string()]);
7771 assert!(namespace.table_exists(exists_request.clone()).await.is_ok());
7772
7773 let mut drop_request = DropTableRequest::new();
7775 drop_request.id = Some(vec!["table_to_drop".to_string()]);
7776 let response = namespace.drop_table(drop_request).await.unwrap();
7777 assert!(response.location.is_some());
7778
7779 assert!(namespace.table_exists(exists_request).await.is_err());
7781 }
7782
7783 #[tokio::test]
7784 async fn test_drop_nonexistent_table() {
7785 let (namespace, _temp_dir) = create_test_namespace().await;
7786
7787 let mut request = DropTableRequest::new();
7788 request.id = Some(vec!["nonexistent".to_string()]);
7789
7790 let result = namespace.drop_table(request).await;
7792 let _ = result;
7795 }
7796
7797 #[tokio::test]
7798 async fn test_root_namespace_operations() {
7799 let (namespace, _temp_dir) = create_test_namespace().await;
7800
7801 let mut request = ListNamespacesRequest::new();
7803 request.id = Some(vec![]);
7804 let result = namespace.list_namespaces(request).await;
7805 assert!(result.is_ok());
7806 assert_eq!(result.unwrap().namespaces.len(), 0);
7807
7808 let mut request = DescribeNamespaceRequest::new();
7810 request.id = Some(vec![]);
7811 let result = namespace.describe_namespace(request).await;
7812 assert!(result.is_ok());
7813
7814 let mut request = NamespaceExistsRequest::new();
7816 request.id = Some(vec![]);
7817 let result = namespace.namespace_exists(request).await;
7818 assert!(result.is_ok());
7819
7820 let mut request = CreateNamespaceRequest::new();
7822 request.id = Some(vec![]);
7823 let result = namespace.create_namespace(request).await;
7824 assert!(result.is_err());
7825 assert!(result.unwrap_err().to_string().contains("already exists"));
7826
7827 let mut request = DropNamespaceRequest::new();
7829 request.id = Some(vec![]);
7830 let result = namespace.drop_namespace(request).await;
7831 assert!(result.is_err());
7832 assert!(
7833 result
7834 .unwrap_err()
7835 .to_string()
7836 .contains("cannot be dropped")
7837 );
7838 }
7839
7840 #[tokio::test]
7841 async fn test_non_root_namespace_operations() {
7842 let (namespace, _temp_dir) = create_test_namespace().await;
7843
7844 let mut request = CreateNamespaceRequest::new();
7847 request.id = Some(vec!["child".to_string()]);
7848 let result = namespace.create_namespace(request).await;
7849 assert!(
7850 result.is_ok(),
7851 "Child namespace creation should succeed with manifest enabled"
7852 );
7853
7854 let mut request = NamespaceExistsRequest::new();
7856 request.id = Some(vec!["child".to_string()]);
7857 let result = namespace.namespace_exists(request).await;
7858 assert!(
7859 result.is_ok(),
7860 "Child namespace should exist after creation"
7861 );
7862
7863 let mut request = DropNamespaceRequest::new();
7865 request.id = Some(vec!["child".to_string()]);
7866 let result = namespace.drop_namespace(request).await;
7867 assert!(
7868 result.is_ok(),
7869 "Child namespace drop should succeed with manifest enabled"
7870 );
7871
7872 let mut request = NamespaceExistsRequest::new();
7874 request.id = Some(vec!["child".to_string()]);
7875 let result = namespace.namespace_exists(request).await;
7876 assert!(
7877 result.is_err(),
7878 "Child namespace should not exist after drop"
7879 );
7880 }
7881
7882 #[tokio::test]
7883 async fn test_config_custom_root() {
7884 let temp_dir = TempStdDir::default();
7885 let custom_path = temp_dir.join("custom");
7886 std::fs::create_dir(&custom_path).unwrap();
7887
7888 let namespace = DirectoryNamespaceBuilder::new(custom_path.to_string_lossy().to_string())
7889 .build()
7890 .await
7891 .unwrap();
7892
7893 let schema = create_test_schema();
7895 let ipc_data = create_test_ipc_data(&schema);
7896
7897 let mut request = CreateTableRequest::new();
7899 request.id = Some(vec!["test_table".to_string()]);
7900
7901 let response = namespace
7902 .create_table(request, bytes::Bytes::from(ipc_data))
7903 .await
7904 .unwrap();
7905
7906 assert!(response.location.unwrap().contains("custom"));
7907 }
7908
7909 #[tokio::test]
7910 async fn test_config_storage_options() {
7911 let temp_dir = TempStdDir::default();
7912
7913 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
7914 .storage_option("option1", "value1")
7915 .storage_option("option2", "value2")
7916 .build()
7917 .await
7918 .unwrap();
7919
7920 let schema = create_test_schema();
7922 let ipc_data = create_test_ipc_data(&schema);
7923
7924 let mut request = CreateTableRequest::new();
7926 request.id = Some(vec!["test_table".to_string()]);
7927
7928 let response = namespace
7929 .create_table(request, bytes::Bytes::from(ipc_data))
7930 .await
7931 .unwrap();
7932
7933 let storage_options = response.storage_options.unwrap();
7934 assert_eq!(storage_options.get("option1"), Some(&"value1".to_string()));
7935 assert_eq!(storage_options.get("option2"), Some(&"value2".to_string()));
7936 }
7937
7938 #[tokio::test]
7942 async fn test_no_storage_options_without_vendor() {
7943 use lance_namespace::models::DeclareTableRequest;
7944
7945 let temp_dir = TempStdDir::default();
7946
7947 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
7949 .manifest_enabled(false)
7950 .storage_option("aws_access_key_id", "AKID")
7951 .storage_option("aws_secret_access_key", "SECRET")
7952 .storage_option("region", "us-east-1")
7953 .build()
7954 .await
7955 .unwrap();
7956
7957 let schema = create_test_schema();
7958 let ipc_data = create_test_ipc_data(&schema);
7959
7960 let mut create_req = CreateTableRequest::new();
7962 create_req.id = Some(vec!["t1".to_string()]);
7963 namespace
7964 .create_table(create_req, bytes::Bytes::from(ipc_data))
7965 .await
7966 .unwrap();
7967
7968 let mut desc_req = DescribeTableRequest::new();
7970 desc_req.id = Some(vec!["t1".to_string()]);
7971 let resp = namespace.describe_table(desc_req).await.unwrap();
7972 assert!(resp.storage_options.is_none());
7973
7974 let mut decl_req = DeclareTableRequest::new();
7976 decl_req.id = Some(vec!["t2".to_string()]);
7977 let resp = namespace.declare_table(decl_req).await.unwrap();
7978 assert!(resp.storage_options.is_none());
7979 }
7980
7981 #[tokio::test]
7983 async fn test_no_storage_options_without_vendor_manifest() {
7984 let temp_dir = TempStdDir::default();
7985
7986 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
7987 .storage_option("aws_access_key_id", "AKID")
7988 .storage_option("aws_secret_access_key", "SECRET")
7989 .storage_option("region", "us-east-1")
7990 .build()
7991 .await
7992 .unwrap();
7993
7994 let schema = create_test_schema();
7995 let ipc_data = create_test_ipc_data(&schema);
7996
7997 let mut create_req = CreateTableRequest::new();
7998 create_req.id = Some(vec!["t1".to_string()]);
7999 namespace
8000 .create_table(create_req, bytes::Bytes::from(ipc_data))
8001 .await
8002 .unwrap();
8003
8004 let mut desc_req = DescribeTableRequest::new();
8006 desc_req.id = Some(vec!["t1".to_string()]);
8007 let resp = namespace.describe_table(desc_req).await.unwrap();
8008 assert!(resp.storage_options.is_none());
8009 }
8010
8011 #[tokio::test]
8012 async fn test_from_properties_manifest_enabled() {
8013 let temp_dir = TempStdDir::default();
8014
8015 let mut properties = HashMap::new();
8016 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
8017 properties.insert("manifest_enabled".to_string(), "true".to_string());
8018 properties.insert("dir_listing_enabled".to_string(), "false".to_string());
8019
8020 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
8021 assert!(builder.manifest_enabled);
8022 assert!(!builder.dir_listing_enabled);
8023
8024 let namespace = builder.build().await.unwrap();
8025
8026 let schema = create_test_schema();
8028 let ipc_data = create_test_ipc_data(&schema);
8029
8030 let mut request = CreateTableRequest::new();
8032 request.id = Some(vec!["test_table".to_string()]);
8033
8034 let response = namespace
8035 .create_table(request, bytes::Bytes::from(ipc_data))
8036 .await
8037 .unwrap();
8038
8039 assert!(response.location.is_some());
8040 }
8041
8042 #[tokio::test]
8043 async fn test_from_properties_dir_listing_enabled() {
8044 let temp_dir = TempStdDir::default();
8045
8046 let mut properties = HashMap::new();
8047 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
8048 properties.insert("manifest_enabled".to_string(), "false".to_string());
8049 properties.insert("dir_listing_enabled".to_string(), "true".to_string());
8050
8051 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
8052 assert!(!builder.manifest_enabled);
8053 assert!(builder.dir_listing_enabled);
8054
8055 let namespace = builder.build().await.unwrap();
8056
8057 let schema = create_test_schema();
8059 let ipc_data = create_test_ipc_data(&schema);
8060
8061 let mut request = CreateTableRequest::new();
8063 request.id = Some(vec!["test_table".to_string()]);
8064
8065 let response = namespace
8066 .create_table(request, bytes::Bytes::from(ipc_data))
8067 .await
8068 .unwrap();
8069
8070 assert!(response.location.is_some());
8071 }
8072
8073 #[tokio::test]
8074 async fn test_from_properties_defaults() {
8075 let temp_dir = TempStdDir::default();
8076
8077 let mut properties = HashMap::new();
8078 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
8079
8080 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
8081 assert!(builder.manifest_enabled);
8083 assert!(builder.dir_listing_enabled);
8084 }
8085
8086 #[tokio::test]
8087 async fn test_from_properties_with_storage_options() {
8088 let temp_dir = TempStdDir::default();
8089
8090 let mut properties = HashMap::new();
8091 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
8092 properties.insert("manifest_enabled".to_string(), "true".to_string());
8093 properties.insert("storage.region".to_string(), "us-west-2".to_string());
8094 properties.insert("storage.bucket".to_string(), "my-bucket".to_string());
8095
8096 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
8097 assert!(builder.manifest_enabled);
8098 assert!(builder.storage_options.is_some());
8099
8100 let storage_options = builder.storage_options.unwrap();
8101 assert_eq!(
8102 storage_options.get("region"),
8103 Some(&"us-west-2".to_string())
8104 );
8105 assert_eq!(
8106 storage_options.get("bucket"),
8107 Some(&"my-bucket".to_string())
8108 );
8109 }
8110
8111 #[tokio::test]
8112 async fn test_various_arrow_types() {
8113 let (namespace, _temp_dir) = create_test_namespace().await;
8114
8115 let fields = vec![
8117 JsonArrowField {
8118 name: "bool_col".to_string(),
8119 r#type: Box::new(JsonArrowDataType::new("bool".to_string())),
8120 nullable: true,
8121 metadata: None,
8122 },
8123 JsonArrowField {
8124 name: "int8_col".to_string(),
8125 r#type: Box::new(JsonArrowDataType::new("int8".to_string())),
8126 nullable: true,
8127 metadata: None,
8128 },
8129 JsonArrowField {
8130 name: "float64_col".to_string(),
8131 r#type: Box::new(JsonArrowDataType::new("float64".to_string())),
8132 nullable: true,
8133 metadata: None,
8134 },
8135 JsonArrowField {
8136 name: "binary_col".to_string(),
8137 r#type: Box::new(JsonArrowDataType::new("binary".to_string())),
8138 nullable: true,
8139 metadata: None,
8140 },
8141 ];
8142
8143 let schema = JsonArrowSchema {
8144 fields,
8145 metadata: None,
8146 };
8147
8148 let ipc_data = create_test_ipc_data(&schema);
8150
8151 let mut request = CreateTableRequest::new();
8152 request.id = Some(vec!["complex_table".to_string()]);
8153
8154 let response = namespace
8155 .create_table(request, bytes::Bytes::from(ipc_data))
8156 .await
8157 .unwrap();
8158
8159 assert!(response.location.is_some());
8160 }
8161
8162 #[tokio::test]
8163 async fn test_connect_dir() {
8164 let temp_dir = TempStdDir::default();
8165
8166 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
8167 .build()
8168 .await
8169 .unwrap();
8170
8171 let mut request = ListTablesRequest::new();
8173 request.id = Some(vec![]);
8174 let response = namespace.list_tables(request).await.unwrap();
8175 assert_eq!(response.tables.len(), 0);
8176 }
8177
8178 #[tokio::test]
8179 async fn test_create_table_with_ipc_data() {
8180 use arrow::array::{Int32Array, StringArray};
8181 use arrow::ipc::writer::StreamWriter;
8182
8183 let (namespace, _temp_dir) = create_test_namespace().await;
8184
8185 let schema = create_test_schema();
8187
8188 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8190 let arrow_schema = Arc::new(arrow_schema);
8191
8192 let id_array = Int32Array::from(vec![1, 2, 3]);
8194 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
8195 let batch = arrow::record_batch::RecordBatch::try_new(
8196 arrow_schema.clone(),
8197 vec![Arc::new(id_array), Arc::new(name_array)],
8198 )
8199 .unwrap();
8200
8201 let mut buffer = Vec::new();
8203 {
8204 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8205 writer.write(&batch).unwrap();
8206 writer.finish().unwrap();
8207 }
8208
8209 let mut request = CreateTableRequest::new();
8211 request.id = Some(vec!["test_table_with_data".to_string()]);
8212
8213 let response = namespace
8214 .create_table(request, Bytes::from(buffer))
8215 .await
8216 .unwrap();
8217
8218 assert_eq!(response.version, Some(1));
8219 assert!(
8220 response
8221 .location
8222 .unwrap()
8223 .contains("test_table_with_data.lance")
8224 );
8225
8226 let mut exists_request = TableExistsRequest::new();
8228 exists_request.id = Some(vec!["test_table_with_data".to_string()]);
8229 namespace.table_exists(exists_request).await.unwrap();
8230 }
8231
8232 #[tokio::test]
8233 async fn test_child_namespace_create_and_list() {
8234 let (namespace, _temp_dir) = create_test_namespace().await;
8235
8236 for i in 1..=3 {
8238 let mut create_req = CreateNamespaceRequest::new();
8239 create_req.id = Some(vec![format!("ns{}", i)]);
8240 let result = namespace.create_namespace(create_req).await;
8241 assert!(result.is_ok(), "Failed to create child namespace ns{}", i);
8242 }
8243
8244 let list_req = ListNamespacesRequest {
8246 id: Some(vec![]),
8247 ..Default::default()
8248 };
8249 let result = namespace.list_namespaces(list_req).await;
8250 assert!(result.is_ok());
8251 let namespaces = result.unwrap().namespaces;
8252 assert_eq!(namespaces.len(), 3);
8253 assert!(namespaces.contains(&"ns1".to_string()));
8254 assert!(namespaces.contains(&"ns2".to_string()));
8255 assert!(namespaces.contains(&"ns3".to_string()));
8256 }
8257
8258 #[tokio::test]
8259 async fn test_nested_namespace_hierarchy() {
8260 let (namespace, _temp_dir) = create_test_namespace().await;
8261
8262 let mut create_req = CreateNamespaceRequest::new();
8264 create_req.id = Some(vec!["parent".to_string()]);
8265 namespace.create_namespace(create_req).await.unwrap();
8266
8267 let mut create_req = CreateNamespaceRequest::new();
8269 create_req.id = Some(vec!["parent".to_string(), "child1".to_string()]);
8270 namespace.create_namespace(create_req).await.unwrap();
8271
8272 let mut create_req = CreateNamespaceRequest::new();
8273 create_req.id = Some(vec!["parent".to_string(), "child2".to_string()]);
8274 namespace.create_namespace(create_req).await.unwrap();
8275
8276 let list_req = ListNamespacesRequest {
8278 id: Some(vec!["parent".to_string()]),
8279 ..Default::default()
8280 };
8281 let result = namespace.list_namespaces(list_req).await;
8282 assert!(result.is_ok());
8283 let children = result.unwrap().namespaces;
8284 assert_eq!(children.len(), 2);
8285 assert!(children.contains(&"child1".to_string()));
8286 assert!(children.contains(&"child2".to_string()));
8287
8288 let list_req = ListNamespacesRequest {
8290 id: Some(vec![]),
8291 ..Default::default()
8292 };
8293 let result = namespace.list_namespaces(list_req).await;
8294 assert!(result.is_ok());
8295 let root_namespaces = result.unwrap().namespaces;
8296 assert_eq!(root_namespaces.len(), 1);
8297 assert_eq!(root_namespaces[0], "parent");
8298 }
8299
8300 #[tokio::test]
8301 async fn test_table_in_child_namespace() {
8302 let (namespace, _temp_dir) = create_test_namespace().await;
8303
8304 let mut create_ns_req = CreateNamespaceRequest::new();
8306 create_ns_req.id = Some(vec!["test_ns".to_string()]);
8307 namespace.create_namespace(create_ns_req).await.unwrap();
8308
8309 let schema = create_test_schema();
8311 let ipc_data = create_test_ipc_data(&schema);
8312 let mut create_table_req = CreateTableRequest::new();
8313 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8314 let result = namespace
8315 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
8316 .await;
8317 assert!(result.is_ok(), "Failed to create table in child namespace");
8318
8319 let list_req = ListTablesRequest {
8321 id: Some(vec!["test_ns".to_string()]),
8322 ..Default::default()
8323 };
8324 let result = namespace.list_tables(list_req).await;
8325 assert!(result.is_ok());
8326 let tables = result.unwrap().tables;
8327 assert_eq!(tables.len(), 1);
8328 assert_eq!(tables[0], "table1");
8329
8330 let mut exists_req = TableExistsRequest::new();
8332 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8333 let result = namespace.table_exists(exists_req).await;
8334 assert!(result.is_ok());
8335
8336 let mut describe_req = DescribeTableRequest::new();
8338 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8339 let result = namespace.describe_table(describe_req).await;
8340 assert!(result.is_ok());
8341 let response = result.unwrap();
8342 assert!(response.location.is_some());
8343 }
8344
8345 #[tokio::test]
8346 async fn test_multiple_tables_in_child_namespace() {
8347 let (namespace, _temp_dir) = create_test_namespace().await;
8348
8349 let mut create_ns_req = CreateNamespaceRequest::new();
8351 create_ns_req.id = Some(vec!["test_ns".to_string()]);
8352 namespace.create_namespace(create_ns_req).await.unwrap();
8353
8354 let schema = create_test_schema();
8356 let ipc_data = create_test_ipc_data(&schema);
8357 for i in 1..=3 {
8358 let mut create_table_req = CreateTableRequest::new();
8359 create_table_req.id = Some(vec!["test_ns".to_string(), format!("table{}", i)]);
8360 namespace
8361 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
8362 .await
8363 .unwrap();
8364 }
8365
8366 let list_req = ListTablesRequest {
8368 id: Some(vec!["test_ns".to_string()]),
8369 ..Default::default()
8370 };
8371 let result = namespace.list_tables(list_req).await;
8372 assert!(result.is_ok());
8373 let tables = result.unwrap().tables;
8374 assert_eq!(tables.len(), 3);
8375 assert!(tables.contains(&"table1".to_string()));
8376 assert!(tables.contains(&"table2".to_string()));
8377 assert!(tables.contains(&"table3".to_string()));
8378 }
8379
8380 #[tokio::test]
8381 async fn test_drop_table_in_child_namespace() {
8382 let (namespace, _temp_dir) = create_test_namespace().await;
8383
8384 let mut create_ns_req = CreateNamespaceRequest::new();
8386 create_ns_req.id = Some(vec!["test_ns".to_string()]);
8387 namespace.create_namespace(create_ns_req).await.unwrap();
8388
8389 let schema = create_test_schema();
8391 let ipc_data = create_test_ipc_data(&schema);
8392 let mut create_table_req = CreateTableRequest::new();
8393 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8394 namespace
8395 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
8396 .await
8397 .unwrap();
8398
8399 let mut drop_req = DropTableRequest::new();
8401 drop_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8402 let result = namespace.drop_table(drop_req).await;
8403 assert!(result.is_ok(), "Failed to drop table in child namespace");
8404
8405 let mut exists_req = TableExistsRequest::new();
8407 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8408 let result = namespace.table_exists(exists_req).await;
8409 assert!(result.is_err());
8410 }
8411
8412 #[tokio::test]
8413 async fn test_deeply_nested_namespace() {
8414 let (namespace, _temp_dir) = create_test_namespace().await;
8415
8416 let mut create_req = CreateNamespaceRequest::new();
8418 create_req.id = Some(vec!["level1".to_string()]);
8419 namespace.create_namespace(create_req).await.unwrap();
8420
8421 let mut create_req = CreateNamespaceRequest::new();
8422 create_req.id = Some(vec!["level1".to_string(), "level2".to_string()]);
8423 namespace.create_namespace(create_req).await.unwrap();
8424
8425 let mut create_req = CreateNamespaceRequest::new();
8426 create_req.id = Some(vec![
8427 "level1".to_string(),
8428 "level2".to_string(),
8429 "level3".to_string(),
8430 ]);
8431 namespace.create_namespace(create_req).await.unwrap();
8432
8433 let schema = create_test_schema();
8435 let ipc_data = create_test_ipc_data(&schema);
8436 let mut create_table_req = CreateTableRequest::new();
8437 create_table_req.id = Some(vec![
8438 "level1".to_string(),
8439 "level2".to_string(),
8440 "level3".to_string(),
8441 "table1".to_string(),
8442 ]);
8443 let result = namespace
8444 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
8445 .await;
8446 assert!(
8447 result.is_ok(),
8448 "Failed to create table in deeply nested namespace"
8449 );
8450
8451 let mut exists_req = TableExistsRequest::new();
8453 exists_req.id = Some(vec![
8454 "level1".to_string(),
8455 "level2".to_string(),
8456 "level3".to_string(),
8457 "table1".to_string(),
8458 ]);
8459 let result = namespace.table_exists(exists_req).await;
8460 assert!(result.is_ok());
8461 }
8462
8463 #[tokio::test]
8464 async fn test_namespace_with_properties() {
8465 let (namespace, _temp_dir) = create_test_namespace().await;
8466
8467 let mut properties = HashMap::new();
8469 properties.insert("owner".to_string(), "test_user".to_string());
8470 properties.insert("description".to_string(), "Test namespace".to_string());
8471
8472 let mut create_req = CreateNamespaceRequest::new();
8473 create_req.id = Some(vec!["test_ns".to_string()]);
8474 create_req.properties = Some(properties.clone());
8475 namespace.create_namespace(create_req).await.unwrap();
8476
8477 let describe_req = DescribeNamespaceRequest {
8479 id: Some(vec!["test_ns".to_string()]),
8480 ..Default::default()
8481 };
8482 let result = namespace.describe_namespace(describe_req).await;
8483 assert!(result.is_ok());
8484 let response = result.unwrap();
8485 assert!(response.properties.is_some());
8486 let props = response.properties.unwrap();
8487 assert_eq!(props.get("owner"), Some(&"test_user".to_string()));
8488 assert_eq!(
8489 props.get("description"),
8490 Some(&"Test namespace".to_string())
8491 );
8492 }
8493
8494 #[tokio::test]
8495 async fn test_cannot_drop_namespace_with_tables() {
8496 let (namespace, _temp_dir) = create_test_namespace().await;
8497
8498 let mut create_ns_req = CreateNamespaceRequest::new();
8500 create_ns_req.id = Some(vec!["test_ns".to_string()]);
8501 namespace.create_namespace(create_ns_req).await.unwrap();
8502
8503 let schema = create_test_schema();
8505 let ipc_data = create_test_ipc_data(&schema);
8506 let mut create_table_req = CreateTableRequest::new();
8507 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
8508 namespace
8509 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
8510 .await
8511 .unwrap();
8512
8513 let mut drop_req = DropNamespaceRequest::new();
8515 drop_req.id = Some(vec!["test_ns".to_string()]);
8516 let result = namespace.drop_namespace(drop_req).await;
8517 assert!(
8518 result.is_err(),
8519 "Should not be able to drop namespace with tables"
8520 );
8521 }
8522
8523 #[tokio::test]
8524 async fn test_isolation_between_namespaces() {
8525 let (namespace, _temp_dir) = create_test_namespace().await;
8526
8527 let mut create_req = CreateNamespaceRequest::new();
8529 create_req.id = Some(vec!["ns1".to_string()]);
8530 namespace.create_namespace(create_req).await.unwrap();
8531
8532 let mut create_req = CreateNamespaceRequest::new();
8533 create_req.id = Some(vec!["ns2".to_string()]);
8534 namespace.create_namespace(create_req).await.unwrap();
8535
8536 let schema = create_test_schema();
8538 let ipc_data = create_test_ipc_data(&schema);
8539
8540 let mut create_table_req = CreateTableRequest::new();
8541 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
8542 namespace
8543 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
8544 .await
8545 .unwrap();
8546
8547 let mut create_table_req = CreateTableRequest::new();
8548 create_table_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
8549 namespace
8550 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
8551 .await
8552 .unwrap();
8553
8554 let list_req = ListTablesRequest {
8556 id: Some(vec!["ns1".to_string()]),
8557 page_token: None,
8558 limit: None,
8559 ..Default::default()
8560 };
8561 let result = namespace.list_tables(list_req).await.unwrap();
8562 assert_eq!(result.tables.len(), 1);
8563 assert_eq!(result.tables[0], "table1");
8564
8565 let list_req = ListTablesRequest {
8566 id: Some(vec!["ns2".to_string()]),
8567 page_token: None,
8568 limit: None,
8569 ..Default::default()
8570 };
8571 let result = namespace.list_tables(list_req).await.unwrap();
8572 assert_eq!(result.tables.len(), 1);
8573 assert_eq!(result.tables[0], "table1");
8574
8575 let mut drop_req = DropTableRequest::new();
8577 drop_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
8578 namespace.drop_table(drop_req).await.unwrap();
8579
8580 let mut exists_req = TableExistsRequest::new();
8582 exists_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
8583 assert!(namespace.table_exists(exists_req).await.is_err());
8584
8585 let mut exists_req = TableExistsRequest::new();
8586 exists_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
8587 assert!(namespace.table_exists(exists_req).await.is_ok());
8588 }
8589
8590 #[tokio::test]
8591 async fn test_migrate_directory_tables() {
8592 let temp_dir = TempStdDir::default();
8593 let temp_path = temp_dir.to_str().unwrap();
8594
8595 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
8597 .manifest_enabled(false)
8598 .dir_listing_enabled(true)
8599 .build()
8600 .await
8601 .unwrap();
8602
8603 let schema = create_test_schema();
8605 let ipc_data = create_test_ipc_data(&schema);
8606
8607 for i in 1..=3 {
8608 let mut create_req = CreateTableRequest::new();
8609 create_req.id = Some(vec![format!("table{}", i)]);
8610 dir_only_ns
8611 .create_table(create_req, bytes::Bytes::from(ipc_data.clone()))
8612 .await
8613 .unwrap();
8614 }
8615
8616 drop(dir_only_ns);
8617
8618 let dual_mode_ns = DirectoryNamespaceBuilder::new(temp_path)
8620 .manifest_enabled(true)
8621 .dir_listing_enabled(true)
8622 .build()
8623 .await
8624 .unwrap();
8625
8626 let mut list_req = ListTablesRequest::new();
8628 list_req.id = Some(vec![]);
8629 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
8630 assert_eq!(tables.len(), 3);
8631
8632 let migrated_count = dual_mode_ns.migrate().await.unwrap();
8634 assert_eq!(migrated_count, 3, "Should migrate all 3 tables");
8635
8636 let mut list_req = ListTablesRequest::new();
8638 list_req.id = Some(vec![]);
8639 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
8640 assert_eq!(tables.len(), 3);
8641
8642 let migrated_count = dual_mode_ns.migrate().await.unwrap();
8644 assert_eq!(
8645 migrated_count, 0,
8646 "Should not migrate already-migrated tables"
8647 );
8648
8649 drop(dual_mode_ns);
8650
8651 let manifest_only_ns = DirectoryNamespaceBuilder::new(temp_path)
8653 .manifest_enabled(true)
8654 .dir_listing_enabled(false)
8655 .build()
8656 .await
8657 .unwrap();
8658
8659 let mut list_req = ListTablesRequest::new();
8661 list_req.id = Some(vec![]);
8662 let tables = manifest_only_ns.list_tables(list_req).await.unwrap().tables;
8663 assert_eq!(tables.len(), 3);
8664 assert!(tables.contains(&"table1".to_string()));
8665 assert!(tables.contains(&"table2".to_string()));
8666 assert!(tables.contains(&"table3".to_string()));
8667 }
8668
8669 #[tokio::test]
8670 async fn test_migrate_without_manifest() {
8671 let temp_dir = TempStdDir::default();
8672 let temp_path = temp_dir.to_str().unwrap();
8673
8674 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8676 .manifest_enabled(false)
8677 .dir_listing_enabled(true)
8678 .build()
8679 .await
8680 .unwrap();
8681
8682 let migrated_count = namespace.migrate().await.unwrap();
8684 assert_eq!(migrated_count, 0);
8685 }
8686
8687 #[tokio::test]
8688 async fn test_register_table() {
8689 use lance_namespace::models::{RegisterTableRequest, TableExistsRequest};
8690
8691 let temp_dir = TempStdDir::default();
8692 let temp_path = temp_dir.to_str().unwrap();
8693
8694 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8695 .dir_listing_to_manifest_migration_enabled(true)
8696 .build()
8697 .await
8698 .unwrap();
8699
8700 let schema = create_test_schema();
8702 let ipc_data = create_test_ipc_data(&schema);
8703
8704 let table_uri = format!("{}/external_table.lance", temp_path);
8705 let cursor = Cursor::new(ipc_data);
8706 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
8707 let batches: Vec<_> = stream_reader
8708 .collect::<std::result::Result<Vec<_>, _>>()
8709 .unwrap();
8710 let schema = batches[0].schema();
8711 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
8712 let reader = RecordBatchIterator::new(batch_results, schema);
8713 Dataset::write(Box::new(reader), &table_uri, None)
8714 .await
8715 .unwrap();
8716
8717 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
8719 register_req.id = Some(vec!["registered_table".to_string()]);
8720
8721 let response = namespace.register_table(register_req).await.unwrap();
8722 assert_eq!(response.location, Some("external_table.lance".to_string()));
8723
8724 let mut exists_req = TableExistsRequest::new();
8726 exists_req.id = Some(vec!["registered_table".to_string()]);
8727 assert!(namespace.table_exists(exists_req).await.is_ok());
8728
8729 let mut list_req = ListTablesRequest::new();
8731 list_req.id = Some(vec![]);
8732 let tables = namespace.list_tables(list_req).await.unwrap();
8733 assert!(tables.tables.contains(&"registered_table".to_string()));
8734 }
8735
8736 #[tokio::test]
8737 async fn test_register_table_duplicate_fails() {
8738 use lance_namespace::models::RegisterTableRequest;
8739
8740 let temp_dir = TempStdDir::default();
8741 let temp_path = temp_dir.to_str().unwrap();
8742
8743 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8744 .build()
8745 .await
8746 .unwrap();
8747
8748 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
8750 register_req.id = Some(vec!["test_table".to_string()]);
8751
8752 namespace
8753 .register_table(register_req.clone())
8754 .await
8755 .unwrap();
8756
8757 let result = namespace.register_table(register_req).await;
8759 assert!(result.is_err());
8760 assert!(result.unwrap_err().to_string().contains("already exists"));
8761 }
8762
8763 #[tokio::test]
8764 async fn test_deregister_table() {
8765 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
8766
8767 let temp_dir = TempStdDir::default();
8768 let temp_path = temp_dir.to_str().unwrap();
8769
8770 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8773 .manifest_enabled(true)
8774 .dir_listing_enabled(false)
8775 .build()
8776 .await
8777 .unwrap();
8778
8779 let schema = create_test_schema();
8781 let ipc_data = create_test_ipc_data(&schema);
8782
8783 let mut create_req = CreateTableRequest::new();
8784 create_req.id = Some(vec!["test_table".to_string()]);
8785 namespace
8786 .create_table(create_req, bytes::Bytes::from(ipc_data))
8787 .await
8788 .unwrap();
8789
8790 let mut exists_req = TableExistsRequest::new();
8792 exists_req.id = Some(vec!["test_table".to_string()]);
8793 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
8794
8795 let mut deregister_req = DeregisterTableRequest::new();
8797 deregister_req.id = Some(vec!["test_table".to_string()]);
8798 let response = namespace.deregister_table(deregister_req).await.unwrap();
8799
8800 assert!(
8802 response.location.is_some(),
8803 "Deregister should return location"
8804 );
8805 let location = response.location.as_ref().unwrap();
8806 let expected_url = lance_io::object_store::uri_to_url(temp_path)
8809 .expect("Failed to convert temp path to URL");
8810 let expected_prefix = expected_url.to_string();
8811 assert!(
8812 location.starts_with(&expected_prefix),
8813 "Location should start with '{}', got: {}",
8814 expected_prefix,
8815 location
8816 );
8817 assert!(
8818 location.contains("test_table"),
8819 "Location should contain table name: {}",
8820 location
8821 );
8822 assert_eq!(response.id, Some(vec!["test_table".to_string()]));
8823
8824 assert!(namespace.table_exists(exists_req).await.is_err());
8826
8827 let dataset = Dataset::open(location).await;
8829 assert!(
8830 dataset.is_ok(),
8831 "Physical table data should still exist at {}",
8832 location
8833 );
8834 }
8835
8836 #[tokio::test]
8837 async fn test_deregister_table_in_child_namespace() {
8838 use lance_namespace::models::{
8839 CreateNamespaceRequest, DeregisterTableRequest, TableExistsRequest,
8840 };
8841
8842 let temp_dir = TempStdDir::default();
8843 let temp_path = temp_dir.to_str().unwrap();
8844
8845 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8846 .build()
8847 .await
8848 .unwrap();
8849
8850 let mut create_ns_req = CreateNamespaceRequest::new();
8852 create_ns_req.id = Some(vec!["test_ns".to_string()]);
8853 namespace.create_namespace(create_ns_req).await.unwrap();
8854
8855 let schema = create_test_schema();
8857 let ipc_data = create_test_ipc_data(&schema);
8858
8859 let mut create_req = CreateTableRequest::new();
8860 create_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
8861 namespace
8862 .create_table(create_req, bytes::Bytes::from(ipc_data))
8863 .await
8864 .unwrap();
8865
8866 let mut deregister_req = DeregisterTableRequest::new();
8868 deregister_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
8869 let response = namespace.deregister_table(deregister_req).await.unwrap();
8870
8871 assert!(
8873 response.location.is_some(),
8874 "Deregister should return location"
8875 );
8876 let location = response.location.as_ref().unwrap();
8877 let expected_url = lance_io::object_store::uri_to_url(temp_path)
8880 .expect("Failed to convert temp path to URL");
8881 let expected_prefix = expected_url.to_string();
8882 assert!(
8883 location.starts_with(&expected_prefix),
8884 "Location should start with '{}', got: {}",
8885 expected_prefix,
8886 location
8887 );
8888 assert!(
8889 location.contains("test_ns") && location.contains("test_table"),
8890 "Location should contain namespace and table name: {}",
8891 location
8892 );
8893 assert_eq!(
8894 response.id,
8895 Some(vec!["test_ns".to_string(), "test_table".to_string()])
8896 );
8897
8898 let mut exists_req = TableExistsRequest::new();
8900 exists_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
8901 assert!(namespace.table_exists(exists_req).await.is_err());
8902 }
8903
8904 #[tokio::test]
8905 async fn test_register_without_manifest_fails() {
8906 use lance_namespace::models::RegisterTableRequest;
8907
8908 let temp_dir = TempStdDir::default();
8909 let temp_path = temp_dir.to_str().unwrap();
8910
8911 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8913 .manifest_enabled(false)
8914 .build()
8915 .await
8916 .unwrap();
8917
8918 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
8920 register_req.id = Some(vec!["test_table".to_string()]);
8921 let result = namespace.register_table(register_req).await;
8922 assert!(result.is_err());
8923 assert!(
8924 result
8925 .unwrap_err()
8926 .to_string()
8927 .contains("manifest mode is enabled")
8928 );
8929
8930 }
8933
8934 #[tokio::test]
8935 async fn test_register_table_rejects_absolute_uri() {
8936 use lance_namespace::models::RegisterTableRequest;
8937
8938 let temp_dir = TempStdDir::default();
8939 let temp_path = temp_dir.to_str().unwrap();
8940
8941 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8942 .build()
8943 .await
8944 .unwrap();
8945
8946 let mut register_req = RegisterTableRequest::new("s3://bucket/table.lance".to_string());
8948 register_req.id = Some(vec!["test_table".to_string()]);
8949 let result = namespace.register_table(register_req).await;
8950 assert!(result.is_err());
8951 let err_msg = result.unwrap_err().to_string();
8952 assert!(err_msg.contains("Absolute URIs are not allowed"));
8953 }
8954
8955 #[tokio::test]
8956 async fn test_register_table_rejects_absolute_path() {
8957 use lance_namespace::models::RegisterTableRequest;
8958
8959 let temp_dir = TempStdDir::default();
8960 let temp_path = temp_dir.to_str().unwrap();
8961
8962 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8963 .build()
8964 .await
8965 .unwrap();
8966
8967 let mut register_req = RegisterTableRequest::new("/tmp/table.lance".to_string());
8969 register_req.id = Some(vec!["test_table".to_string()]);
8970 let result = namespace.register_table(register_req).await;
8971 assert!(result.is_err());
8972 let err_msg = result.unwrap_err().to_string();
8973 assert!(err_msg.contains("Absolute paths are not allowed"));
8974 }
8975
8976 #[tokio::test]
8977 async fn test_register_table_rejects_path_traversal() {
8978 use lance_namespace::models::RegisterTableRequest;
8979
8980 let temp_dir = TempStdDir::default();
8981 let temp_path = temp_dir.to_str().unwrap();
8982
8983 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8984 .build()
8985 .await
8986 .unwrap();
8987
8988 let mut register_req = RegisterTableRequest::new("../outside/table.lance".to_string());
8990 register_req.id = Some(vec!["test_table".to_string()]);
8991 let result = namespace.register_table(register_req).await;
8992 assert!(result.is_err());
8993 let err_msg = result.unwrap_err().to_string();
8994 assert!(err_msg.contains("Path traversal is not allowed"));
8995 }
8996
8997 #[tokio::test]
8998 async fn test_namespace_write() {
8999 use arrow::array::Int32Array;
9000 use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
9001 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
9002 use lance::dataset::{Dataset, WriteMode, WriteParams};
9003 use lance_namespace::LanceNamespace;
9004
9005 let (namespace, _temp_dir) = create_test_namespace().await;
9006 let namespace = Arc::new(namespace) as Arc<dyn LanceNamespace>;
9007
9008 let table_id = vec!["test_ns".to_string(), "test_table".to_string()];
9010 let schema = Arc::new(ArrowSchema::new(vec![
9011 ArrowField::new("a", DataType::Int32, false),
9012 ArrowField::new("b", DataType::Int32, false),
9013 ]));
9014
9015 let data1 = RecordBatch::try_new(
9017 schema.clone(),
9018 vec![
9019 Arc::new(Int32Array::from(vec![1, 2, 3])),
9020 Arc::new(Int32Array::from(vec![10, 20, 30])),
9021 ],
9022 )
9023 .unwrap();
9024
9025 let reader1 = RecordBatchIterator::new(vec![data1].into_iter().map(Ok), schema.clone());
9026 let dataset =
9027 Dataset::write_into_namespace(reader1, namespace.clone(), table_id.clone(), None)
9028 .await
9029 .unwrap();
9030
9031 assert_eq!(dataset.count_rows(None).await.unwrap(), 3);
9032 assert_eq!(dataset.version().version, 1);
9033
9034 let data2 = RecordBatch::try_new(
9036 schema.clone(),
9037 vec![
9038 Arc::new(Int32Array::from(vec![4, 5])),
9039 Arc::new(Int32Array::from(vec![40, 50])),
9040 ],
9041 )
9042 .unwrap();
9043
9044 let params_append = WriteParams {
9045 mode: WriteMode::Append,
9046 ..Default::default()
9047 };
9048
9049 let reader2 = RecordBatchIterator::new(vec![data2].into_iter().map(Ok), schema.clone());
9050 let dataset = Dataset::write_into_namespace(
9051 reader2,
9052 namespace.clone(),
9053 table_id.clone(),
9054 Some(params_append),
9055 )
9056 .await
9057 .unwrap();
9058
9059 assert_eq!(dataset.count_rows(None).await.unwrap(), 5);
9060 assert_eq!(dataset.version().version, 2);
9061
9062 let data3 = RecordBatch::try_new(
9064 schema.clone(),
9065 vec![
9066 Arc::new(Int32Array::from(vec![100, 200])),
9067 Arc::new(Int32Array::from(vec![1000, 2000])),
9068 ],
9069 )
9070 .unwrap();
9071
9072 let params_overwrite = WriteParams {
9073 mode: WriteMode::Overwrite,
9074 ..Default::default()
9075 };
9076
9077 let reader3 = RecordBatchIterator::new(vec![data3].into_iter().map(Ok), schema.clone());
9078 let dataset = Dataset::write_into_namespace(
9079 reader3,
9080 namespace.clone(),
9081 table_id.clone(),
9082 Some(params_overwrite),
9083 )
9084 .await
9085 .unwrap();
9086
9087 assert_eq!(dataset.count_rows(None).await.unwrap(), 2);
9088 assert_eq!(dataset.version().version, 3);
9089
9090 let result = dataset.scan().try_into_batch().await.unwrap();
9092 let a_col = result
9093 .column_by_name("a")
9094 .unwrap()
9095 .as_any()
9096 .downcast_ref::<Int32Array>()
9097 .unwrap();
9098 assert_eq!(a_col.values(), &[100, 200]);
9099 }
9100
9101 #[tokio::test]
9106 async fn test_declare_table_v1_mode() {
9107 use lance_namespace::models::{
9108 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
9109 };
9110
9111 let temp_dir = TempStdDir::default();
9112 let temp_path = temp_dir.to_str().unwrap();
9113
9114 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9116 .manifest_enabled(false)
9117 .build()
9118 .await
9119 .unwrap();
9120
9121 let mut declare_req = DeclareTableRequest::new();
9123 declare_req.id = Some(vec!["test_table".to_string()]);
9124 let response = namespace.declare_table(declare_req).await.unwrap();
9125
9126 assert!(response.location.is_some());
9128 let location = response.location.as_ref().unwrap();
9129 assert!(location.ends_with("test_table.lance"));
9130
9131 let mut exists_req = TableExistsRequest::new();
9133 exists_req.id = Some(vec!["test_table".to_string()]);
9134 assert!(namespace.table_exists(exists_req).await.is_ok());
9135
9136 let mut describe_req = DescribeTableRequest::new();
9138 describe_req.id = Some(vec!["test_table".to_string()]);
9139 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9140 assert!(describe_response.location.is_some());
9141 assert!(describe_response.version.is_none()); assert!(describe_response.schema.is_none()); assert_eq!(describe_response.is_only_declared, None);
9144
9145 let mut describe_req = DescribeTableRequest::new();
9146 describe_req.id = Some(vec!["test_table".to_string()]);
9147 describe_req.check_declared = Some(true);
9148 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9149 assert_eq!(describe_response.is_only_declared, Some(true));
9150
9151 let mut list_req = ListTablesRequest::new();
9152 list_req.id = Some(vec![]);
9153 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
9154 assert_eq!(list_response.tables, vec!["test_table".to_string()]);
9155
9156 list_req.include_declared = Some(false);
9157 let list_response = namespace.list_tables(list_req).await.unwrap();
9158 assert!(list_response.tables.is_empty());
9159 }
9160
9161 #[tokio::test]
9162 async fn test_insert_into_declared_table_promotes_it_from_declared_state() {
9163 use lance_namespace::models::{
9164 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest,
9165 };
9166
9167 let temp_dir = TempStdDir::default();
9168 let temp_path = temp_dir.to_str().unwrap();
9169
9170 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9171 .manifest_enabled(false)
9172 .build()
9173 .await
9174 .unwrap();
9175
9176 let mut declare_req = DeclareTableRequest::new();
9177 declare_req.id = Some(vec!["test_table".to_string()]);
9178 namespace.declare_table(declare_req).await.unwrap();
9179
9180 let schema = create_test_schema();
9181 let ipc_data = create_test_ipc_data(&schema);
9182 let mut insert_req = InsertIntoTableRequest::new();
9183 insert_req.id = Some(vec!["test_table".to_string()]);
9184 namespace
9185 .insert_into_table(insert_req, bytes::Bytes::from(ipc_data))
9186 .await
9187 .unwrap();
9188
9189 let mut describe_req = DescribeTableRequest::new();
9190 describe_req.id = Some(vec!["test_table".to_string()]);
9191 describe_req.load_detailed_metadata = Some(true);
9192 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9193
9194 assert_eq!(describe_response.is_only_declared, Some(false));
9195 assert_eq!(describe_response.version, Some(1));
9196 assert!(describe_response.schema.is_some());
9197
9198 let mut list_req = ListTablesRequest::new();
9199 list_req.id = Some(vec![]);
9200 list_req.include_declared = Some(false);
9201 assert_eq!(
9202 namespace.list_tables(list_req).await.unwrap().tables,
9203 vec!["test_table".to_string()]
9204 );
9205 }
9206
9207 #[tokio::test]
9208 async fn test_create_table_after_declare_table_v1_mode_creates_table() {
9209 use lance_namespace::models::{
9210 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
9211 };
9212
9213 let temp_dir = TempStdDir::default();
9214 let temp_path = temp_dir.to_str().unwrap();
9215
9216 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9217 .manifest_enabled(false)
9218 .build()
9219 .await
9220 .unwrap();
9221
9222 let mut declare_req = DeclareTableRequest::new();
9223 declare_req.id = Some(vec!["test_table".to_string()]);
9224 namespace.declare_table(declare_req).await.unwrap();
9225
9226 let mut create_req = CreateTableRequest::new();
9227 create_req.id = Some(vec!["test_table".to_string()]);
9228 let response = namespace
9229 .create_table(
9230 create_req,
9231 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9232 )
9233 .await
9234 .unwrap();
9235
9236 assert_eq!(response.version, Some(1));
9237
9238 let mut describe_req = DescribeTableRequest::new();
9239 describe_req.id = Some(vec!["test_table".to_string()]);
9240 describe_req.load_detailed_metadata = Some(true);
9241 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9242 assert_eq!(describe_response.is_only_declared, Some(false));
9243 assert_eq!(describe_response.version, Some(1));
9244
9245 let mut list_req = ListTablesRequest::new();
9246 list_req.id = Some(vec![]);
9247 list_req.include_declared = Some(false);
9248 assert_eq!(
9249 namespace.list_tables(list_req).await.unwrap().tables,
9250 vec!["test_table".to_string()]
9251 );
9252 }
9253
9254 #[tokio::test]
9255 async fn test_insert_into_declared_table_with_manifest_promotes_it() {
9256 use lance_namespace::models::{
9257 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest, ListTablesRequest,
9258 };
9259
9260 let temp_dir = TempStdDir::default();
9261 let temp_path = temp_dir.to_str().unwrap();
9262
9263 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9264 .manifest_enabled(true)
9265 .dir_listing_enabled(false)
9266 .build()
9267 .await
9268 .unwrap();
9269
9270 let mut declare_req = DeclareTableRequest::new();
9271 declare_req.id = Some(vec!["test_table".to_string()]);
9272 namespace.declare_table(declare_req).await.unwrap();
9273
9274 let mut insert_req = InsertIntoTableRequest::new();
9275 insert_req.id = Some(vec!["test_table".to_string()]);
9276 namespace
9277 .insert_into_table(
9278 insert_req,
9279 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9280 )
9281 .await
9282 .unwrap();
9283
9284 let mut describe_req = DescribeTableRequest::new();
9285 describe_req.id = Some(vec!["test_table".to_string()]);
9286 describe_req.load_detailed_metadata = Some(true);
9287 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9288 assert_eq!(describe_response.is_only_declared, Some(false));
9289 assert_eq!(describe_response.version, Some(1));
9290
9291 let mut list_req = ListTablesRequest::new();
9292 list_req.id = Some(vec![]);
9293 list_req.include_declared = Some(false);
9294 assert_eq!(
9295 namespace.list_tables(list_req).await.unwrap().tables,
9296 vec!["test_table".to_string()]
9297 );
9298 }
9299
9300 #[tokio::test]
9301 async fn test_create_table_after_declare_table_with_manifest_creates_table() {
9302 use lance_namespace::models::{
9303 CreateTableRequest, DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
9304 };
9305
9306 let temp_dir = TempStdDir::default();
9307 let temp_path = temp_dir.to_str().unwrap();
9308
9309 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9310 .manifest_enabled(true)
9311 .dir_listing_enabled(false)
9312 .build()
9313 .await
9314 .unwrap();
9315
9316 let mut declare_req = DeclareTableRequest::new();
9317 declare_req.id = Some(vec!["test_table".to_string()]);
9318 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
9319 namespace.declare_table(declare_req).await.unwrap();
9320
9321 let mut create_req = CreateTableRequest::new();
9322 create_req.id = Some(vec!["test_table".to_string()]);
9323 create_req.mode = Some("Overwrite".to_string());
9324 let response = namespace
9325 .create_table(
9326 create_req,
9327 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9328 )
9329 .await
9330 .unwrap();
9331
9332 assert_eq!(response.version, Some(1));
9333 assert_eq!(
9334 response
9335 .properties
9336 .as_ref()
9337 .and_then(|properties| properties.get("owner")),
9338 Some(&"alice".to_string())
9339 );
9340
9341 let mut describe_req = DescribeTableRequest::new();
9342 describe_req.id = Some(vec!["test_table".to_string()]);
9343 describe_req.load_detailed_metadata = Some(true);
9344 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9345 assert_eq!(describe_response.is_only_declared, Some(false));
9346 assert_eq!(describe_response.version, Some(1));
9347 assert_eq!(
9348 describe_response
9349 .properties
9350 .as_ref()
9351 .and_then(|properties| properties.get("owner")),
9352 Some(&"alice".to_string())
9353 );
9354
9355 let mut list_req = ListTablesRequest::new();
9356 list_req.id = Some(vec![]);
9357 list_req.include_declared = Some(false);
9358 assert_eq!(
9359 namespace.list_tables(list_req).await.unwrap().tables,
9360 vec!["test_table".to_string()]
9361 );
9362 }
9363
9364 #[tokio::test]
9365 async fn test_create_table_after_declare_table_with_manifest_rejects_new_properties() {
9366 use lance_namespace::models::{CreateTableRequest, DeclareTableRequest};
9367
9368 let temp_dir = TempStdDir::default();
9369 let temp_path = temp_dir.to_str().unwrap();
9370
9371 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9372 .manifest_enabled(true)
9373 .dir_listing_enabled(false)
9374 .build()
9375 .await
9376 .unwrap();
9377
9378 let mut declare_req = DeclareTableRequest::new();
9379 declare_req.id = Some(vec!["test_table".to_string()]);
9380 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
9381 namespace.declare_table(declare_req).await.unwrap();
9382
9383 let mut create_req = CreateTableRequest::new();
9384 create_req.id = Some(vec!["test_table".to_string()]);
9385 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
9386
9387 let result = namespace
9388 .create_table(
9389 create_req,
9390 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9391 )
9392 .await;
9393
9394 assert!(result.is_err());
9395 assert!(
9396 result
9397 .unwrap_err()
9398 .to_string()
9399 .contains("cannot set properties for already declared table")
9400 );
9401 }
9402
9403 #[tokio::test]
9404 async fn test_create_table_with_manifest_exist_ok_keeps_existing_table() {
9405 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
9406
9407 let temp_dir = TempStdDir::default();
9408 let temp_path = temp_dir.to_str().unwrap();
9409
9410 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9411 .manifest_enabled(true)
9412 .dir_listing_enabled(false)
9413 .build()
9414 .await
9415 .unwrap();
9416
9417 let mut create_req = CreateTableRequest::new();
9418 create_req.id = Some(vec!["test_table".to_string()]);
9419 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
9420 namespace
9421 .create_table(
9422 create_req,
9423 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9424 )
9425 .await
9426 .unwrap();
9427
9428 let mut create_req = CreateTableRequest::new();
9429 create_req.id = Some(vec!["test_table".to_string()]);
9430 create_req.mode = Some("ExistOk".to_string());
9431 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
9432 let response = namespace
9433 .create_table(
9434 create_req,
9435 bytes::Bytes::from(create_single_row_test_ipc_data()),
9436 )
9437 .await
9438 .unwrap();
9439
9440 assert_eq!(
9441 response
9442 .properties
9443 .as_ref()
9444 .and_then(|properties| properties.get("owner")),
9445 Some(&"alice".to_string())
9446 );
9447 assert_eq!(
9448 open_dataset(&namespace, "test_table")
9449 .await
9450 .count_rows(None)
9451 .await
9452 .unwrap(),
9453 2
9454 );
9455
9456 let mut describe_req = DescribeTableRequest::new();
9457 describe_req.id = Some(vec!["test_table".to_string()]);
9458 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9459 assert_eq!(
9460 describe_response
9461 .properties
9462 .as_ref()
9463 .and_then(|properties| properties.get("owner")),
9464 Some(&"alice".to_string())
9465 );
9466 }
9467
9468 #[tokio::test]
9469 async fn test_create_table_with_manifest_overwrite_replaces_existing_table() {
9470 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
9471
9472 let temp_dir = TempStdDir::default();
9473 let temp_path = temp_dir.to_str().unwrap();
9474
9475 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9476 .manifest_enabled(true)
9477 .dir_listing_enabled(false)
9478 .build()
9479 .await
9480 .unwrap();
9481
9482 let mut create_req = CreateTableRequest::new();
9483 create_req.id = Some(vec!["test_table".to_string()]);
9484 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
9485 namespace
9486 .create_table(
9487 create_req,
9488 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9489 )
9490 .await
9491 .unwrap();
9492
9493 let mut create_req = CreateTableRequest::new();
9494 create_req.id = Some(vec!["test_table".to_string()]);
9495 create_req.mode = Some("overwrite".to_string());
9496 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
9497 let response = namespace
9498 .create_table(
9499 create_req,
9500 bytes::Bytes::from(create_single_row_test_ipc_data()),
9501 )
9502 .await
9503 .unwrap();
9504
9505 assert_eq!(response.version, Some(2));
9506 assert_eq!(
9507 response
9508 .properties
9509 .as_ref()
9510 .and_then(|properties| properties.get("owner")),
9511 Some(&"bob".to_string())
9512 );
9513 assert_eq!(
9514 open_dataset(&namespace, "test_table")
9515 .await
9516 .count_rows(None)
9517 .await
9518 .unwrap(),
9519 1
9520 );
9521
9522 let mut describe_req = DescribeTableRequest::new();
9523 describe_req.id = Some(vec!["test_table".to_string()]);
9524 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9525 assert_eq!(
9526 describe_response
9527 .properties
9528 .as_ref()
9529 .and_then(|properties| properties.get("owner")),
9530 Some(&"bob".to_string())
9531 );
9532 }
9533
9534 #[tokio::test]
9535 async fn test_create_table_with_manifest_invalid_mode_rejected() {
9536 use lance_namespace::models::CreateTableRequest;
9537
9538 let temp_dir = TempStdDir::default();
9539 let temp_path = temp_dir.to_str().unwrap();
9540
9541 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9542 .manifest_enabled(true)
9543 .dir_listing_enabled(false)
9544 .build()
9545 .await
9546 .unwrap();
9547
9548 let mut create_req = CreateTableRequest::new();
9549 create_req.id = Some(vec!["test_table".to_string()]);
9550 create_req.mode = Some("append".to_string());
9551 let result = namespace
9552 .create_table(
9553 create_req,
9554 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9555 )
9556 .await;
9557
9558 assert!(result.is_err());
9559 assert!(
9560 result
9561 .unwrap_err()
9562 .to_string()
9563 .contains("Unsupported create_table mode")
9564 );
9565 }
9566
9567 #[tokio::test]
9568 async fn test_merge_insert_into_declared_table_v1_mode_creates_table() {
9569 use lance_namespace::models::{
9570 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
9571 MergeInsertIntoTableRequest,
9572 };
9573
9574 let temp_dir = TempStdDir::default();
9575 let temp_path = temp_dir.to_str().unwrap();
9576
9577 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9578 .manifest_enabled(false)
9579 .build()
9580 .await
9581 .unwrap();
9582
9583 let mut declare_req = DeclareTableRequest::new();
9584 declare_req.id = Some(vec!["test_table".to_string()]);
9585 namespace.declare_table(declare_req).await.unwrap();
9586
9587 let mut merge_req = MergeInsertIntoTableRequest::new();
9588 merge_req.id = Some(vec!["test_table".to_string()]);
9589 merge_req.on = Some("id".to_string());
9590 let response = namespace
9591 .merge_insert_into_table(
9592 merge_req,
9593 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9594 )
9595 .await
9596 .unwrap();
9597
9598 assert_eq!(response.num_inserted_rows, Some(2));
9599 assert_eq!(response.num_updated_rows, Some(0));
9600
9601 let mut describe_req = DescribeTableRequest::new();
9602 describe_req.id = Some(vec!["test_table".to_string()]);
9603 describe_req.load_detailed_metadata = Some(true);
9604 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9605 assert_eq!(describe_response.is_only_declared, Some(false));
9606 assert_eq!(describe_response.version, Some(1));
9607
9608 let mut list_req = ListTablesRequest::new();
9609 list_req.id = Some(vec![]);
9610 list_req.include_declared = Some(false);
9611 assert_eq!(
9612 namespace.list_tables(list_req).await.unwrap().tables,
9613 vec!["test_table".to_string()]
9614 );
9615 }
9616
9617 #[tokio::test]
9618 async fn test_merge_insert_into_declared_table_with_manifest_creates_table() {
9619 use lance_namespace::models::{
9620 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
9621 MergeInsertIntoTableRequest,
9622 };
9623
9624 let temp_dir = TempStdDir::default();
9625 let temp_path = temp_dir.to_str().unwrap();
9626
9627 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9628 .manifest_enabled(true)
9629 .dir_listing_enabled(false)
9630 .build()
9631 .await
9632 .unwrap();
9633
9634 let mut declare_req = DeclareTableRequest::new();
9635 declare_req.id = Some(vec!["test_table".to_string()]);
9636 namespace.declare_table(declare_req).await.unwrap();
9637
9638 let mut merge_req = MergeInsertIntoTableRequest::new();
9639 merge_req.id = Some(vec!["test_table".to_string()]);
9640 merge_req.on = Some("id".to_string());
9641 let response = namespace
9642 .merge_insert_into_table(
9643 merge_req,
9644 bytes::Bytes::from(create_non_empty_test_ipc_data()),
9645 )
9646 .await
9647 .unwrap();
9648
9649 assert_eq!(response.num_inserted_rows, Some(2));
9650 assert_eq!(response.num_updated_rows, Some(0));
9651
9652 let mut describe_req = DescribeTableRequest::new();
9653 describe_req.id = Some(vec!["test_table".to_string()]);
9654 describe_req.load_detailed_metadata = Some(true);
9655 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9656 assert_eq!(describe_response.is_only_declared, Some(false));
9657 assert_eq!(describe_response.version, Some(1));
9658
9659 let mut list_req = ListTablesRequest::new();
9660 list_req.id = Some(vec![]);
9661 list_req.include_declared = Some(false);
9662 assert_eq!(
9663 namespace.list_tables(list_req).await.unwrap().tables,
9664 vec!["test_table".to_string()]
9665 );
9666 }
9667
9668 #[tokio::test]
9669 async fn test_declare_table_with_manifest() {
9670 use lance_namespace::models::{
9671 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
9672 };
9673
9674 let temp_dir = TempStdDir::default();
9675 let temp_path = temp_dir.to_str().unwrap();
9676
9677 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9679 .manifest_enabled(true)
9680 .dir_listing_enabled(false)
9681 .build()
9682 .await
9683 .unwrap();
9684
9685 let mut declare_req = DeclareTableRequest::new();
9687 declare_req.id = Some(vec!["test_table".to_string()]);
9688 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
9689 let response = namespace.declare_table(declare_req).await.unwrap();
9690
9691 assert!(response.location.is_some());
9693 assert_eq!(
9694 response
9695 .properties
9696 .as_ref()
9697 .and_then(|properties| properties.get("owner")),
9698 Some(&"alice".to_string())
9699 );
9700
9701 let mut exists_req = TableExistsRequest::new();
9703 exists_req.id = Some(vec!["test_table".to_string()]);
9704 assert!(namespace.table_exists(exists_req).await.is_ok());
9705
9706 let mut describe_req = DescribeTableRequest::new();
9707 describe_req.id = Some(vec!["test_table".to_string()]);
9708 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9709 assert_eq!(describe_response.is_only_declared, None);
9710
9711 let mut describe_req = DescribeTableRequest::new();
9712 describe_req.id = Some(vec!["test_table".to_string()]);
9713 describe_req.check_declared = Some(true);
9714 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9715 assert_eq!(describe_response.is_only_declared, Some(true));
9716 assert_eq!(
9717 describe_response
9718 .properties
9719 .as_ref()
9720 .and_then(|properties| properties.get("owner")),
9721 Some(&"alice".to_string())
9722 );
9723
9724 let mut list_req = ListTablesRequest::new();
9725 list_req.id = Some(vec![]);
9726 assert_eq!(
9727 namespace
9728 .list_tables(list_req.clone())
9729 .await
9730 .unwrap()
9731 .tables,
9732 vec!["test_table".to_string()]
9733 );
9734 list_req.include_declared = Some(false);
9735 assert!(
9736 namespace
9737 .list_tables(list_req)
9738 .await
9739 .unwrap()
9740 .tables
9741 .is_empty()
9742 );
9743 }
9744
9745 #[tokio::test]
9746 async fn test_declare_table_when_table_exists() {
9747 use lance_namespace::models::DeclareTableRequest;
9748
9749 let temp_dir = TempStdDir::default();
9750 let temp_path = temp_dir.to_str().unwrap();
9751
9752 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9753 .manifest_enabled(false)
9754 .build()
9755 .await
9756 .unwrap();
9757
9758 let schema = create_test_schema();
9760 let ipc_data = create_test_ipc_data(&schema);
9761 let mut create_req = CreateTableRequest::new();
9762 create_req.id = Some(vec!["test_table".to_string()]);
9763 namespace
9764 .create_table(create_req, bytes::Bytes::from(ipc_data))
9765 .await
9766 .unwrap();
9767
9768 let mut declare_req = DeclareTableRequest::new();
9770 declare_req.id = Some(vec!["test_table".to_string()]);
9771 let result = namespace.declare_table(declare_req).await;
9772 assert!(result.is_err());
9773 }
9774
9775 #[tokio::test]
9780 async fn test_deregister_table_v1_mode() {
9781 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
9782
9783 let temp_dir = TempStdDir::default();
9784 let temp_path = temp_dir.to_str().unwrap();
9785
9786 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9788 .manifest_enabled(false)
9789 .dir_listing_enabled(true)
9790 .build()
9791 .await
9792 .unwrap();
9793
9794 let schema = create_test_schema();
9796 let ipc_data = create_test_ipc_data(&schema);
9797 let mut create_req = CreateTableRequest::new();
9798 create_req.id = Some(vec!["test_table".to_string()]);
9799 namespace
9800 .create_table(create_req, bytes::Bytes::from(ipc_data))
9801 .await
9802 .unwrap();
9803
9804 let mut exists_req = TableExistsRequest::new();
9806 exists_req.id = Some(vec!["test_table".to_string()]);
9807 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
9808
9809 let mut deregister_req = DeregisterTableRequest::new();
9811 deregister_req.id = Some(vec!["test_table".to_string()]);
9812 let response = namespace.deregister_table(deregister_req).await.unwrap();
9813
9814 assert!(response.location.is_some());
9816 let location = response.location.as_ref().unwrap();
9817 assert!(location.contains("test_table"));
9818
9819 let result = namespace.table_exists(exists_req).await;
9821 assert!(result.is_err());
9822 assert!(result.unwrap_err().to_string().contains("deregistered"));
9823
9824 let dataset = Dataset::open(location).await;
9826 assert!(dataset.is_ok(), "Physical table data should still exist");
9827 }
9828
9829 #[tokio::test]
9830 async fn test_deregister_table_v1_already_deregistered() {
9831 use lance_namespace::models::DeregisterTableRequest;
9832
9833 let temp_dir = TempStdDir::default();
9834 let temp_path = temp_dir.to_str().unwrap();
9835
9836 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9837 .manifest_enabled(false)
9838 .dir_listing_enabled(true)
9839 .build()
9840 .await
9841 .unwrap();
9842
9843 let schema = create_test_schema();
9845 let ipc_data = create_test_ipc_data(&schema);
9846 let mut create_req = CreateTableRequest::new();
9847 create_req.id = Some(vec!["test_table".to_string()]);
9848 namespace
9849 .create_table(create_req, bytes::Bytes::from(ipc_data))
9850 .await
9851 .unwrap();
9852
9853 let mut deregister_req = DeregisterTableRequest::new();
9855 deregister_req.id = Some(vec!["test_table".to_string()]);
9856 namespace
9857 .deregister_table(deregister_req.clone())
9858 .await
9859 .unwrap();
9860
9861 let result = namespace.deregister_table(deregister_req).await;
9863 assert!(result.is_err());
9864 assert!(
9865 result
9866 .unwrap_err()
9867 .to_string()
9868 .contains("already deregistered")
9869 );
9870 }
9871
9872 #[tokio::test]
9877 async fn test_list_tables_skips_deregistered_v1() {
9878 use lance_namespace::models::DeregisterTableRequest;
9879
9880 let temp_dir = TempStdDir::default();
9881 let temp_path = temp_dir.to_str().unwrap();
9882
9883 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9884 .manifest_enabled(false)
9885 .dir_listing_enabled(true)
9886 .build()
9887 .await
9888 .unwrap();
9889
9890 let schema = create_test_schema();
9892 let ipc_data = create_test_ipc_data(&schema);
9893
9894 let mut create_req1 = CreateTableRequest::new();
9895 create_req1.id = Some(vec!["table1".to_string()]);
9896 namespace
9897 .create_table(create_req1, bytes::Bytes::from(ipc_data.clone()))
9898 .await
9899 .unwrap();
9900
9901 let mut create_req2 = CreateTableRequest::new();
9902 create_req2.id = Some(vec!["table2".to_string()]);
9903 namespace
9904 .create_table(create_req2, bytes::Bytes::from(ipc_data))
9905 .await
9906 .unwrap();
9907
9908 let mut list_req = ListTablesRequest::new();
9910 list_req.id = Some(vec![]);
9911 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
9912 assert_eq!(list_response.tables.len(), 2);
9913
9914 let mut deregister_req = DeregisterTableRequest::new();
9916 deregister_req.id = Some(vec!["table1".to_string()]);
9917 namespace.deregister_table(deregister_req).await.unwrap();
9918
9919 let list_response = namespace.list_tables(list_req).await.unwrap();
9921 assert_eq!(list_response.tables.len(), 1);
9922 assert!(list_response.tables.contains(&"table2".to_string()));
9923 assert!(!list_response.tables.contains(&"table1".to_string()));
9924 }
9925
9926 #[tokio::test]
9931 async fn test_describe_table_fails_for_deregistered_v1() {
9932 use lance_namespace::models::{DeregisterTableRequest, DescribeTableRequest};
9933
9934 let temp_dir = TempStdDir::default();
9935 let temp_path = temp_dir.to_str().unwrap();
9936
9937 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9938 .manifest_enabled(false)
9939 .dir_listing_enabled(true)
9940 .build()
9941 .await
9942 .unwrap();
9943
9944 let schema = create_test_schema();
9946 let ipc_data = create_test_ipc_data(&schema);
9947 let mut create_req = CreateTableRequest::new();
9948 create_req.id = Some(vec!["test_table".to_string()]);
9949 namespace
9950 .create_table(create_req, bytes::Bytes::from(ipc_data))
9951 .await
9952 .unwrap();
9953
9954 let mut describe_req = DescribeTableRequest::new();
9956 describe_req.id = Some(vec!["test_table".to_string()]);
9957 assert!(namespace.describe_table(describe_req.clone()).await.is_ok());
9958
9959 let mut deregister_req = DeregisterTableRequest::new();
9961 deregister_req.id = Some(vec!["test_table".to_string()]);
9962 namespace.deregister_table(deregister_req).await.unwrap();
9963
9964 let result = namespace.describe_table(describe_req).await;
9966 assert!(result.is_err());
9967 let err = result.unwrap_err();
9968 assert!(matches!(err, Error::Namespace { .. }));
9969 let err_msg = err.to_string();
9970 assert!(err_msg.contains("deregistered"));
9971 assert!(err_msg.contains("table id 'test_table'"));
9972 }
9973
9974 #[tokio::test]
9975 async fn test_table_exists_fails_for_deregistered_v1() {
9976 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
9977
9978 let temp_dir = TempStdDir::default();
9979 let temp_path = temp_dir.to_str().unwrap();
9980
9981 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9982 .manifest_enabled(false)
9983 .dir_listing_enabled(true)
9984 .build()
9985 .await
9986 .unwrap();
9987
9988 let schema = create_test_schema();
9990 let ipc_data = create_test_ipc_data(&schema);
9991 let mut create_req = CreateTableRequest::new();
9992 create_req.id = Some(vec!["test_table".to_string()]);
9993 namespace
9994 .create_table(create_req, bytes::Bytes::from(ipc_data))
9995 .await
9996 .unwrap();
9997
9998 let mut exists_req = TableExistsRequest::new();
10000 exists_req.id = Some(vec!["test_table".to_string()]);
10001 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
10002
10003 let mut deregister_req = DeregisterTableRequest::new();
10005 deregister_req.id = Some(vec!["test_table".to_string()]);
10006 namespace.deregister_table(deregister_req).await.unwrap();
10007
10008 let result = namespace.table_exists(exists_req).await;
10010 assert!(result.is_err());
10011 let err = result.unwrap_err();
10012 assert!(matches!(err, Error::Namespace { .. }));
10013 let err_msg = err.to_string();
10014 assert!(err_msg.contains("deregistered"));
10015 assert!(err_msg.contains("table id 'test_table'"));
10016 }
10017
10018 #[tokio::test]
10019 async fn test_atomic_table_status_check() {
10020 let temp_dir = TempStdDir::default();
10024 let temp_path = temp_dir.to_str().unwrap();
10025
10026 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10027 .manifest_enabled(false)
10028 .dir_listing_enabled(true)
10029 .build()
10030 .await
10031 .unwrap();
10032
10033 let schema = create_test_schema();
10035 let ipc_data = create_test_ipc_data(&schema);
10036 let mut create_req = CreateTableRequest::new();
10037 create_req.id = Some(vec!["test_table".to_string()]);
10038 namespace
10039 .create_table(create_req, bytes::Bytes::from(ipc_data))
10040 .await
10041 .unwrap();
10042
10043 let status = namespace.check_table_status("test_table").await;
10045 assert!(status.exists);
10046 assert!(!status.is_deregistered);
10047 assert!(!status.has_reserved_file);
10048 }
10049
10050 #[tokio::test]
10051 async fn test_table_version_tracking_enabled_managed_versioning() {
10052 use lance_namespace::models::DescribeTableRequest;
10053
10054 let temp_dir = TempStdDir::default();
10055 let temp_path = temp_dir.to_str().unwrap();
10056
10057 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10059 .table_version_tracking_enabled(true)
10060 .build()
10061 .await
10062 .unwrap();
10063
10064 let schema = create_test_schema();
10066 let ipc_data = create_test_ipc_data(&schema);
10067 let mut create_req = CreateTableRequest::new();
10068 create_req.id = Some(vec!["test_table".to_string()]);
10069 namespace
10070 .create_table(create_req, bytes::Bytes::from(ipc_data))
10071 .await
10072 .unwrap();
10073
10074 let mut describe_req = DescribeTableRequest::new();
10076 describe_req.id = Some(vec!["test_table".to_string()]);
10077 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
10078
10079 assert_eq!(
10081 describe_resp.managed_versioning,
10082 Some(true),
10083 "managed_versioning should be true when table_version_tracking_enabled=true"
10084 );
10085 }
10086
10087 #[tokio::test]
10088 async fn test_table_version_tracking_disabled_no_managed_versioning() {
10089 use lance_namespace::models::DescribeTableRequest;
10090
10091 let temp_dir = TempStdDir::default();
10092 let temp_path = temp_dir.to_str().unwrap();
10093
10094 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10096 .table_version_tracking_enabled(false)
10097 .build()
10098 .await
10099 .unwrap();
10100
10101 let schema = create_test_schema();
10103 let ipc_data = create_test_ipc_data(&schema);
10104 let mut create_req = CreateTableRequest::new();
10105 create_req.id = Some(vec!["test_table".to_string()]);
10106 namespace
10107 .create_table(create_req, bytes::Bytes::from(ipc_data))
10108 .await
10109 .unwrap();
10110
10111 let mut describe_req = DescribeTableRequest::new();
10113 describe_req.id = Some(vec!["test_table".to_string()]);
10114 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
10115
10116 assert!(
10118 describe_resp.managed_versioning.is_none(),
10119 "managed_versioning should be None when table_version_tracking_enabled=false, got: {:?}",
10120 describe_resp.managed_versioning
10121 );
10122 }
10123
10124 #[tokio::test]
10125 async fn test_list_table_versions() {
10126 use arrow::array::{Int32Array, RecordBatchIterator};
10127 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
10128 use arrow::record_batch::RecordBatch;
10129 use lance::dataset::{Dataset, WriteMode, WriteParams};
10130 use lance_namespace::models::{CreateNamespaceRequest, ListTableVersionsRequest};
10131
10132 let temp_dir = TempStrDir::default();
10133 let temp_path: &str = &temp_dir;
10134
10135 let namespace: Arc<dyn LanceNamespace> = Arc::new(
10136 DirectoryNamespaceBuilder::new(temp_path)
10137 .table_version_tracking_enabled(true)
10138 .build()
10139 .await
10140 .unwrap(),
10141 );
10142
10143 let mut create_ns_req = CreateNamespaceRequest::new();
10145 create_ns_req.id = Some(vec!["workspace".to_string()]);
10146 namespace.create_namespace(create_ns_req).await.unwrap();
10147
10148 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
10150 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
10151 "id",
10152 DataType::Int32,
10153 false,
10154 )]));
10155 let batch = RecordBatch::try_new(
10156 arrow_schema.clone(),
10157 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
10158 )
10159 .unwrap();
10160 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
10161 let write_params = WriteParams {
10162 mode: WriteMode::Create,
10163 ..Default::default()
10164 };
10165 let mut dataset = Dataset::write_into_namespace(
10166 batches,
10167 namespace.clone(),
10168 table_id.clone(),
10169 Some(write_params),
10170 )
10171 .await
10172 .unwrap();
10173
10174 let batch2 = RecordBatch::try_new(
10176 arrow_schema.clone(),
10177 vec![Arc::new(Int32Array::from(vec![100, 200]))],
10178 )
10179 .unwrap();
10180 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
10181 dataset.append(batches, None).await.unwrap();
10182
10183 let batch3 = RecordBatch::try_new(
10185 arrow_schema.clone(),
10186 vec![Arc::new(Int32Array::from(vec![300, 400]))],
10187 )
10188 .unwrap();
10189 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
10190 dataset.append(batches, None).await.unwrap();
10191
10192 let mut list_req = ListTableVersionsRequest::new();
10194 list_req.id = Some(table_id.clone());
10195 let list_resp = namespace.list_table_versions(list_req).await.unwrap();
10196
10197 assert_eq!(
10198 list_resp.versions.len(),
10199 3,
10200 "Should have 3 versions, got: {:?}",
10201 list_resp.versions
10202 );
10203
10204 for expected_version in 1..=3 {
10206 let version = list_resp
10207 .versions
10208 .iter()
10209 .find(|v| v.version == expected_version)
10210 .unwrap_or_else(|| panic!("Expected version {}", expected_version));
10211
10212 assert!(
10213 !version.manifest_path.is_empty(),
10214 "manifest_path should be set for version {}",
10215 expected_version
10216 );
10217 assert!(
10218 version.manifest_path.contains(".manifest"),
10219 "manifest_path should contain .manifest for version {}",
10220 expected_version
10221 );
10222 assert!(
10223 version.manifest_size.is_some(),
10224 "manifest_size should be set for version {}",
10225 expected_version
10226 );
10227 assert!(
10228 version.manifest_size.unwrap() > 0,
10229 "manifest_size should be > 0 for version {}",
10230 expected_version
10231 );
10232 assert!(
10233 version.timestamp_millis.is_some(),
10234 "timestamp_millis should be set for version {}",
10235 expected_version
10236 );
10237 }
10238 }
10239
10240 #[tokio::test]
10241 async fn test_describe_table_version() {
10242 use arrow::array::{Int32Array, RecordBatchIterator};
10243 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
10244 use arrow::record_batch::RecordBatch;
10245 use lance::dataset::{Dataset, WriteMode, WriteParams};
10246 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
10247
10248 let temp_dir = TempStrDir::default();
10249 let temp_path: &str = &temp_dir;
10250
10251 let namespace: Arc<dyn LanceNamespace> = Arc::new(
10252 DirectoryNamespaceBuilder::new(temp_path)
10253 .table_version_tracking_enabled(true)
10254 .build()
10255 .await
10256 .unwrap(),
10257 );
10258
10259 let mut create_ns_req = CreateNamespaceRequest::new();
10261 create_ns_req.id = Some(vec!["workspace".to_string()]);
10262 namespace.create_namespace(create_ns_req).await.unwrap();
10263
10264 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
10266 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
10267 "id",
10268 DataType::Int32,
10269 false,
10270 )]));
10271 let batch = RecordBatch::try_new(
10272 arrow_schema.clone(),
10273 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
10274 )
10275 .unwrap();
10276 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
10277 let write_params = WriteParams {
10278 mode: WriteMode::Create,
10279 ..Default::default()
10280 };
10281 let mut dataset = Dataset::write_into_namespace(
10282 batches,
10283 namespace.clone(),
10284 table_id.clone(),
10285 Some(write_params),
10286 )
10287 .await
10288 .unwrap();
10289
10290 let batch2 = RecordBatch::try_new(
10292 arrow_schema.clone(),
10293 vec![Arc::new(Int32Array::from(vec![100, 200]))],
10294 )
10295 .unwrap();
10296 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
10297 dataset.append(batches, None).await.unwrap();
10298
10299 let mut describe_req = DescribeTableVersionRequest::new();
10301 describe_req.id = Some(table_id.clone());
10302 describe_req.version = Some(1);
10303 let describe_resp = namespace
10304 .describe_table_version(describe_req)
10305 .await
10306 .unwrap();
10307
10308 let version = &describe_resp.version;
10309 assert_eq!(version.version, 1);
10310 assert!(version.timestamp_millis.is_some());
10311 assert!(
10312 !version.manifest_path.is_empty(),
10313 "manifest_path should be set"
10314 );
10315 assert!(
10316 version.manifest_path.contains(".manifest"),
10317 "manifest_path should contain .manifest"
10318 );
10319 assert!(
10320 version.manifest_size.is_some(),
10321 "manifest_size should be set"
10322 );
10323 assert!(
10324 version.manifest_size.unwrap() > 0,
10325 "manifest_size should be > 0"
10326 );
10327
10328 let mut describe_req = DescribeTableVersionRequest::new();
10330 describe_req.id = Some(table_id.clone());
10331 describe_req.version = Some(2);
10332 let describe_resp = namespace
10333 .describe_table_version(describe_req)
10334 .await
10335 .unwrap();
10336
10337 let version = &describe_resp.version;
10338 assert_eq!(version.version, 2);
10339 assert!(version.timestamp_millis.is_some());
10340 assert!(
10341 !version.manifest_path.is_empty(),
10342 "manifest_path should be set"
10343 );
10344 assert!(
10345 version.manifest_size.is_some(),
10346 "manifest_size should be set"
10347 );
10348 assert!(
10349 version.manifest_size.unwrap() > 0,
10350 "manifest_size should be > 0"
10351 );
10352 }
10353
10354 #[tokio::test]
10355 async fn test_describe_table_version_latest() {
10356 use arrow::array::{Int32Array, RecordBatchIterator};
10357 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
10358 use arrow::record_batch::RecordBatch;
10359 use lance::dataset::{Dataset, WriteMode, WriteParams};
10360 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
10361
10362 let temp_dir = TempStrDir::default();
10363 let temp_path: &str = &temp_dir;
10364
10365 let namespace: Arc<dyn LanceNamespace> = Arc::new(
10366 DirectoryNamespaceBuilder::new(temp_path)
10367 .table_version_tracking_enabled(true)
10368 .build()
10369 .await
10370 .unwrap(),
10371 );
10372
10373 let mut create_ns_req = CreateNamespaceRequest::new();
10375 create_ns_req.id = Some(vec!["workspace".to_string()]);
10376 namespace.create_namespace(create_ns_req).await.unwrap();
10377
10378 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
10380 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
10381 "id",
10382 DataType::Int32,
10383 false,
10384 )]));
10385 let batch = RecordBatch::try_new(
10386 arrow_schema.clone(),
10387 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
10388 )
10389 .unwrap();
10390 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
10391 let write_params = WriteParams {
10392 mode: WriteMode::Create,
10393 ..Default::default()
10394 };
10395 let mut dataset = Dataset::write_into_namespace(
10396 batches,
10397 namespace.clone(),
10398 table_id.clone(),
10399 Some(write_params),
10400 )
10401 .await
10402 .unwrap();
10403
10404 let batch2 = RecordBatch::try_new(
10406 arrow_schema.clone(),
10407 vec![Arc::new(Int32Array::from(vec![100, 200]))],
10408 )
10409 .unwrap();
10410 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
10411 dataset.append(batches, None).await.unwrap();
10412
10413 let batch3 = RecordBatch::try_new(
10415 arrow_schema.clone(),
10416 vec![Arc::new(Int32Array::from(vec![300, 400]))],
10417 )
10418 .unwrap();
10419 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
10420 dataset.append(batches, None).await.unwrap();
10421
10422 let mut describe_req = DescribeTableVersionRequest::new();
10424 describe_req.id = Some(table_id.clone());
10425 describe_req.version = None;
10426 let describe_resp = namespace
10427 .describe_table_version(describe_req)
10428 .await
10429 .unwrap();
10430
10431 assert_eq!(describe_resp.version.version, 3);
10433 }
10434
10435 #[tokio::test]
10436 async fn test_create_table_version() {
10437 use futures::TryStreamExt;
10438 use lance::dataset::builder::DatasetBuilder;
10439 use lance_namespace::models::CreateTableVersionRequest;
10440
10441 let temp_dir = TempStrDir::default();
10442 let temp_path: &str = &temp_dir;
10443
10444 let namespace: Arc<dyn LanceNamespace> = Arc::new(
10445 DirectoryNamespaceBuilder::new(temp_path)
10446 .table_version_tracking_enabled(true)
10447 .build()
10448 .await
10449 .unwrap(),
10450 );
10451
10452 let schema = create_test_schema();
10454 let ipc_data = create_test_ipc_data(&schema);
10455 let mut create_req = CreateTableRequest::new();
10456 create_req.id = Some(vec!["test_table".to_string()]);
10457 namespace
10458 .create_table(create_req, bytes::Bytes::from(ipc_data))
10459 .await
10460 .unwrap();
10461
10462 let table_id = vec!["test_table".to_string()];
10464 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
10465 .await
10466 .unwrap()
10467 .load()
10468 .await
10469 .unwrap();
10470
10471 let versions_path = dataset.versions_dir();
10473 let manifest_metas: Vec<_> = dataset
10474 .object_store(None)
10475 .await
10476 .unwrap()
10477 .inner
10478 .list(Some(&versions_path))
10479 .try_collect()
10480 .await
10481 .unwrap();
10482
10483 let manifest_meta = manifest_metas
10484 .iter()
10485 .find(|m| {
10486 m.location
10487 .filename()
10488 .map(|f| f.ends_with(".manifest"))
10489 .unwrap_or(false)
10490 })
10491 .expect("No manifest file found");
10492
10493 let manifest_data = dataset
10495 .object_store(None)
10496 .await
10497 .unwrap()
10498 .inner
10499 .get(&manifest_meta.location)
10500 .await
10501 .unwrap()
10502 .bytes()
10503 .await
10504 .unwrap();
10505
10506 let staging_path = dataset.versions_dir().join("staging_manifest");
10508 dataset
10509 .object_store(None)
10510 .await
10511 .unwrap()
10512 .inner
10513 .put(&staging_path, manifest_data.into())
10514 .await
10515 .unwrap();
10516
10517 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
10520 create_version_req.id = Some(table_id.clone());
10521 create_version_req.naming_scheme = Some("V2".to_string());
10522
10523 let result = namespace.create_table_version(create_version_req).await;
10524 assert!(
10525 result.is_ok(),
10526 "create_table_version should succeed: {:?}",
10527 result
10528 );
10529
10530 let response = result.unwrap();
10532 let version_info = response
10533 .version
10534 .expect("response should contain version info");
10535 let version_2_path = Path::parse(&version_info.manifest_path).unwrap();
10536 let head_result = dataset
10537 .object_store(None)
10538 .await
10539 .unwrap()
10540 .inner
10541 .head(&version_2_path)
10542 .await;
10543 assert!(
10544 head_result.is_ok(),
10545 "Version 2 manifest should exist at {}",
10546 version_2_path
10547 );
10548
10549 let staging_head_result = dataset
10551 .object_store(None)
10552 .await
10553 .unwrap()
10554 .inner
10555 .head(&staging_path)
10556 .await;
10557 assert!(
10558 staging_head_result.is_err(),
10559 "Staging manifest should have been deleted after create_table_version"
10560 );
10561 }
10562
10563 #[tokio::test]
10564 async fn test_create_table_version_conflict() {
10565 use futures::TryStreamExt;
10568 use lance::dataset::builder::DatasetBuilder;
10569 use lance_namespace::models::CreateTableVersionRequest;
10570
10571 let temp_dir = TempStrDir::default();
10572 let temp_path: &str = &temp_dir;
10573
10574 let namespace: Arc<dyn LanceNamespace> = Arc::new(
10575 DirectoryNamespaceBuilder::new(temp_path)
10576 .table_version_tracking_enabled(true)
10577 .build()
10578 .await
10579 .unwrap(),
10580 );
10581
10582 let schema = create_test_schema();
10584 let ipc_data = create_test_ipc_data(&schema);
10585 let mut create_req = CreateTableRequest::new();
10586 create_req.id = Some(vec!["test_table".to_string()]);
10587 namespace
10588 .create_table(create_req, bytes::Bytes::from(ipc_data))
10589 .await
10590 .unwrap();
10591
10592 let table_id = vec!["test_table".to_string()];
10594 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
10595 .await
10596 .unwrap()
10597 .load()
10598 .await
10599 .unwrap();
10600
10601 let versions_path = dataset.versions_dir();
10603 let manifest_metas: Vec<_> = dataset
10604 .object_store(None)
10605 .await
10606 .unwrap()
10607 .inner
10608 .list(Some(&versions_path))
10609 .try_collect()
10610 .await
10611 .unwrap();
10612
10613 let manifest_meta = manifest_metas
10614 .iter()
10615 .find(|m| {
10616 m.location
10617 .filename()
10618 .map(|f| f.ends_with(".manifest"))
10619 .unwrap_or(false)
10620 })
10621 .expect("No manifest file found");
10622
10623 let manifest_data = dataset
10625 .object_store(None)
10626 .await
10627 .unwrap()
10628 .inner
10629 .get(&manifest_meta.location)
10630 .await
10631 .unwrap()
10632 .bytes()
10633 .await
10634 .unwrap();
10635
10636 let staging_path = dataset.versions_dir().join("staging_manifest");
10638 dataset
10639 .object_store(None)
10640 .await
10641 .unwrap()
10642 .inner
10643 .put(&staging_path, manifest_data.into())
10644 .await
10645 .unwrap();
10646
10647 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
10649 create_version_req.id = Some(table_id.clone());
10650 create_version_req.naming_scheme = Some("V2".to_string());
10651 let first_result = namespace.create_table_version(create_version_req).await;
10652 assert!(
10653 first_result.is_ok(),
10654 "First create_table_version for version 2 should succeed: {:?}",
10655 first_result
10656 );
10657
10658 let version_2_path = Path::parse(
10660 &first_result
10661 .unwrap()
10662 .version
10663 .expect("response should contain version info")
10664 .manifest_path,
10665 )
10666 .unwrap();
10667
10668 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
10670 create_version_req.id = Some(table_id.clone());
10671 create_version_req.naming_scheme = Some("V2".to_string());
10672
10673 let result = namespace.create_table_version(create_version_req).await;
10674 assert!(
10675 result.is_err(),
10676 "create_table_version should fail for existing version"
10677 );
10678
10679 let head_result = dataset
10681 .object_store(None)
10682 .await
10683 .unwrap()
10684 .inner
10685 .head(&version_2_path)
10686 .await;
10687 assert!(
10688 head_result.is_ok(),
10689 "Version 2 manifest should still exist at {}",
10690 version_2_path
10691 );
10692 }
10693
10694 #[tokio::test]
10695 async fn test_create_table_version_table_not_found() {
10696 use lance_namespace::models::CreateTableVersionRequest;
10697
10698 let temp_dir = TempStdDir::default();
10699 let temp_path = temp_dir.to_str().unwrap();
10700
10701 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10702 .table_version_tracking_enabled(true)
10703 .build()
10704 .await
10705 .unwrap();
10706
10707 let mut create_version_req =
10709 CreateTableVersionRequest::new(1, "/some/staging/path".to_string());
10710 create_version_req.id = Some(vec!["non_existent_table".to_string()]);
10711
10712 let result = namespace.create_table_version(create_version_req).await;
10713 assert!(
10714 result.is_err(),
10715 "create_table_version should fail for non-existent table"
10716 );
10717 let err_msg = result.unwrap_err().to_string();
10718 assert!(
10719 err_msg.contains("Table not found"),
10720 "Error should mention table not found, got: {}",
10721 err_msg
10722 );
10723 }
10724
10725 mod e2e_table_version_tracking {
10727 use super::*;
10728 use std::sync::atomic::{AtomicUsize, Ordering};
10729
10730 struct TrackingNamespace {
10732 inner: DirectoryNamespace,
10733 create_table_version_count: AtomicUsize,
10734 describe_table_version_count: AtomicUsize,
10735 list_table_versions_count: AtomicUsize,
10736 }
10737
10738 impl TrackingNamespace {
10739 fn new(inner: DirectoryNamespace) -> Self {
10740 Self {
10741 inner,
10742 create_table_version_count: AtomicUsize::new(0),
10743 describe_table_version_count: AtomicUsize::new(0),
10744 list_table_versions_count: AtomicUsize::new(0),
10745 }
10746 }
10747
10748 fn create_table_version_calls(&self) -> usize {
10749 self.create_table_version_count.load(Ordering::SeqCst)
10750 }
10751
10752 fn describe_table_version_calls(&self) -> usize {
10753 self.describe_table_version_count.load(Ordering::SeqCst)
10754 }
10755
10756 fn list_table_versions_calls(&self) -> usize {
10757 self.list_table_versions_count.load(Ordering::SeqCst)
10758 }
10759 }
10760
10761 impl std::fmt::Debug for TrackingNamespace {
10762 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
10763 f.debug_struct("TrackingNamespace")
10764 .field(
10765 "create_table_version_calls",
10766 &self.create_table_version_calls(),
10767 )
10768 .finish()
10769 }
10770 }
10771
10772 #[async_trait]
10773 impl LanceNamespace for TrackingNamespace {
10774 async fn create_namespace(
10775 &self,
10776 request: CreateNamespaceRequest,
10777 ) -> Result<CreateNamespaceResponse> {
10778 self.inner.create_namespace(request).await
10779 }
10780
10781 async fn describe_namespace(
10782 &self,
10783 request: DescribeNamespaceRequest,
10784 ) -> Result<DescribeNamespaceResponse> {
10785 self.inner.describe_namespace(request).await
10786 }
10787
10788 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
10789 self.inner.namespace_exists(request).await
10790 }
10791
10792 async fn list_namespaces(
10793 &self,
10794 request: ListNamespacesRequest,
10795 ) -> Result<ListNamespacesResponse> {
10796 self.inner.list_namespaces(request).await
10797 }
10798
10799 async fn drop_namespace(
10800 &self,
10801 request: DropNamespaceRequest,
10802 ) -> Result<DropNamespaceResponse> {
10803 self.inner.drop_namespace(request).await
10804 }
10805
10806 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
10807 self.inner.list_tables(request).await
10808 }
10809
10810 async fn describe_table(
10811 &self,
10812 request: DescribeTableRequest,
10813 ) -> Result<DescribeTableResponse> {
10814 self.inner.describe_table(request).await
10815 }
10816
10817 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
10818 self.inner.table_exists(request).await
10819 }
10820
10821 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
10822 self.inner.drop_table(request).await
10823 }
10824
10825 async fn create_table(
10826 &self,
10827 request: CreateTableRequest,
10828 request_data: Bytes,
10829 ) -> Result<CreateTableResponse> {
10830 self.inner.create_table(request, request_data).await
10831 }
10832
10833 async fn declare_table(
10834 &self,
10835 request: DeclareTableRequest,
10836 ) -> Result<DeclareTableResponse> {
10837 self.inner.declare_table(request).await
10838 }
10839
10840 async fn list_table_versions(
10841 &self,
10842 request: ListTableVersionsRequest,
10843 ) -> Result<ListTableVersionsResponse> {
10844 self.list_table_versions_count
10845 .fetch_add(1, Ordering::SeqCst);
10846 self.inner.list_table_versions(request).await
10847 }
10848
10849 async fn create_table_version(
10850 &self,
10851 request: CreateTableVersionRequest,
10852 ) -> Result<CreateTableVersionResponse> {
10853 self.create_table_version_count
10854 .fetch_add(1, Ordering::SeqCst);
10855 self.inner.create_table_version(request).await
10856 }
10857
10858 async fn describe_table_version(
10859 &self,
10860 request: DescribeTableVersionRequest,
10861 ) -> Result<DescribeTableVersionResponse> {
10862 self.describe_table_version_count
10863 .fetch_add(1, Ordering::SeqCst);
10864 self.inner.describe_table_version(request).await
10865 }
10866
10867 async fn batch_delete_table_versions(
10868 &self,
10869 request: BatchDeleteTableVersionsRequest,
10870 ) -> Result<BatchDeleteTableVersionsResponse> {
10871 self.inner.batch_delete_table_versions(request).await
10872 }
10873
10874 fn namespace_id(&self) -> String {
10875 self.inner.namespace_id()
10876 }
10877 }
10878
10879 #[tokio::test]
10880 async fn test_describe_table_returns_managed_versioning() {
10881 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableRequest};
10882
10883 let temp_dir = TempStdDir::default();
10884 let temp_path = temp_dir.to_str().unwrap();
10885
10886 let ns = DirectoryNamespaceBuilder::new(temp_path)
10888 .table_version_tracking_enabled(true)
10889 .manifest_enabled(true)
10890 .build()
10891 .await
10892 .unwrap();
10893
10894 let mut create_ns_req = CreateNamespaceRequest::new();
10896 create_ns_req.id = Some(vec!["workspace".to_string()]);
10897 ns.create_namespace(create_ns_req).await.unwrap();
10898
10899 let schema = create_test_schema();
10901 let ipc_data = create_test_ipc_data(&schema);
10902 let mut create_req = CreateTableRequest::new();
10903 create_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
10904 ns.create_table(create_req, bytes::Bytes::from(ipc_data))
10905 .await
10906 .unwrap();
10907
10908 let mut describe_req = DescribeTableRequest::new();
10910 describe_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
10911 let describe_resp = ns.describe_table(describe_req).await.unwrap();
10912
10913 assert_eq!(
10915 describe_resp.managed_versioning,
10916 Some(true),
10917 "managed_versioning should be true when table_version_tracking_enabled=true"
10918 );
10919 }
10920
10921 #[tokio::test]
10922 async fn test_external_manifest_store_invokes_namespace_apis() {
10923 use arrow::array::{Int32Array, StringArray};
10924 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
10925 use arrow::record_batch::RecordBatch;
10926 use lance::Dataset;
10927 use lance::dataset::builder::DatasetBuilder;
10928 use lance::dataset::{WriteMode, WriteParams};
10929 use lance_namespace::models::CreateNamespaceRequest;
10930
10931 let temp_dir = TempStdDir::default();
10932 let temp_path = temp_dir.to_str().unwrap();
10933
10934 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
10936 .table_version_tracking_enabled(true)
10937 .manifest_enabled(true)
10938 .build()
10939 .await
10940 .unwrap();
10941
10942 let tracking_ns = Arc::new(TrackingNamespace::new(inner_ns));
10943 let ns: Arc<dyn LanceNamespace> = tracking_ns.clone();
10944
10945 let mut create_ns_req = CreateNamespaceRequest::new();
10947 create_ns_req.id = Some(vec!["workspace".to_string()]);
10948 ns.create_namespace(create_ns_req).await.unwrap();
10949
10950 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
10952
10953 let arrow_schema = Arc::new(ArrowSchema::new(vec![
10955 Field::new("id", DataType::Int32, false),
10956 Field::new("name", DataType::Utf8, true),
10957 ]));
10958 let batch = RecordBatch::try_new(
10959 arrow_schema.clone(),
10960 vec![
10961 Arc::new(Int32Array::from(vec![1, 2, 3])),
10962 Arc::new(StringArray::from(vec!["a", "b", "c"])),
10963 ],
10964 )
10965 .unwrap();
10966
10967 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
10969 let write_params = WriteParams {
10970 mode: WriteMode::Create,
10971 ..Default::default()
10972 };
10973 let mut dataset = Dataset::write_into_namespace(
10974 batches,
10975 ns.clone(),
10976 table_id.clone(),
10977 Some(write_params),
10978 )
10979 .await
10980 .unwrap();
10981 assert_eq!(dataset.version().version, 1);
10982
10983 assert_eq!(
10985 tracking_ns.create_table_version_calls(),
10986 1,
10987 "create_table_version should have been called once during initial write_into_namespace"
10988 );
10989
10990 let append_batch = RecordBatch::try_new(
10992 arrow_schema.clone(),
10993 vec![
10994 Arc::new(Int32Array::from(vec![4, 5, 6])),
10995 Arc::new(StringArray::from(vec!["d", "e", "f"])),
10996 ],
10997 )
10998 .unwrap();
10999 let append_batches = RecordBatchIterator::new(vec![Ok(append_batch)], arrow_schema);
11000 dataset.append(append_batches, None).await.unwrap();
11001
11002 assert_eq!(
11003 tracking_ns.create_table_version_calls(),
11004 2,
11005 "create_table_version should have been called twice (once for create, once for append)"
11006 );
11007
11008 let initial_list_calls = tracking_ns.list_table_versions_calls();
11010 let latest_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
11011 .await
11012 .unwrap()
11013 .load()
11014 .await
11015 .unwrap();
11016 assert_eq!(latest_dataset.version().version, 2);
11017 assert_eq!(
11018 tracking_ns.list_table_versions_calls(),
11019 initial_list_calls + 1,
11020 "list_table_versions should have been called exactly once during checkout_latest"
11021 );
11022
11023 let initial_describe_calls = tracking_ns.describe_table_version_calls();
11025 let v1_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
11026 .await
11027 .unwrap()
11028 .with_version(1)
11029 .load()
11030 .await
11031 .unwrap();
11032 assert_eq!(v1_dataset.version().version, 1);
11033 assert_eq!(
11034 tracking_ns.describe_table_version_calls(),
11035 initial_describe_calls + 1,
11036 "describe_table_version should have been called exactly once during checkout to version 1"
11037 );
11038 }
11039
11040 #[tokio::test]
11041 async fn test_dataset_commit_with_external_manifest_store() {
11042 use arrow::array::{Int32Array, StringArray};
11043 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11044 use arrow::record_batch::RecordBatch;
11045 use futures::TryStreamExt;
11046 use lance::dataset::{Dataset, WriteMode, WriteParams};
11047 use lance_namespace::models::CreateNamespaceRequest;
11048 use lance_table::io::commit::ManifestNamingScheme;
11049
11050 let temp_dir = TempStdDir::default();
11051 let temp_path = temp_dir.to_str().unwrap();
11052
11053 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
11055 .table_version_tracking_enabled(true)
11056 .manifest_enabled(true)
11057 .build()
11058 .await
11059 .unwrap();
11060
11061 let tracking_ns: Arc<dyn LanceNamespace> = Arc::new(TrackingNamespace::new(inner_ns));
11062
11063 let mut create_ns_req = CreateNamespaceRequest::new();
11065 create_ns_req.id = Some(vec!["workspace".to_string()]);
11066 tracking_ns.create_namespace(create_ns_req).await.unwrap();
11067
11068 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11070 let arrow_schema = Arc::new(ArrowSchema::new(vec![
11071 Field::new("id", DataType::Int32, false),
11072 Field::new("name", DataType::Utf8, true),
11073 ]));
11074 let batch = RecordBatch::try_new(
11075 arrow_schema.clone(),
11076 vec![
11077 Arc::new(Int32Array::from(vec![1, 2, 3])),
11078 Arc::new(StringArray::from(vec!["a", "b", "c"])),
11079 ],
11080 )
11081 .unwrap();
11082 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
11083 let write_params = WriteParams {
11084 mode: WriteMode::Create,
11085 ..Default::default()
11086 };
11087 let dataset = Dataset::write_into_namespace(
11088 batches,
11089 tracking_ns.clone(),
11090 table_id.clone(),
11091 Some(write_params),
11092 )
11093 .await
11094 .unwrap();
11095 assert_eq!(dataset.version().version, 1);
11096
11097 let batch2 = RecordBatch::try_new(
11099 arrow_schema.clone(),
11100 vec![
11101 Arc::new(Int32Array::from(vec![4, 5, 6])),
11102 Arc::new(StringArray::from(vec!["d", "e", "f"])),
11103 ],
11104 )
11105 .unwrap();
11106 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
11107 let write_params = WriteParams {
11108 mode: WriteMode::Append,
11109 ..Default::default()
11110 };
11111 Dataset::write_into_namespace(
11112 batches,
11113 tracking_ns.clone(),
11114 table_id.clone(),
11115 Some(write_params),
11116 )
11117 .await
11118 .unwrap();
11119
11120 let manifest_metas: Vec<_> = dataset
11123 .object_store(None)
11124 .await
11125 .unwrap()
11126 .inner
11127 .list(Some(&dataset.versions_dir()))
11128 .try_collect()
11129 .await
11130 .unwrap();
11131 let version_2_found = manifest_metas.iter().any(|m| {
11132 m.location
11133 .filename()
11134 .map(|f| {
11135 f.ends_with(".manifest")
11136 && ManifestNamingScheme::V2.parse_version(f) == Some(2)
11137 })
11138 .unwrap_or(false)
11139 });
11140 assert!(
11141 version_2_found,
11142 "Version 2 manifest should exist in versions directory"
11143 );
11144 }
11145
11146 async fn create_ns_with_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
11148 use arrow::array::{Int32Array, StringArray};
11149 use arrow::ipc::writer::StreamWriter;
11150
11151 let (namespace, temp_dir) = create_test_namespace().await;
11152
11153 let schema = create_test_schema();
11154 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
11155 let arrow_schema = Arc::new(arrow_schema);
11156
11157 let id_array = Int32Array::from(vec![1, 2, 3]);
11158 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
11159 let batch = arrow::record_batch::RecordBatch::try_new(
11160 arrow_schema.clone(),
11161 vec![Arc::new(id_array), Arc::new(name_array)],
11162 )
11163 .unwrap();
11164
11165 let mut buffer = Vec::new();
11166 {
11167 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11168 writer.write(&batch).unwrap();
11169 writer.finish().unwrap();
11170 }
11171
11172 let mut request = CreateTableRequest::new();
11173 let table_id = vec!["test_ops_table".to_string()];
11174 request.id = Some(table_id.clone());
11175
11176 namespace
11177 .create_table(request, Bytes::from(buffer))
11178 .await
11179 .unwrap();
11180
11181 (namespace, temp_dir, table_id)
11182 }
11183
11184 #[tokio::test]
11185 async fn test_count_table_rows_basic() {
11186 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11187
11188 let request = CountTableRowsRequest {
11189 id: Some(table_id),
11190 version: None,
11191 predicate: None,
11192 ..Default::default()
11193 };
11194
11195 let count = namespace.count_table_rows(request).await.unwrap();
11196 assert_eq!(count, 3);
11197 }
11198
11199 #[tokio::test]
11200 async fn test_count_table_rows_with_predicate() {
11201 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11202
11203 let request = CountTableRowsRequest {
11204 id: Some(table_id),
11205 version: None,
11206 predicate: Some("id > 1".to_string()),
11207 ..Default::default()
11208 };
11209
11210 let count = namespace.count_table_rows(request).await.unwrap();
11211 assert_eq!(count, 2);
11212 }
11213
11214 #[tokio::test]
11215 async fn test_query_table_invalid_distance_type() {
11216 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
11217
11218 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11219 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
11220 multi_vector: None,
11221 });
11222
11223 let request = QueryTableRequest {
11224 id: Some(table_id),
11225 k: 2,
11226 vector,
11227 vector_column: Some("vector".to_string()),
11228 distance_type: Some("invalid_metric".to_string()),
11229 filter: None,
11230 offset: None,
11231 version: None,
11232 ..Default::default()
11233 };
11234
11235 let result = namespace.query_table(request).await;
11236 assert!(result.is_err());
11237 let err_msg = result.unwrap_err().to_string();
11238 assert!(
11239 err_msg.contains("Unknown distance type"),
11240 "Expected error about unknown distance type, got: {}",
11241 err_msg
11242 );
11243 }
11244
11245 #[tokio::test]
11246 async fn test_insert_into_table_append() {
11247 use arrow::array::{Int32Array, StringArray};
11248 use arrow::ipc::writer::StreamWriter;
11249
11250 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11251
11252 let schema = create_test_schema();
11254 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
11255 let arrow_schema = Arc::new(arrow_schema);
11256
11257 let id_array = Int32Array::from(vec![4, 5]);
11258 let name_array = StringArray::from(vec!["Dave", "Eve"]);
11259 let batch = arrow::record_batch::RecordBatch::try_new(
11260 arrow_schema.clone(),
11261 vec![Arc::new(id_array), Arc::new(name_array)],
11262 )
11263 .unwrap();
11264
11265 let mut buffer = Vec::new();
11266 {
11267 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11268 writer.write(&batch).unwrap();
11269 writer.finish().unwrap();
11270 }
11271
11272 let request = InsertIntoTableRequest {
11273 id: Some(table_id.clone()),
11274 mode: Some("append".to_string()),
11275 ..Default::default()
11276 };
11277
11278 let response = namespace
11279 .insert_into_table(request, Bytes::from(buffer))
11280 .await
11281 .unwrap();
11282 assert!(response.transaction_id.is_none());
11283
11284 let count_req = CountTableRowsRequest {
11286 id: Some(table_id),
11287 version: None,
11288 predicate: None,
11289 ..Default::default()
11290 };
11291 let count = namespace.count_table_rows(count_req).await.unwrap();
11292 assert_eq!(count, 5);
11293 }
11294
11295 #[tokio::test]
11296 async fn test_insert_into_table_overwrite() {
11297 use arrow::array::{Int32Array, StringArray};
11298 use arrow::ipc::writer::StreamWriter;
11299
11300 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11301
11302 let schema = create_test_schema();
11303 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
11304 let arrow_schema = Arc::new(arrow_schema);
11305
11306 let id_array = Int32Array::from(vec![10, 20]);
11307 let name_array = StringArray::from(vec!["X", "Y"]);
11308 let batch = arrow::record_batch::RecordBatch::try_new(
11309 arrow_schema.clone(),
11310 vec![Arc::new(id_array), Arc::new(name_array)],
11311 )
11312 .unwrap();
11313
11314 let mut buffer = Vec::new();
11315 {
11316 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11317 writer.write(&batch).unwrap();
11318 writer.finish().unwrap();
11319 }
11320
11321 let request = InsertIntoTableRequest {
11322 id: Some(table_id.clone()),
11323 mode: Some("overwrite".to_string()),
11324 ..Default::default()
11325 };
11326
11327 namespace
11328 .insert_into_table(request, Bytes::from(buffer))
11329 .await
11330 .unwrap();
11331
11332 let count_req = CountTableRowsRequest {
11334 id: Some(table_id),
11335 version: None,
11336 predicate: None,
11337 ..Default::default()
11338 };
11339 let count = namespace.count_table_rows(count_req).await.unwrap();
11340 assert_eq!(count, 2);
11341 }
11342
11343 #[tokio::test]
11344 async fn test_insert_into_table_empty_data() {
11345 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11346
11347 let request = InsertIntoTableRequest {
11348 id: Some(table_id),
11349 mode: None,
11350 ..Default::default()
11351 };
11352
11353 let result = namespace.insert_into_table(request, Bytes::new()).await;
11354 assert!(result.is_err());
11355 assert!(
11356 result
11357 .unwrap_err()
11358 .to_string()
11359 .contains("Arrow IPC stream) is required")
11360 );
11361 }
11362
11363 #[tokio::test]
11364 async fn test_insert_into_table_with_storage_options() {
11365 use arrow::array::{Int32Array, StringArray};
11366 use arrow::ipc::writer::StreamWriter;
11367
11368 let temp_dir = TempStdDir::default();
11369
11370 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
11372 .storage_option("allow_http", "true")
11373 .build()
11374 .await
11375 .unwrap();
11376
11377 let schema = create_test_schema();
11379 let ipc_data = create_test_ipc_data(&schema);
11380 let mut create_req = CreateTableRequest::new();
11381 let table_id = vec!["so_table".to_string()];
11382 create_req.id = Some(table_id.clone());
11383 namespace
11384 .create_table(create_req, Bytes::from(ipc_data))
11385 .await
11386 .unwrap();
11387
11388 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
11390 let arrow_schema = Arc::new(arrow_schema);
11391
11392 let id_array = Int32Array::from(vec![10, 20]);
11393 let name_array = StringArray::from(vec!["X", "Y"]);
11394 let batch = arrow::record_batch::RecordBatch::try_new(
11395 arrow_schema.clone(),
11396 vec![Arc::new(id_array), Arc::new(name_array)],
11397 )
11398 .unwrap();
11399
11400 let mut buffer = Vec::new();
11401 {
11402 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11403 writer.write(&batch).unwrap();
11404 writer.finish().unwrap();
11405 }
11406
11407 let request = InsertIntoTableRequest {
11408 id: Some(table_id.clone()),
11409 mode: Some("append".to_string()),
11410 ..Default::default()
11411 };
11412
11413 let response = namespace
11414 .insert_into_table(request, Bytes::from(buffer))
11415 .await
11416 .unwrap();
11417 assert!(response.transaction_id.is_none());
11418
11419 let count_req = CountTableRowsRequest {
11421 id: Some(table_id),
11422 version: None,
11423 predicate: None,
11424 ..Default::default()
11425 };
11426 let count = namespace.count_table_rows(count_req).await.unwrap();
11427 assert_eq!(count, 2);
11428 }
11429
11430 #[tokio::test]
11431 async fn test_query_table_basic() {
11432 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11433
11434 let request = QueryTableRequest {
11435 id: Some(table_id),
11436 k: 10,
11437 filter: None,
11438 offset: None,
11439 version: None,
11440 ..Default::default()
11441 };
11442
11443 let bytes = namespace.query_table(request).await.unwrap();
11444
11445 let cursor = Cursor::new(bytes.to_vec());
11447 let reader = FileReader::try_new(cursor, None).unwrap();
11448 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11449 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11450 assert_eq!(total_rows, 3);
11451 }
11452
11453 #[tokio::test]
11454 async fn test_query_table_with_filter() {
11455 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11456
11457 let request = QueryTableRequest {
11458 id: Some(table_id),
11459 k: 10,
11460 filter: Some("id <= 2".to_string()),
11461 offset: None,
11462 version: None,
11463 ..Default::default()
11464 };
11465
11466 let bytes = namespace.query_table(request).await.unwrap();
11467
11468 let cursor = Cursor::new(bytes.to_vec());
11469 let reader = FileReader::try_new(cursor, None).unwrap();
11470 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11471 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11472 assert_eq!(total_rows, 2);
11473 }
11474
11475 #[tokio::test]
11476 async fn test_query_table_with_limit_and_offset() {
11477 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11478
11479 let request = QueryTableRequest {
11480 id: Some(table_id),
11481 k: 2,
11482 filter: None,
11483 offset: Some(1),
11484 version: None,
11485 ..Default::default()
11486 };
11487
11488 let bytes = namespace.query_table(request).await.unwrap();
11489
11490 let cursor = Cursor::new(bytes.to_vec());
11491 let reader = FileReader::try_new(cursor, None).unwrap();
11492 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11493 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11494 assert_eq!(total_rows, 2);
11495 }
11496
11497 #[tokio::test]
11498 async fn test_query_table_no_limit() {
11499 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11500
11501 let request = QueryTableRequest {
11503 id: Some(table_id),
11504 k: 0,
11505 filter: None,
11506 offset: None,
11507 version: None,
11508 ..Default::default()
11509 };
11510
11511 let bytes = namespace.query_table(request).await.unwrap();
11512
11513 let cursor = Cursor::new(bytes.to_vec());
11514 let reader = FileReader::try_new(cursor, None).unwrap();
11515 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11516 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11517 assert_eq!(total_rows, 3);
11518 }
11519
11520 #[tokio::test]
11521 async fn test_query_table_with_columns() {
11522 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11523
11524 let columns = Box::new(lance_namespace::models::QueryTableRequestColumns {
11525 column_names: Some(vec!["id".to_string()]),
11526 column_aliases: None,
11527 });
11528
11529 let request = QueryTableRequest {
11530 id: Some(table_id),
11531 k: 10,
11532 filter: None,
11533 offset: None,
11534 version: None,
11535 columns: Some(columns),
11536 ..Default::default()
11537 };
11538
11539 let bytes = namespace.query_table(request).await.unwrap();
11540
11541 let cursor = Cursor::new(bytes.to_vec());
11542 let reader = FileReader::try_new(cursor, None).unwrap();
11543 let schema = reader.schema();
11544 assert_eq!(schema.fields().len(), 1);
11545 assert_eq!(schema.field(0).name(), "id");
11546 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11547 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11548 assert_eq!(total_rows, 3);
11549 }
11550
11551 #[tokio::test]
11552 async fn test_count_table_rows_with_version() {
11553 use arrow::array::{Int32Array, StringArray};
11554 use arrow::ipc::writer::StreamWriter;
11555
11556 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11557
11558 let schema = create_test_schema();
11560 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
11561 let arrow_schema = Arc::new(arrow_schema);
11562
11563 let id_array = Int32Array::from(vec![4, 5]);
11564 let name_array = StringArray::from(vec!["Dave", "Eve"]);
11565 let batch = arrow::record_batch::RecordBatch::try_new(
11566 arrow_schema.clone(),
11567 vec![Arc::new(id_array), Arc::new(name_array)],
11568 )
11569 .unwrap();
11570
11571 let mut buffer = Vec::new();
11572 {
11573 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11574 writer.write(&batch).unwrap();
11575 writer.finish().unwrap();
11576 }
11577
11578 let request = InsertIntoTableRequest {
11579 id: Some(table_id.clone()),
11580 mode: None,
11581 ..Default::default()
11582 };
11583 namespace
11584 .insert_into_table(request, Bytes::from(buffer))
11585 .await
11586 .unwrap();
11587
11588 let count_req = CountTableRowsRequest {
11590 id: Some(table_id.clone()),
11591 version: Some(1),
11592 predicate: None,
11593 ..Default::default()
11594 };
11595 let count = namespace.count_table_rows(count_req).await.unwrap();
11596 assert_eq!(count, 3);
11597
11598 let count_req = CountTableRowsRequest {
11600 id: Some(table_id),
11601 version: None,
11602 predicate: None,
11603 ..Default::default()
11604 };
11605 let count = namespace.count_table_rows(count_req).await.unwrap();
11606 assert_eq!(count, 5);
11607 }
11608
11609 #[tokio::test]
11610 async fn test_query_table_with_version() {
11611 use arrow::array::{Int32Array, StringArray};
11612 use arrow::ipc::writer::StreamWriter;
11613
11614 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11615
11616 let schema = create_test_schema();
11618 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
11619 let arrow_schema = Arc::new(arrow_schema);
11620
11621 let id_array = Int32Array::from(vec![4, 5]);
11622 let name_array = StringArray::from(vec!["Dave", "Eve"]);
11623 let batch = arrow::record_batch::RecordBatch::try_new(
11624 arrow_schema.clone(),
11625 vec![Arc::new(id_array), Arc::new(name_array)],
11626 )
11627 .unwrap();
11628
11629 let mut buffer = Vec::new();
11630 {
11631 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11632 writer.write(&batch).unwrap();
11633 writer.finish().unwrap();
11634 }
11635
11636 let request = InsertIntoTableRequest {
11637 id: Some(table_id.clone()),
11638 mode: None,
11639 ..Default::default()
11640 };
11641 namespace
11642 .insert_into_table(request, Bytes::from(buffer))
11643 .await
11644 .unwrap();
11645
11646 let request = QueryTableRequest {
11648 id: Some(table_id.clone()),
11649 k: 100,
11650 filter: None,
11651 offset: None,
11652 version: Some(1),
11653 ..Default::default()
11654 };
11655
11656 let bytes = namespace.query_table(request).await.unwrap();
11657 let cursor = Cursor::new(bytes.to_vec());
11658 let reader = FileReader::try_new(cursor, None).unwrap();
11659 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11660 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11661 assert_eq!(total_rows, 3);
11662
11663 let request = QueryTableRequest {
11665 id: Some(table_id),
11666 k: 100,
11667 filter: None,
11668 offset: None,
11669 version: None,
11670 ..Default::default()
11671 };
11672
11673 let bytes = namespace.query_table(request).await.unwrap();
11674 let cursor = Cursor::new(bytes.to_vec());
11675 let reader = FileReader::try_new(cursor, None).unwrap();
11676 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11677 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11678 assert_eq!(total_rows, 5);
11679 }
11680
11681 async fn create_ns_with_vector_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
11684 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
11685 use arrow::ipc::writer::StreamWriter;
11686
11687 let (namespace, temp_dir) = create_test_namespace().await;
11688
11689 let arrow_schema = Arc::new(arrow::datatypes::Schema::new(vec![
11691 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
11692 arrow::datatypes::Field::new(
11693 "vector",
11694 arrow::datatypes::DataType::FixedSizeList(
11695 Arc::new(arrow::datatypes::Field::new(
11696 "item",
11697 arrow::datatypes::DataType::Float32,
11698 true,
11699 )),
11700 4,
11701 ),
11702 true,
11703 ),
11704 ]));
11705
11706 let id_array = Int32Array::from(vec![1, 2, 3]);
11707 let values = Float32Array::from(vec![
11708 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ]);
11712 let vector_array = FixedSizeListArray::try_new(
11713 Arc::new(arrow::datatypes::Field::new(
11714 "item",
11715 arrow::datatypes::DataType::Float32,
11716 true,
11717 )),
11718 4,
11719 Arc::new(values),
11720 None,
11721 )
11722 .unwrap();
11723
11724 let batch = arrow::record_batch::RecordBatch::try_new(
11725 arrow_schema.clone(),
11726 vec![Arc::new(id_array), Arc::new(vector_array)],
11727 )
11728 .unwrap();
11729
11730 let mut buffer = Vec::new();
11731 {
11732 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
11733 writer.write(&batch).unwrap();
11734 writer.finish().unwrap();
11735 }
11736
11737 let table_name = "vector_table";
11739 let table_uri = format!("{}/{}.lance", temp_dir.to_str().unwrap(), table_name);
11740 let reader = arrow::record_batch::RecordBatchIterator::new(
11741 vec![Ok(batch)],
11742 arrow_schema.clone(),
11743 );
11744 Dataset::write(reader, &table_uri, None).await.unwrap();
11745
11746 let table_id = vec![table_name.to_string()];
11747 (namespace, temp_dir, table_id)
11748 }
11749
11750 #[tokio::test]
11751 async fn test_query_table_vector_search() {
11752 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
11753
11754 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11755 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
11756 multi_vector: None,
11757 });
11758
11759 let request = QueryTableRequest {
11760 id: Some(table_id),
11761 k: 2,
11762 vector,
11763 filter: None,
11764 offset: None,
11765 version: None,
11766 ..Default::default()
11767 };
11768
11769 let bytes = namespace.query_table(request).await.unwrap();
11770
11771 let cursor = Cursor::new(bytes.to_vec());
11772 let reader = FileReader::try_new(cursor, None).unwrap();
11773 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11774 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11775 assert_eq!(total_rows, 2);
11776 }
11777
11778 #[tokio::test]
11779 async fn test_query_table_vector_search_with_distance_type() {
11780 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
11781
11782 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11783 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
11784 multi_vector: None,
11785 });
11786
11787 let request = QueryTableRequest {
11788 id: Some(table_id),
11789 k: 3,
11790 vector,
11791 filter: None,
11792 offset: None,
11793 version: None,
11794 distance_type: Some("cosine".to_string()),
11795 ..Default::default()
11796 };
11797
11798 let bytes = namespace.query_table(request).await.unwrap();
11799
11800 let cursor = Cursor::new(bytes.to_vec());
11801 let reader = FileReader::try_new(cursor, None).unwrap();
11802 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11803 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11804 assert_eq!(total_rows, 3);
11805 }
11806
11807 #[tokio::test]
11808 async fn test_query_table_vector_search_with_filter() {
11809 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
11810
11811 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11812 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
11813 multi_vector: None,
11814 });
11815
11816 let request = QueryTableRequest {
11817 id: Some(table_id),
11818 k: 10,
11819 vector,
11820 filter: Some("id <= 2".to_string()),
11821 offset: None,
11822 version: None,
11823 ..Default::default()
11824 };
11825
11826 let bytes = namespace.query_table(request).await.unwrap();
11827
11828 let cursor = Cursor::new(bytes.to_vec());
11829 let reader = FileReader::try_new(cursor, None).unwrap();
11830 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11831 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11832 assert!(total_rows <= 2);
11833 }
11834
11835 #[tokio::test]
11836 async fn test_query_table_vector_search_with_nprobes_and_refine() {
11837 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
11838
11839 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11840 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
11841 multi_vector: None,
11842 });
11843
11844 let request = QueryTableRequest {
11845 id: Some(table_id),
11846 k: 2,
11847 vector,
11848 filter: None,
11849 offset: None,
11850 version: None,
11851 nprobes: Some(1),
11852 refine_factor: Some(1),
11853 prefilter: Some(true),
11854 ..Default::default()
11855 };
11856
11857 let bytes = namespace.query_table(request).await.unwrap();
11858
11859 let cursor = Cursor::new(bytes.to_vec());
11860 let reader = FileReader::try_new(cursor, None).unwrap();
11861 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11862 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11863 assert_eq!(total_rows, 2);
11864 }
11865
11866 #[tokio::test]
11867 async fn test_namespace_id() {
11868 let (namespace, _temp_dir) = create_test_namespace().await;
11869 let id = namespace.namespace_id();
11870 assert!(id.contains("DirectoryNamespace"));
11871 assert!(id.contains("root"));
11872 }
11873
11874 #[tokio::test]
11875 async fn test_query_table_empty_table() {
11876 let (namespace, _temp_dir) = create_test_namespace().await;
11877
11878 let schema = create_test_schema();
11880 let ipc_data = create_test_ipc_data(&schema);
11881 let mut create_request = CreateTableRequest::new();
11882 create_request.id = Some(vec!["empty_table".to_string()]);
11883 namespace
11884 .create_table(create_request, bytes::Bytes::from(ipc_data))
11885 .await
11886 .unwrap();
11887
11888 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11890 single_vector: None,
11891 multi_vector: None,
11892 });
11893 let request = QueryTableRequest {
11894 id: Some(vec!["empty_table".to_string()]),
11895 k: 10,
11896 vector,
11897 ..Default::default()
11898 };
11899 let bytes = namespace.query_table(request).await.unwrap();
11900
11901 let cursor = Cursor::new(bytes.to_vec());
11902 let reader = FileReader::try_new(cursor, None).unwrap();
11903 let batches: Vec<_> = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
11904 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11905 assert_eq!(total_rows, 0, "empty table should yield no rows");
11906 }
11907
11908 #[tokio::test]
11909 async fn test_query_table_with_plain_filter_no_vector() {
11910 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11911
11912 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
11914 single_vector: None,
11915 multi_vector: None,
11916 });
11917 let request = QueryTableRequest {
11918 id: Some(table_id),
11919 k: 0,
11920 vector,
11921 filter: Some("id > 1".to_string()),
11922 ..Default::default()
11923 };
11924 let bytes = namespace.query_table(request).await.unwrap();
11925
11926 let cursor = Cursor::new(bytes.to_vec());
11927 let reader = FileReader::try_new(cursor, None).unwrap();
11928 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
11929 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
11930 assert!(total_rows > 0);
11931 assert!(total_rows < 3);
11932 }
11933
11934 #[tokio::test]
11937 async fn test_update_full_table() {
11938 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11939
11940 let base_version = open_dataset(&namespace, &table_id[0])
11942 .await
11943 .version()
11944 .version;
11945
11946 let request = UpdateTableRequest {
11947 id: Some(table_id.clone()),
11948 updates: vec![vec!["name".to_string(), "'updated'".to_string()]],
11949 predicate: None,
11950 ..Default::default()
11951 };
11952
11953 let response = namespace.update_table(request).await.unwrap();
11954 assert_eq!(response.updated_rows, 3);
11955 assert!(response.version as u64 > base_version);
11956
11957 let count_req = CountTableRowsRequest {
11959 id: Some(table_id),
11960 version: None,
11961 predicate: Some("name = 'updated'".to_string()),
11962 ..Default::default()
11963 };
11964 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 3);
11965 }
11966
11967 #[tokio::test]
11968 async fn test_update_with_predicate() {
11969 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
11970
11971 let request = UpdateTableRequest {
11972 id: Some(table_id.clone()),
11973 updates: vec![vec!["name".to_string(), "'matched'".to_string()]],
11974 predicate: Some("id > 1".to_string()),
11975 ..Default::default()
11976 };
11977
11978 let response = namespace.update_table(request).await.unwrap();
11979 assert_eq!(response.updated_rows, 2);
11980
11981 let untouched = CountTableRowsRequest {
11983 id: Some(table_id.clone()),
11984 version: None,
11985 predicate: Some("name = 'Alice'".to_string()),
11986 ..Default::default()
11987 };
11988 assert_eq!(namespace.count_table_rows(untouched).await.unwrap(), 1);
11989
11990 let touched = CountTableRowsRequest {
11991 id: Some(table_id),
11992 version: None,
11993 predicate: Some("name = 'matched'".to_string()),
11994 ..Default::default()
11995 };
11996 assert_eq!(namespace.count_table_rows(touched).await.unwrap(), 2);
11997 }
11998
11999 #[tokio::test]
12000 async fn test_update_invalid_expression_returns_invalid_input() {
12001 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12002
12003 let request = UpdateTableRequest {
12004 id: Some(table_id),
12005 updates: vec![vec!["name".to_string(), "no_such_column + 1".to_string()]],
12007 predicate: None,
12008 ..Default::default()
12009 };
12010
12011 let err = namespace.update_table(request).await.unwrap_err();
12012 let msg = err.to_string();
12013 assert!(
12014 msg.contains("Invalid input"),
12015 "expected Invalid input error, got: {}",
12016 msg
12017 );
12018 }
12019
12020 #[tokio::test]
12021 async fn test_update_rejects_duplicate_columns() {
12022 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12023
12024 let request = UpdateTableRequest {
12025 id: Some(table_id),
12026 updates: vec![
12027 vec!["name".to_string(), "'a'".to_string()],
12028 vec!["name".to_string(), "'b'".to_string()],
12029 ],
12030 predicate: None,
12031 ..Default::default()
12032 };
12033
12034 let err = namespace.update_table(request).await.unwrap_err();
12035 let msg = err.to_string();
12036 assert!(
12037 msg.contains("Invalid input") && msg.contains("more than once"),
12038 "expected duplicate column InvalidInput error, got: {}",
12039 msg
12040 );
12041 }
12042
12043 #[tokio::test]
12044 async fn test_delete_with_predicate() {
12045 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12046
12047 let request = DeleteFromTableRequest {
12048 id: Some(table_id.clone()),
12049 predicate: "id > 1".to_string(),
12050 ..Default::default()
12051 };
12052
12053 let response = namespace.delete_from_table(request).await.unwrap();
12054 assert!(response.version.is_some());
12055
12056 let count_req = CountTableRowsRequest {
12057 id: Some(table_id),
12058 version: None,
12059 predicate: None,
12060 ..Default::default()
12061 };
12062 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 1);
12064 }
12065
12066 #[tokio::test]
12067 async fn test_delete_empty_predicate_returns_invalid_input() {
12068 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12069
12070 let request = DeleteFromTableRequest {
12071 id: Some(table_id),
12072 predicate: " ".to_string(),
12073 ..Default::default()
12074 };
12075
12076 let err = namespace.delete_from_table(request).await.unwrap_err();
12077 let msg = err.to_string();
12078 assert!(
12079 msg.contains("Invalid input") && msg.contains("non-empty predicate"),
12080 "expected non-empty predicate InvalidInput error, got: {}",
12081 msg
12082 );
12083 }
12084
12085 #[tokio::test]
12086 async fn test_delete_table_not_found() {
12087 let (namespace, _temp_dir) = create_test_namespace().await;
12088
12089 let request = DeleteFromTableRequest {
12090 id: Some(vec!["does_not_exist".to_string()]),
12091 predicate: "id = 1".to_string(),
12092 ..Default::default()
12093 };
12094
12095 let err = namespace.delete_from_table(request).await.unwrap_err();
12096 let msg = err.to_string();
12097 assert!(
12098 msg.contains("Table not found"),
12099 "expected TableNotFound for missing table, got: {}",
12100 msg
12101 );
12102 }
12103
12104 #[tokio::test]
12105 async fn test_delete_invalid_predicate_returns_invalid_input() {
12106 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12107
12108 let request = DeleteFromTableRequest {
12112 id: Some(table_id),
12113 predicate: "no_such_column = 1".to_string(),
12114 ..Default::default()
12115 };
12116
12117 let err = namespace.delete_from_table(request).await.unwrap_err();
12118 let lance_core::Error::Namespace { source, .. } = &err else {
12119 panic!("expected a Namespace error, got: {}", err);
12120 };
12121 let ns_err = source
12122 .downcast_ref::<NamespaceError>()
12123 .expect("expected a NamespaceError source");
12124 assert_eq!(
12125 ns_err.code(),
12126 lance_namespace::ErrorCode::InvalidInput,
12127 "expected InvalidInput for an invalid delete predicate, got: {}",
12128 err
12129 );
12130 }
12131 }
12132
12133 #[tokio::test]
12134 async fn test_list_all_tables() {
12135 use lance_namespace::models::ListTablesRequest;
12136
12137 let (namespace, _temp_dir) = create_test_namespace().await;
12138 create_scalar_table(&namespace, "alpha").await;
12139 create_scalar_table(&namespace, "beta").await;
12140
12141 let request = ListTablesRequest {
12142 id: Some(vec![]),
12143 page_token: None,
12144 limit: None,
12145 ..Default::default()
12146 };
12147 let response = namespace.list_all_tables(request).await.unwrap();
12148 let mut tables = response.tables;
12149 tables.sort();
12150 assert_eq!(tables, vec!["alpha", "beta"]);
12151 }
12152
12153 #[tokio::test]
12154 async fn test_restore_table() {
12155 use lance_namespace::models::RestoreTableRequest;
12156
12157 let (namespace, _temp_dir) = create_test_namespace().await;
12158 create_scalar_table(&namespace, "users").await;
12159
12160 create_scalar_index(&namespace, "users", "users_id_idx").await;
12162
12163 let dataset = open_dataset(&namespace, "users").await;
12164 let current_version = dataset.version().version;
12165 assert!(current_version >= 2, "Should have at least 2 versions");
12166
12167 let mut restore_req = RestoreTableRequest::new(1);
12169 restore_req.id = Some(vec!["users".to_string()]);
12170 let response = namespace.restore_table(restore_req).await.unwrap();
12171
12172 assert!(
12174 response.transaction_id.is_some(),
12175 "restore_table should return a transaction_id"
12176 );
12177
12178 let dataset_after = open_dataset(&namespace, "users").await;
12180 assert!(
12181 dataset_after.version().version > current_version,
12182 "Restore should create a new version"
12183 );
12184 }
12185
12186 #[tokio::test]
12187 async fn test_alter_table_add_columns() {
12188 use lance_namespace::models::{
12189 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
12190 };
12191
12192 let (namespace, _temp_dir) = create_test_namespace().await;
12193
12194 let schema = create_test_schema();
12196 let ipc_data = create_test_ipc_data(&schema);
12197 let mut create_request = CreateTableRequest::new();
12198 create_request.id = Some(vec!["test_table".to_string()]);
12199 namespace
12200 .create_table(create_request, bytes::Bytes::from(ipc_data))
12201 .await
12202 .unwrap();
12203
12204 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
12206 new_col.expression = Some(Some("id * 2".to_string()));
12207 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
12208 add_request.id = Some(vec!["test_table".to_string()]);
12209
12210 let response = namespace
12211 .alter_table_add_columns(add_request)
12212 .await
12213 .unwrap();
12214 assert!(
12215 response.version > 1,
12216 "Version should increment after adding columns"
12217 );
12218
12219 let mut describe_request = DescribeTableRequest::new();
12221 describe_request.id = Some(vec!["test_table".to_string()]);
12222 describe_request.load_detailed_metadata = Some(true);
12223 let describe_response = namespace.describe_table(describe_request).await.unwrap();
12224 assert!(describe_response.schema.is_some());
12225
12226 let resp_schema = describe_response.schema.unwrap();
12227 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
12228 assert!(
12229 field_names.contains(&"doubled_id"),
12230 "Column 'doubled_id' should exist, got: {:?}",
12231 field_names
12232 );
12233 }
12234
12235 #[tokio::test]
12236 async fn test_update_table_schema_metadata() {
12237 use lance_namespace::models::UpdateTableSchemaMetadataRequest;
12238
12239 let (namespace, _temp_dir) = create_test_namespace().await;
12240 create_scalar_table(&namespace, "products").await;
12241
12242 let mut metadata = HashMap::new();
12243 metadata.insert("owner".to_string(), "team_a".to_string());
12244 metadata.insert("version".to_string(), "1.0".to_string());
12245
12246 let mut req = UpdateTableSchemaMetadataRequest::new();
12247 req.id = Some(vec!["products".to_string()]);
12248 req.metadata = Some(metadata.clone());
12249
12250 let response = namespace.update_table_schema_metadata(req).await.unwrap();
12251
12252 assert!(response.metadata.is_some());
12253 let returned = response.metadata.unwrap();
12254 assert_eq!(returned.get("owner"), Some(&"team_a".to_string()));
12255 assert_eq!(returned.get("version"), Some(&"1.0".to_string()));
12256 assert!(
12257 response.transaction_id.is_some(),
12258 "update_table_schema_metadata should return a transaction_id"
12259 );
12260 }
12261
12262 #[tokio::test]
12263 async fn test_alter_table_add_columns_missing_id() {
12264 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
12265
12266 let (namespace, _temp_dir) = create_test_namespace().await;
12267
12268 let new_col = AddColumnsEntry::new("col".to_string());
12269 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
12270 let result = namespace.alter_table_add_columns(request).await;
12271 assert!(result.is_err(), "Should fail when table ID is missing");
12272 }
12273
12274 #[tokio::test]
12275 async fn test_alter_table_alter_columns_rename() {
12276 use lance_namespace::models::{
12277 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
12278 };
12279
12280 let (namespace, _temp_dir) = create_test_namespace().await;
12281
12282 let schema = create_test_schema();
12284 let ipc_data = create_test_ipc_data(&schema);
12285 let mut create_request = CreateTableRequest::new();
12286 create_request.id = Some(vec!["test_table".to_string()]);
12287 namespace
12288 .create_table(create_request, bytes::Bytes::from(ipc_data))
12289 .await
12290 .unwrap();
12291
12292 let mut entry = AlterColumnsEntry::new("name".to_string());
12294 entry.rename = Some(Some("full_name".to_string()));
12295 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
12296 alter_request.id = Some(vec!["test_table".to_string()]);
12297
12298 let response = namespace
12299 .alter_table_alter_columns(alter_request)
12300 .await
12301 .unwrap();
12302 assert!(
12303 response.version > 1,
12304 "Version should increment after altering columns"
12305 );
12306
12307 let mut describe_request = DescribeTableRequest::new();
12309 describe_request.id = Some(vec!["test_table".to_string()]);
12310 describe_request.load_detailed_metadata = Some(true);
12311 let describe_response = namespace.describe_table(describe_request).await.unwrap();
12312 assert!(describe_response.schema.is_some());
12313
12314 let resp_schema = describe_response.schema.unwrap();
12315 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
12316 assert!(
12317 field_names.contains(&"full_name"),
12318 "Column should be renamed to 'full_name', got: {:?}",
12319 field_names
12320 );
12321 assert!(
12322 !field_names.contains(&"name"),
12323 "Old column 'name' should not exist, got: {:?}",
12324 field_names
12325 );
12326 }
12327
12328 #[tokio::test]
12329 async fn test_get_table_stats() {
12330 use lance_namespace::models::GetTableStatsRequest;
12331
12332 let (namespace, _temp_dir) = create_test_namespace().await;
12333 create_scalar_table(&namespace, "items").await;
12334 create_scalar_index(&namespace, "items", "items_id_idx").await;
12335
12336 let mut req = GetTableStatsRequest::new();
12337 req.id = Some(vec!["items".to_string()]);
12338
12339 let response = namespace.get_table_stats(req).await.unwrap();
12340 assert_eq!(response.num_rows, 3);
12341 assert_eq!(response.num_indices, 1);
12342 }
12343
12344 #[tokio::test]
12345 async fn test_explain_table_query_plan() {
12346 use lance_namespace::models::QueryTableRequestVector;
12347 use lance_namespace::models::{ExplainTableQueryPlanRequest, QueryTableRequest};
12348
12349 let (namespace, _temp_dir) = create_test_namespace().await;
12350 create_scalar_table(&namespace, "catalog").await;
12351
12352 let mut query = QueryTableRequest::new(1, QueryTableRequestVector::new());
12353 query.filter = Some("id > 1".to_string());
12354 query.columns = Some(Box::new(QueryTableRequestColumns {
12355 column_names: Some(vec!["id".to_string(), "name".to_string()]),
12356 column_aliases: None,
12357 }));
12358 query.with_row_id = Some(true);
12359
12360 let mut req = ExplainTableQueryPlanRequest::new(query);
12361 req.id = Some(vec!["catalog".to_string()]);
12362
12363 let plan_str = namespace.explain_table_query_plan(req).await.unwrap();
12364 assert_plan_contains_all(
12365 &plan_str,
12366 &[
12367 "ProjectionExec: expr=[id@0 as id, name@2 as name",
12368 "projection=[name], source=stream(_rowid)",
12369 "LanceRead: uri=",
12370 "projection=[id]",
12371 "row_id=true, row_addr=false",
12372 "full_filter=id > Int32(1)",
12373 "refine_filter=id > Int32(1)",
12374 ],
12375 "Filtered explain plan should preserve late materialization and filter pushdown",
12376 );
12377 }
12378
12379 #[tokio::test]
12380 async fn test_alter_table_alter_columns_missing_id() {
12381 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
12382
12383 let (namespace, _temp_dir) = create_test_namespace().await;
12384
12385 let entry = AlterColumnsEntry::new("name".to_string());
12386 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
12387 let result = namespace.alter_table_alter_columns(request).await;
12388 assert!(result.is_err(), "Should fail when table ID is missing");
12389 }
12390
12391 #[tokio::test]
12392 async fn test_alter_table_drop_columns() {
12393 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
12394
12395 let (namespace, _temp_dir) = create_test_namespace().await;
12396
12397 let schema = create_test_schema();
12399 let ipc_data = create_test_ipc_data(&schema);
12400 let mut create_request = CreateTableRequest::new();
12401 create_request.id = Some(vec!["test_table".to_string()]);
12402 namespace
12403 .create_table(create_request, bytes::Bytes::from(ipc_data))
12404 .await
12405 .unwrap();
12406
12407 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
12409 drop_request.id = Some(vec!["test_table".to_string()]);
12410
12411 let response = namespace
12412 .alter_table_drop_columns(drop_request)
12413 .await
12414 .unwrap();
12415 assert!(
12416 response.version > 1,
12417 "Version should increment after dropping columns"
12418 );
12419
12420 let mut describe_request = DescribeTableRequest::new();
12422 describe_request.id = Some(vec!["test_table".to_string()]);
12423 describe_request.load_detailed_metadata = Some(true);
12424 let describe_response = namespace.describe_table(describe_request).await.unwrap();
12425 assert!(describe_response.schema.is_some());
12426
12427 let resp_schema = describe_response.schema.unwrap();
12428 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
12429 assert!(
12430 !field_names.contains(&"name"),
12431 "Column 'name' should be dropped, got: {:?}",
12432 field_names
12433 );
12434 assert!(
12435 field_names.contains(&"id"),
12436 "Column 'id' should still exist, got: {:?}",
12437 field_names
12438 );
12439 }
12440
12441 #[tokio::test]
12442 async fn test_analyze_table_query_plan() {
12443 use lance_namespace::models::AnalyzeTableQueryPlanRequest;
12444 use lance_namespace::models::QueryTableRequestVector;
12445
12446 let (namespace, _temp_dir) = create_test_namespace().await;
12447 create_scalar_table(&namespace, "catalog").await;
12448
12449 let mut req = AnalyzeTableQueryPlanRequest::new(1, QueryTableRequestVector::new());
12450 req.id = Some(vec!["catalog".to_string()]);
12451 req.filter = Some("id > 0".to_string());
12452 req.columns = Some(Box::new(QueryTableRequestColumns {
12453 column_names: Some(vec!["id".to_string(), "name".to_string()]),
12454 column_aliases: None,
12455 }));
12456 req.with_row_id = Some(true);
12457
12458 let analysis_str = namespace.analyze_table_query_plan(req).await.unwrap();
12459 assert_plan_contains_all(
12460 &analysis_str,
12461 &[
12462 "AnalyzeExec verbose=true",
12463 "ProjectionExec: elapsed=",
12464 "expr=[id@0 as id, name@2 as name",
12465 "projection=[name], source=stream(_rowid)",
12466 "LanceRead: elapsed=",
12467 "projection=[id]",
12468 "row_id=true, row_addr=false",
12469 "full_filter=id > Int32(0)",
12470 "refine_filter=id > Int32(0)",
12471 "metrics=[output_rows=",
12472 ],
12473 "Filtered analyze plan should preserve late materialization and filter pushdown",
12474 );
12475 }
12476
12477 #[tokio::test]
12478 async fn test_dir_listing_no_extra_calls_without_migration() {
12479 let temp_dir = TempStdDir::default();
12480 let temp_path = temp_dir.to_str().unwrap();
12481 let root_uri = file_object_store_uri(temp_path);
12482 let listing_count = Arc::new(AtomicUsize::new(0));
12483 let session = build_listing_counting_session(listing_count.clone());
12484
12485 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
12487 .session(session.clone())
12488 .manifest_enabled(false)
12489 .dir_listing_enabled(true)
12490 .build()
12491 .await
12492 .unwrap();
12493
12494 let schema = create_test_schema();
12495 let ipc_data = create_test_ipc_data(&schema);
12496 let mut create_req = CreateTableRequest::new();
12497 create_req.id = Some(vec!["test_table".to_string()]);
12498 dir_only_ns
12499 .create_table(create_req, Bytes::from(ipc_data))
12500 .await
12501 .unwrap();
12502
12503 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
12505 .session(session)
12506 .manifest_enabled(true)
12507 .dir_listing_enabled(true)
12508 .dir_listing_to_manifest_migration_enabled(false)
12509 .build()
12510 .await
12511 .unwrap();
12512
12513 listing_count.store(0, Ordering::SeqCst);
12515
12516 let mut exists_req = TableExistsRequest::new();
12518 exists_req.id = Some(vec!["test_table".to_string()]);
12519 hybrid_ns.table_exists(exists_req).await.unwrap();
12520
12521 let count = listing_count.load(Ordering::SeqCst);
12522 assert_eq!(
12523 count, 1,
12524 "Expected exactly 1 listing call for table_exists \
12525 without migration mode, but got {}",
12526 count
12527 );
12528
12529 listing_count.store(0, Ordering::SeqCst);
12531
12532 let mut describe_req = DescribeTableRequest::new();
12533 describe_req.id = Some(vec!["test_table".to_string()]);
12534 hybrid_ns.describe_table(describe_req).await.unwrap();
12535
12536 let count = listing_count.load(Ordering::SeqCst);
12537 assert_eq!(
12538 count, 1,
12539 "Expected exactly 1 listing call for describe_table \
12540 without migration mode, but got {}",
12541 count
12542 );
12543 }
12544
12545 #[tokio::test]
12546 async fn test_build_and_root_reads_do_not_create_manifest() {
12547 let temp_dir = TempStdDir::default();
12548 let temp_path = temp_dir.to_str().unwrap();
12549 let manifest_path = std::path::Path::new(temp_path).join("__manifest");
12550
12551 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
12552 .manifest_enabled(false)
12553 .dir_listing_enabled(true)
12554 .build()
12555 .await
12556 .unwrap();
12557 create_scalar_table(&dir_only_ns, "catalog").await;
12558 assert!(!manifest_path.exists());
12559
12560 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12561 .manifest_enabled(true)
12562 .dir_listing_enabled(true)
12563 .build()
12564 .await
12565 .unwrap();
12566 assert!(!manifest_path.exists());
12567
12568 let mut exists_req = TableExistsRequest::new();
12569 exists_req.id = Some(vec!["catalog".to_string()]);
12570 namespace.table_exists(exists_req).await.unwrap();
12571 assert!(!manifest_path.exists());
12572
12573 let mut describe_req = DescribeTableRequest::new();
12574 describe_req.id = Some(vec!["catalog".to_string()]);
12575 namespace.describe_table(describe_req).await.unwrap();
12576 assert!(!manifest_path.exists());
12577
12578 let list_response = namespace
12579 .list_tables(ListTablesRequest {
12580 id: Some(vec![]),
12581 ..Default::default()
12582 })
12583 .await
12584 .unwrap();
12585 assert_eq!(list_response.tables, vec!["catalog".to_string()]);
12586 assert!(!manifest_path.exists());
12587
12588 let mut list_namespaces_req = ListNamespacesRequest::new();
12589 list_namespaces_req.id = Some(vec!["workspace".to_string()]);
12590 let err = namespace
12591 .list_namespaces(list_namespaces_req)
12592 .await
12593 .unwrap_err();
12594 assert!(err.to_string().contains("__manifest"));
12595 assert!(!manifest_path.exists());
12596
12597 let err = namespace
12598 .list_tables(ListTablesRequest {
12599 id: Some(vec!["workspace".to_string()]),
12600 ..Default::default()
12601 })
12602 .await
12603 .unwrap_err();
12604 assert!(err.to_string().contains("__manifest"));
12605 assert!(!manifest_path.exists());
12606
12607 let mut child_describe_req = DescribeTableRequest::new();
12608 child_describe_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
12609 let err = namespace
12610 .describe_table(child_describe_req)
12611 .await
12612 .unwrap_err();
12613 assert!(err.to_string().contains("__manifest"));
12614 assert!(!manifest_path.exists());
12615
12616 let mut child_exists_req = TableExistsRequest::new();
12617 child_exists_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
12618 let err = namespace.table_exists(child_exists_req).await.unwrap_err();
12619 assert!(err.to_string().contains("__manifest"));
12620 assert!(!manifest_path.exists());
12621
12622 let mut create_ns_req = CreateNamespaceRequest::new();
12623 create_ns_req.id = Some(vec!["workspace".to_string()]);
12624 namespace.create_namespace(create_ns_req).await.unwrap();
12625 assert!(manifest_path.exists());
12626 }
12627
12628 #[tokio::test]
12629 async fn test_migrate_updates_read_opened_legacy_manifest() {
12630 let temp_dir = TempStdDir::default();
12631 let temp_path = temp_dir.to_str().unwrap();
12632 create_legacy_manifest_without_primary_key_metadata(temp_path).await;
12633 assert!(!manifest_has_primary_key_metadata(temp_path).await);
12634
12635 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12636 .manifest_enabled(true)
12637 .dir_listing_enabled(true)
12638 .build()
12639 .await
12640 .unwrap();
12641 assert!(!manifest_has_primary_key_metadata(temp_path).await);
12642
12643 let migrated = namespace.migrate().await.unwrap();
12644 assert_eq!(migrated, 0);
12645 assert!(manifest_has_primary_key_metadata(temp_path).await);
12646 }
12647
12648 #[tokio::test]
12649 async fn test_describe_declared_table_checks_versions_only_when_requested() {
12650 let temp_dir = TempStdDir::default();
12651 let temp_path = temp_dir.to_str().unwrap();
12652 let root_uri = file_object_store_uri(temp_path);
12653 let listing_count = Arc::new(AtomicUsize::new(0));
12654 let session = build_listing_counting_session(listing_count.clone());
12655
12656 let namespace = DirectoryNamespaceBuilder::new(root_uri)
12657 .session(session)
12658 .manifest_enabled(false)
12659 .dir_listing_enabled(true)
12660 .build()
12661 .await
12662 .unwrap();
12663
12664 let mut declare_req = DeclareTableRequest::new();
12665 declare_req.id = Some(vec!["test_table".to_string()]);
12666 namespace.declare_table(declare_req).await.unwrap();
12667
12668 listing_count.store(0, Ordering::SeqCst);
12669
12670 let mut describe_req = DescribeTableRequest::new();
12671 describe_req.id = Some(vec!["test_table".to_string()]);
12672 let describe_response = namespace.describe_table(describe_req).await.unwrap();
12673
12674 assert_eq!(describe_response.is_only_declared, None);
12675 assert_eq!(
12676 listing_count.load(Ordering::SeqCst),
12677 1,
12678 "Default describe_table should only list the table directory"
12679 );
12680
12681 listing_count.store(0, Ordering::SeqCst);
12682
12683 let mut describe_req = DescribeTableRequest::new();
12684 describe_req.id = Some(vec!["test_table".to_string()]);
12685 describe_req.check_declared = Some(true);
12686 let describe_response = namespace.describe_table(describe_req).await.unwrap();
12687
12688 assert_eq!(describe_response.is_only_declared, Some(true));
12689 assert_eq!(
12690 listing_count.load(Ordering::SeqCst),
12691 2,
12692 "check_declared describe_table should list the table directory and _versions"
12693 );
12694 }
12695
12696 #[tokio::test]
12697 async fn test_dir_listing_extra_calls_with_migration() {
12698 let temp_dir = TempStdDir::default();
12699 let temp_path = temp_dir.to_str().unwrap();
12700 let root_uri = file_object_store_uri(temp_path);
12701 let listing_count = Arc::new(AtomicUsize::new(0));
12702 let session = build_listing_counting_session(listing_count.clone());
12703
12704 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
12706 .session(session.clone())
12707 .manifest_enabled(false)
12708 .dir_listing_enabled(true)
12709 .build()
12710 .await
12711 .unwrap();
12712
12713 let schema = create_test_schema();
12714 let ipc_data = create_test_ipc_data(&schema);
12715 let mut create_req = CreateTableRequest::new();
12716 create_req.id = Some(vec!["test_table".to_string()]);
12717 dir_only_ns
12718 .create_table(create_req, Bytes::from(ipc_data))
12719 .await
12720 .unwrap();
12721
12722 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
12723 .session(session)
12724 .manifest_enabled(true)
12725 .dir_listing_enabled(true)
12726 .dir_listing_to_manifest_migration_enabled(true)
12727 .build()
12728 .await
12729 .unwrap();
12730
12731 listing_count.store(0, Ordering::SeqCst);
12735
12736 let mut exists_req = TableExistsRequest::new();
12737 exists_req.id = Some(vec!["test_table".to_string()]);
12738 hybrid_ns.table_exists(exists_req).await.unwrap();
12739
12740 let count = listing_count.load(Ordering::SeqCst);
12741 assert_eq!(
12742 count, 1,
12743 "Expected exactly 1 listing call for table_exists with migration mode \
12744 (table directory fallback; manifest reload uses the version hint), but got {}",
12745 count
12746 );
12747
12748 listing_count.store(0, Ordering::SeqCst);
12750
12751 let mut describe_req = DescribeTableRequest::new();
12752 describe_req.id = Some(vec!["test_table".to_string()]);
12753 hybrid_ns.describe_table(describe_req).await.unwrap();
12754
12755 let count = listing_count.load(Ordering::SeqCst);
12756 assert_eq!(
12757 count, 1,
12758 "Expected exactly 1 listing call for describe_table with migration mode \
12759 (table directory fallback; manifest reload uses the version hint), but got {}",
12760 count
12761 );
12762 }
12763
12764 #[tokio::test]
12765 async fn test_manifest_reload_observes_new_version_from_other_namespace() {
12766 let temp_dir = TempStdDir::default();
12767 let temp_path = temp_dir.to_str().unwrap();
12768
12769 let namespace_a = DirectoryNamespaceBuilder::new(temp_path)
12770 .manifest_enabled(true)
12771 .dir_listing_enabled(false)
12772 .build()
12773 .await
12774 .unwrap();
12775 create_scalar_table(&namespace_a, "alpha").await;
12776
12777 let namespace_b = DirectoryNamespaceBuilder::new(temp_path)
12778 .manifest_enabled(true)
12779 .dir_listing_enabled(false)
12780 .build()
12781 .await
12782 .unwrap();
12783 create_scalar_table(&namespace_b, "beta").await;
12784
12785 let response = namespace_a
12786 .list_tables(ListTablesRequest {
12787 id: Some(vec![]),
12788 ..Default::default()
12789 })
12790 .await
12791 .unwrap();
12792
12793 let mut tables = response.tables;
12794 tables.sort();
12795 assert_eq!(tables, vec!["alpha", "beta"]);
12796 }
12797
12798 #[tokio::test]
12799 async fn test_migration_not_found_errors_include_table_id() {
12800 let temp_dir = TempStdDir::default();
12801 let temp_path = temp_dir.to_str().unwrap();
12802
12803 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12804 .manifest_enabled(true)
12805 .dir_listing_enabled(true)
12806 .dir_listing_to_manifest_migration_enabled(true)
12807 .build()
12808 .await
12809 .unwrap();
12810
12811 let mut exists_req = TableExistsRequest::new();
12812 exists_req.id = Some(vec!["missing_table".to_string()]);
12813 let err = namespace.table_exists(exists_req).await.unwrap_err();
12814 assert!(matches!(err, Error::Namespace { .. }));
12815 let err_msg = err.to_string();
12816 assert!(err_msg.contains("Table not found"));
12817 assert!(err_msg.contains("table id 'missing_table'"));
12818
12819 let mut describe_req = DescribeTableRequest::new();
12820 describe_req.id = Some(vec!["missing_table".to_string()]);
12821 let err = namespace.describe_table(describe_req).await.unwrap_err();
12822 assert!(matches!(err, Error::Namespace { .. }));
12823 let err_msg = err.to_string();
12824 assert!(err_msg.contains("Table not found"));
12825 assert!(err_msg.contains("table id 'missing_table'"));
12826 }
12827
12828 #[tokio::test]
12829 async fn test_manifest_not_found_errors_include_full_table_id() {
12830 use lance_namespace::models::CreateNamespaceRequest;
12831
12832 let temp_dir = TempStdDir::default();
12833 let temp_path = temp_dir.to_str().unwrap();
12834
12835 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12836 .manifest_enabled(true)
12837 .dir_listing_enabled(true)
12838 .build()
12839 .await
12840 .unwrap();
12841
12842 let mut create_ns_req = CreateNamespaceRequest::new();
12843 create_ns_req.id = Some(vec!["workspace".to_string()]);
12844 namespace.create_namespace(create_ns_req).await.unwrap();
12845
12846 let missing_table_id = vec!["workspace".to_string(), "missing_table".to_string()];
12847
12848 let mut exists_req = TableExistsRequest::new();
12849 exists_req.id = Some(missing_table_id.clone());
12850 let err = namespace.table_exists(exists_req).await.unwrap_err();
12851 assert!(matches!(err, Error::Namespace { .. }));
12852 let err_msg = err.to_string();
12853 assert!(err_msg.contains("Table not found"));
12854 assert!(err_msg.contains("table id 'workspace$missing_table'"));
12855
12856 let mut describe_req = DescribeTableRequest::new();
12857 describe_req.id = Some(missing_table_id);
12858 let err = namespace.describe_table(describe_req).await.unwrap_err();
12859 assert!(matches!(err, Error::Namespace { .. }));
12860 let err_msg = err.to_string();
12861 assert!(err_msg.contains("Table not found"));
12862 assert!(err_msg.contains("table id 'workspace$missing_table'"));
12863 }
12864
12865 async fn create_tagged_test_table(
12868 versions: u32,
12869 ) -> (Arc<DirectoryNamespace>, TempStdDir, Vec<String>) {
12870 use arrow::array::{Int32Array, RecordBatchIterator};
12871 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12872 use arrow::record_batch::RecordBatch;
12873 use lance::dataset::{Dataset, WriteMode, WriteParams};
12874
12875 assert!(versions >= 1, "versions must be at least 1");
12876
12877 let temp_dir = TempStdDir::default();
12878 let temp_path = temp_dir.to_str().unwrap();
12879
12880 let namespace = Arc::new(
12881 DirectoryNamespaceBuilder::new(temp_path)
12882 .build()
12883 .await
12884 .unwrap(),
12885 );
12886 let table_id = vec!["tag_table".to_string()];
12887 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12888 "id",
12889 DataType::Int32,
12890 false,
12891 )]));
12892 let initial_batch = RecordBatch::try_new(
12893 arrow_schema.clone(),
12894 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12895 )
12896 .unwrap();
12897 let batches = RecordBatchIterator::new(vec![Ok(initial_batch)], arrow_schema.clone());
12898 let write_params = WriteParams {
12899 mode: WriteMode::Create,
12900 ..Default::default()
12901 };
12902
12903 let mut dataset = Dataset::write_into_namespace(
12904 batches,
12905 namespace.clone() as Arc<dyn LanceNamespace>,
12906 table_id.clone(),
12907 Some(write_params),
12908 )
12909 .await
12910 .unwrap();
12911
12912 for i in 1..versions {
12913 let value_start = (i as i32) * 10;
12914 let batch = RecordBatch::try_new(
12915 arrow_schema.clone(),
12916 vec![Arc::new(Int32Array::from(vec![
12917 value_start,
12918 value_start + 1,
12919 ]))],
12920 )
12921 .unwrap();
12922 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12923 dataset.append(batches, None).await.unwrap();
12924 }
12925
12926 (namespace, temp_dir, table_id)
12927 }
12928
12929 fn namespace_code(err: &Error) -> Option<ErrorCode> {
12931 match err {
12932 Error::Namespace { source, .. } => {
12933 source.downcast_ref::<NamespaceError>().map(|e| e.code())
12934 }
12935 _ => None,
12936 }
12937 }
12938
12939 #[tokio::test]
12940 async fn test_create_and_list_branches() {
12941 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
12942
12943 namespace
12944 .create_table_branch(CreateTableBranchRequest {
12945 id: Some(table_id.clone()),
12946 name: "dev".to_string(),
12947 ..Default::default()
12948 })
12949 .await
12950 .unwrap();
12951 namespace
12952 .create_table_branch(CreateTableBranchRequest {
12953 id: Some(table_id.clone()),
12954 name: "staging".to_string(),
12955 ..Default::default()
12956 })
12957 .await
12958 .unwrap();
12959
12960 let resp = namespace
12961 .list_table_branches(ListTableBranchesRequest {
12962 id: Some(table_id.clone()),
12963 ..Default::default()
12964 })
12965 .await
12966 .unwrap();
12967 assert_eq!(
12968 resp.branches.len(),
12969 2,
12970 "expected 2 branches, got: {:?}",
12971 resp.branches
12972 );
12973 assert!(resp.branches.contains_key("dev"));
12974 assert!(resp.branches.contains_key("staging"));
12975 assert!(resp.page_token.is_none());
12976
12977 namespace
12979 .delete_table_branch(DeleteTableBranchRequest {
12980 id: Some(table_id.clone()),
12981 name: "dev".to_string(),
12982 ..Default::default()
12983 })
12984 .await
12985 .unwrap();
12986
12987 let resp = namespace
12988 .list_table_branches(ListTableBranchesRequest {
12989 id: Some(table_id),
12990 ..Default::default()
12991 })
12992 .await
12993 .unwrap();
12994 assert_eq!(resp.branches.len(), 1, "expected 1 branch after delete");
12995 assert!(!resp.branches.contains_key("dev"));
12996 assert!(resp.branches.contains_key("staging"));
12997 }
12998
12999 #[tokio::test]
13000 async fn test_create_branch_from_version() {
13001 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
13002
13003 namespace
13005 .create_table_branch(CreateTableBranchRequest {
13006 id: Some(table_id.clone()),
13007 name: "from-v1".to_string(),
13008 from_version: Some(1),
13009 ..Default::default()
13010 })
13011 .await
13012 .unwrap();
13013
13014 let resp = namespace
13015 .list_table_branches(ListTableBranchesRequest {
13016 id: Some(table_id),
13017 ..Default::default()
13018 })
13019 .await
13020 .unwrap();
13021 let branch = resp
13022 .branches
13023 .get("from-v1")
13024 .expect("forked branch should be listed");
13025 assert_eq!(
13026 branch.parent_version, 1,
13027 "branch should fork from version 1"
13028 );
13029 assert!(
13030 branch.parent_branch.is_none(),
13031 "a branch forked from main has no parent branch"
13032 );
13033 }
13034
13035 #[tokio::test]
13040 async fn test_create_branch_from_other_branch() {
13041 use lance::dataset::builder::DatasetBuilder;
13042
13043 let (namespace, _temp_dir) = create_test_namespace().await;
13044 create_scalar_table(&namespace, "users").await; create_branch_with_commits(&namespace, "users", "dev", 1).await;
13047 let main_ds = open_dataset(&namespace, "users").await;
13049 append_scalar_version(main_ds.uri(), 500).await; namespace
13052 .create_table_branch(CreateTableBranchRequest {
13053 id: Some(vec!["users".to_string()]),
13054 name: "child".to_string(),
13055 from_branch: Some("dev".to_string()),
13056 from_version: Some(2),
13057 ..Default::default()
13058 })
13059 .await
13060 .unwrap();
13061
13062 let child_ds = DatasetBuilder::from_uri(main_ds.uri())
13063 .with_branch("child", None)
13064 .load()
13065 .await
13066 .unwrap();
13067 let ids = scan_id_column(&child_ds).await;
13068 assert!(
13069 ids.contains(&100) && ids.contains(&101),
13070 "child must contain dev's appended rows, got: {:?}",
13071 ids
13072 );
13073 assert!(
13074 !ids.contains(&500),
13075 "child must not contain main's diverged rows, got: {:?}",
13076 ids
13077 );
13078
13079 let listed = namespace
13081 .list_table_branches(ListTableBranchesRequest {
13082 id: Some(vec!["users".to_string()]),
13083 ..Default::default()
13084 })
13085 .await
13086 .unwrap();
13087 assert_eq!(
13088 listed
13089 .branches
13090 .get("child")
13091 .unwrap()
13092 .parent_branch
13093 .as_deref(),
13094 Some("dev")
13095 );
13096 }
13097
13098 #[tokio::test]
13099 async fn test_create_existing_branch_conflict() {
13100 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13101
13102 namespace
13103 .create_table_branch(CreateTableBranchRequest {
13104 id: Some(table_id.clone()),
13105 name: "dev".to_string(),
13106 ..Default::default()
13107 })
13108 .await
13109 .unwrap();
13110
13111 let err = namespace
13112 .create_table_branch(CreateTableBranchRequest {
13113 id: Some(table_id),
13114 name: "dev".to_string(),
13115 ..Default::default()
13116 })
13117 .await
13118 .unwrap_err();
13119 assert_eq!(
13120 namespace_code(&err),
13121 Some(ErrorCode::TableBranchAlreadyExists),
13122 "expected TableBranchAlreadyExists, got: {}",
13123 err
13124 );
13125 assert!(
13126 err.to_string().to_lowercase().contains("already exists"),
13127 "expected already-exists message, got: {}",
13128 err
13129 );
13130 }
13131
13132 #[tokio::test]
13133 async fn test_delete_unknown_branch() {
13134 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13135
13136 let err = namespace
13137 .delete_table_branch(DeleteTableBranchRequest {
13138 id: Some(table_id),
13139 name: "does-not-exist".to_string(),
13140 ..Default::default()
13141 })
13142 .await
13143 .unwrap_err();
13144 assert_eq!(
13145 namespace_code(&err),
13146 Some(ErrorCode::TableBranchNotFound),
13147 "expected TableBranchNotFound, got: {}",
13148 err
13149 );
13150 assert!(
13151 err.to_string().to_lowercase().contains("not found"),
13152 "expected not-found message, got: {}",
13153 err
13154 );
13155 }
13156
13157 #[tokio::test]
13158 async fn test_delete_referenced_branch_conflict() {
13159 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13160
13161 namespace
13163 .create_table_branch(CreateTableBranchRequest {
13164 id: Some(table_id.clone()),
13165 name: "parent".to_string(),
13166 ..Default::default()
13167 })
13168 .await
13169 .unwrap();
13170 namespace
13171 .create_table_branch(CreateTableBranchRequest {
13172 id: Some(table_id.clone()),
13173 name: "child".to_string(),
13174 from_branch: Some("parent".to_string()),
13175 ..Default::default()
13176 })
13177 .await
13178 .unwrap();
13179
13180 let listed = namespace
13182 .list_table_branches(ListTableBranchesRequest {
13183 id: Some(table_id.clone()),
13184 ..Default::default()
13185 })
13186 .await
13187 .unwrap();
13188 let child = listed
13189 .branches
13190 .get("child")
13191 .expect("child branch should be listed");
13192 assert_eq!(
13193 child.parent_branch.as_deref(),
13194 Some("parent"),
13195 "child should record parent branch as its fork point"
13196 );
13197 assert!(
13198 child.parent_version >= 1,
13199 "child should record the parent version it forked from, got {}",
13200 child.parent_version
13201 );
13202
13203 let err = namespace
13206 .delete_table_branch(DeleteTableBranchRequest {
13207 id: Some(table_id),
13208 name: "parent".to_string(),
13209 ..Default::default()
13210 })
13211 .await
13212 .unwrap_err();
13213 assert_eq!(
13214 namespace_code(&err),
13215 Some(ErrorCode::InvalidInput),
13216 "expected InvalidInput for deleting a referenced branch, got: {}",
13217 err
13218 );
13219 assert!(
13220 err.to_string().to_lowercase().contains("referenced"),
13221 "error should explain the branch is still referenced, got: {}",
13222 err
13223 );
13224 }
13225
13226 #[tokio::test]
13227 async fn test_branch_name_required() {
13228 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13229
13230 let create_err = namespace
13231 .create_table_branch(CreateTableBranchRequest {
13232 id: Some(table_id.clone()),
13233 name: String::new(),
13234 ..Default::default()
13235 })
13236 .await
13237 .unwrap_err();
13238 assert_eq!(
13239 namespace_code(&create_err),
13240 Some(ErrorCode::InvalidInput),
13241 "empty name on create should be InvalidInput, got: {}",
13242 create_err
13243 );
13244 assert!(
13245 create_err
13246 .to_string()
13247 .to_lowercase()
13248 .contains("must not be empty")
13249 );
13250
13251 let delete_err = namespace
13252 .delete_table_branch(DeleteTableBranchRequest {
13253 id: Some(table_id),
13254 name: String::new(),
13255 ..Default::default()
13256 })
13257 .await
13258 .unwrap_err();
13259 assert_eq!(
13260 namespace_code(&delete_err),
13261 Some(ErrorCode::InvalidInput),
13262 "empty name on delete should be InvalidInput, got: {}",
13263 delete_err
13264 );
13265 assert!(
13266 delete_err
13267 .to_string()
13268 .to_lowercase()
13269 .contains("must not be empty")
13270 );
13271 }
13272
13273 #[tokio::test]
13274 async fn test_create_branch_rejects_negative_from_version() {
13275 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13276
13277 let err = namespace
13278 .create_table_branch(CreateTableBranchRequest {
13279 id: Some(table_id),
13280 name: "dev".to_string(),
13281 from_version: Some(-1),
13282 ..Default::default()
13283 })
13284 .await
13285 .unwrap_err();
13286 assert_eq!(
13287 namespace_code(&err),
13288 Some(ErrorCode::InvalidInput),
13289 "negative from_version should be InvalidInput, got: {}",
13290 err
13291 );
13292 assert!(err.to_string().to_lowercase().contains("from_version"));
13293 }
13294
13295 #[tokio::test]
13296 async fn test_create_branch_nonexistent_from_version() {
13297 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13298
13299 let err = namespace
13302 .create_table_branch(CreateTableBranchRequest {
13303 id: Some(table_id),
13304 name: "dev".to_string(),
13305 from_version: Some(999),
13306 ..Default::default()
13307 })
13308 .await
13309 .unwrap_err();
13310 assert_eq!(
13311 namespace_code(&err),
13312 Some(ErrorCode::InvalidInput),
13313 "non-existent from_version should map to InvalidInput, got: {}",
13314 err
13315 );
13316 assert!(
13317 err.to_string().to_lowercase().contains("does not exist"),
13318 "error should name the missing source, got: {}",
13319 err
13320 );
13321 }
13322
13323 #[tokio::test]
13324 async fn test_create_and_list_tags() {
13325 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
13326
13327 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
13328 req.id = Some(table_id.clone());
13329 namespace.create_table_tag(req).await.unwrap();
13330
13331 let mut req = CreateTableTagRequest::new("v2".to_string(), 2);
13332 req.id = Some(table_id.clone());
13333 namespace.create_table_tag(req).await.unwrap();
13334
13335 let mut list_req = ListTableTagsRequest::new();
13336 list_req.id = Some(table_id);
13337 let resp = namespace.list_table_tags(list_req).await.unwrap();
13338
13339 assert_eq!(resp.tags.len(), 2, "expected 2 tags, got: {:?}", resp.tags);
13340 assert_eq!(resp.tags.get("v1").unwrap().version, 1);
13341 assert_eq!(resp.tags.get("v2").unwrap().version, 2);
13342 assert!(resp.page_token.is_none());
13343 }
13344
13345 #[tokio::test]
13346 async fn test_create_existing_tag_conflict() {
13347 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13348
13349 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
13350 req.id = Some(table_id.clone());
13351 namespace.create_table_tag(req).await.unwrap();
13352
13353 let mut req = CreateTableTagRequest::new("v1".to_string(), 2);
13354 req.id = Some(table_id);
13355 let err = namespace.create_table_tag(req).await.unwrap_err();
13356 assert!(
13357 err.to_string().to_lowercase().contains("already exists"),
13358 "expected already-exists error, got: {}",
13359 err
13360 );
13361 }
13362
13363 #[tokio::test]
13364 async fn test_get_tag_version() {
13365 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
13366
13367 let mut req = CreateTableTagRequest::new("release".to_string(), 2);
13368 req.id = Some(table_id.clone());
13369 namespace.create_table_tag(req).await.unwrap();
13370
13371 let mut get_req = GetTableTagVersionRequest::new("release".to_string());
13372 get_req.id = Some(table_id);
13373 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
13374 assert_eq!(resp.version, 2);
13375 assert_eq!(resp.branch, None);
13376 }
13377
13378 #[tokio::test]
13379 async fn test_get_unknown_tag() {
13380 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13381
13382 let mut get_req = GetTableTagVersionRequest::new("does-not-exist".to_string());
13383 get_req.id = Some(table_id);
13384 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
13385 assert!(
13386 err.to_string().to_lowercase().contains("not found"),
13387 "expected not-found error, got: {}",
13388 err
13389 );
13390 }
13391
13392 #[tokio::test]
13393 async fn test_update_tag_to_new_version() {
13394 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
13395
13396 let mut req = CreateTableTagRequest::new("rolling".to_string(), 1);
13397 req.id = Some(table_id.clone());
13398 namespace.create_table_tag(req).await.unwrap();
13399
13400 let mut update_req = UpdateTableTagRequest::new("rolling".to_string(), 3);
13401 update_req.id = Some(table_id.clone());
13402 namespace.update_table_tag(update_req).await.unwrap();
13403
13404 let mut get_req = GetTableTagVersionRequest::new("rolling".to_string());
13405 get_req.id = Some(table_id);
13406 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
13407 assert_eq!(resp.version, 3);
13408 }
13409
13410 #[tokio::test]
13411 async fn test_update_unknown_tag() {
13412 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13413
13414 let mut update_req = UpdateTableTagRequest::new("ghost".to_string(), 1);
13415 update_req.id = Some(table_id);
13416 let err = namespace.update_table_tag(update_req).await.unwrap_err();
13417 assert!(
13418 err.to_string().to_lowercase().contains("not found"),
13419 "expected not-found error, got: {}",
13420 err
13421 );
13422 }
13423
13424 #[tokio::test]
13425 async fn test_delete_tag() {
13426 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13427
13428 let mut req = CreateTableTagRequest::new("doomed".to_string(), 1);
13429 req.id = Some(table_id.clone());
13430 namespace.create_table_tag(req).await.unwrap();
13431
13432 let mut delete_req = DeleteTableTagRequest::new("doomed".to_string());
13433 delete_req.id = Some(table_id.clone());
13434 namespace.delete_table_tag(delete_req).await.unwrap();
13435
13436 let mut list_req = ListTableTagsRequest::new();
13437 list_req.id = Some(table_id.clone());
13438 let resp = namespace.list_table_tags(list_req).await.unwrap();
13439 assert!(resp.tags.is_empty(), "tag should be removed after delete");
13440
13441 let mut get_req = GetTableTagVersionRequest::new("doomed".to_string());
13443 get_req.id = Some(table_id);
13444 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
13445 assert!(err.to_string().to_lowercase().contains("not found"));
13446 }
13447
13448 #[tokio::test]
13449 async fn test_delete_unknown_tag() {
13450 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13451
13452 let mut delete_req = DeleteTableTagRequest::new("nope".to_string());
13453 delete_req.id = Some(table_id);
13454 let err = namespace.delete_table_tag(delete_req).await.unwrap_err();
13455 assert!(
13456 err.to_string().to_lowercase().contains("not found"),
13457 "expected not-found error, got: {}",
13458 err
13459 );
13460 }
13461
13462 #[tokio::test]
13463 async fn test_create_tag_invalid_version() {
13464 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
13465
13466 let mut req = CreateTableTagRequest::new("v0".to_string(), 0);
13468 req.id = Some(table_id.clone());
13469 let err = namespace.create_table_tag(req).await.unwrap_err();
13470 assert!(
13471 err.to_string().to_lowercase().contains("positive"),
13472 "expected positive-version error, got: {}",
13473 err
13474 );
13475
13476 let mut req = CreateTableTagRequest::new(String::new(), 1);
13478 req.id = Some(table_id);
13479 let err = namespace.create_table_tag(req).await.unwrap_err();
13480 assert!(
13481 err.to_string().to_lowercase().contains("must not be empty"),
13482 "expected empty-tag-name error, got: {}",
13483 err
13484 );
13485 }
13486
13487 #[tokio::test]
13488 async fn test_create_tag_table_not_found() {
13489 let (namespace, _temp_dir) = create_test_namespace().await;
13490
13491 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
13492 req.id = Some(vec!["does_not_exist".to_string()]);
13493 let err = namespace.create_table_tag(req).await.unwrap_err();
13494 let msg = err.to_string();
13495 assert!(
13496 msg.contains("Table") && msg.to_lowercase().contains("not found"),
13497 "expected TableNotFound error, got: {}",
13498 err
13499 );
13500 }
13501 #[tokio::test]
13502 async fn test_alter_table_drop_columns_missing_id() {
13503 use lance_namespace::models::AlterTableDropColumnsRequest;
13504
13505 let (namespace, _temp_dir) = create_test_namespace().await;
13506
13507 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
13508 let result = namespace.alter_table_drop_columns(request).await;
13509 assert!(result.is_err(), "Should fail when table ID is missing");
13510 }
13511
13512 #[tokio::test]
13513 async fn test_alter_table_drop_columns_nonexistent_table() {
13514 use lance_namespace::models::AlterTableDropColumnsRequest;
13515
13516 let (namespace, _temp_dir) = create_test_namespace().await;
13517
13518 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
13519 request.id = Some(vec!["nonexistent".to_string()]);
13520 let result = namespace.alter_table_drop_columns(request).await;
13521 assert!(result.is_err(), "Should fail when table does not exist");
13522 }
13523
13524 #[tokio::test]
13525 async fn test_create_branch_on_managed_dataset_succeeds() {
13526 use lance::dataset::builder::DatasetBuilder;
13527
13528 let temp = TempStdDir::default();
13529 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
13530 let table_id = vec!["t".to_string()];
13531 let mut main = create_managed_table(&ns, &table_id).await;
13532
13533 let fork_version = main.version().version;
13534 let branch = main
13535 .create_branch("exp", fork_version, None)
13536 .await
13537 .expect("create_branch failed");
13538 assert_eq!(branch.manifest.branch.as_deref(), Some("exp"));
13539 assert_eq!(scan_id_column(&branch).await, vec![1, 2]);
13540
13541 let reopened = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
13542 .await
13543 .unwrap()
13544 .with_branch("exp", None)
13545 .load()
13546 .await
13547 .expect("reopen branch failed");
13548 assert_eq!(scan_id_column(&reopened).await, vec![1, 2]);
13549 }
13550
13551 #[tokio::test]
13552 async fn test_alter_transaction_set_status() {
13553 use lance_namespace::models::{
13554 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
13555 DescribeTransactionRequest,
13556 };
13557
13558 let (namespace, _temp_dir) = create_test_namespace().await;
13559 create_scalar_table(&namespace, "users").await;
13560 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
13561 .await
13562 .expect("create_scalar_index should return a transaction id");
13563
13564 let describe_resp = namespace
13566 .describe_transaction(DescribeTransactionRequest {
13567 id: Some(vec!["users".to_string(), txn_id.clone()]),
13568 ..Default::default()
13569 })
13570 .await
13571 .unwrap();
13572 assert_eq!(describe_resp.status, "SUCCEEDED");
13573
13574 let response = namespace
13576 .alter_transaction(AlterTransactionRequest {
13577 id: Some(vec!["users".to_string(), txn_id.clone()]),
13578 actions: vec![AlterTransactionAction {
13579 set_status_action: Some(Box::new(AlterTransactionSetStatus {
13580 status: Some("Canceled".to_string()),
13581 })),
13582 set_property_action: None,
13583 unset_property_action: None,
13584 }],
13585 ..Default::default()
13586 })
13587 .await
13588 .unwrap();
13589 assert_eq!(response.status, "Canceled");
13590 assert!(response.properties.is_some());
13591 let props = response.properties.unwrap();
13592 assert_eq!(props.get("uuid"), Some(&txn_id));
13593 assert_eq!(props.get("operation"), Some(&"CreateIndex".to_string()));
13594 }
13595
13596 #[tokio::test]
13597 async fn test_alter_transaction_set_property() {
13598 use lance_namespace::models::{
13599 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
13600 };
13601
13602 let (namespace, _temp_dir) = create_test_namespace().await;
13603 create_scalar_table(&namespace, "users").await;
13604 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
13605 .await
13606 .expect("create_scalar_index should return a transaction id");
13607
13608 let response = namespace
13609 .alter_transaction(AlterTransactionRequest {
13610 id: Some(vec!["users".to_string(), txn_id.clone()]),
13611 actions: vec![AlterTransactionAction {
13612 set_status_action: None,
13613 set_property_action: Some(Box::new(AlterTransactionSetProperty {
13614 key: Some("custom_key".to_string()),
13615 value: Some("custom_value".to_string()),
13616 mode: None,
13617 })),
13618 unset_property_action: None,
13619 }],
13620 ..Default::default()
13621 })
13622 .await
13623 .unwrap();
13624 assert_eq!(response.status, "SUCCEEDED");
13625 let props = response.properties.unwrap();
13626 assert_eq!(props.get("custom_key"), Some(&"custom_value".to_string()));
13627 }
13628
13629 #[tokio::test]
13630 async fn test_alter_transaction_set_property_fail_mode() {
13631 use lance_namespace::models::{
13632 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
13633 };
13634
13635 let (namespace, _temp_dir) = create_test_namespace().await;
13636 create_scalar_table(&namespace, "users").await;
13637 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
13638 .await
13639 .expect("create_scalar_index should return a transaction id");
13640
13641 namespace
13643 .alter_transaction(AlterTransactionRequest {
13644 id: Some(vec!["users".to_string(), txn_id.clone()]),
13645 actions: vec![AlterTransactionAction {
13646 set_status_action: None,
13647 set_property_action: Some(Box::new(AlterTransactionSetProperty {
13648 key: Some("custom_key".to_string()),
13649 value: Some("initial_value".to_string()),
13650 mode: None,
13651 })),
13652 unset_property_action: None,
13653 }],
13654 ..Default::default()
13655 })
13656 .await
13657 .unwrap();
13658
13659 let result = namespace
13662 .alter_transaction(AlterTransactionRequest {
13663 id: Some(vec!["users".to_string(), txn_id.clone()]),
13664 actions: vec![AlterTransactionAction {
13665 set_status_action: None,
13666 set_property_action: Some(Box::new(AlterTransactionSetProperty {
13667 key: Some("custom_key".to_string()),
13668 value: Some("new_value".to_string()),
13669 mode: Some("Fail".to_string()),
13670 })),
13671 unset_property_action: None,
13672 }],
13673 ..Default::default()
13674 })
13675 .await;
13676 assert!(result.is_err());
13677 }
13678
13679 #[tokio::test]
13680 async fn test_alter_transaction_unset_property() {
13681 use lance_namespace::models::{
13682 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
13683 AlterTransactionUnsetProperty,
13684 };
13685
13686 let (namespace, _temp_dir) = create_test_namespace().await;
13687 create_scalar_table(&namespace, "users").await;
13688 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
13689 .await
13690 .expect("create_scalar_index should return a transaction id");
13691
13692 let response = namespace
13694 .alter_transaction(AlterTransactionRequest {
13695 id: Some(vec!["users".to_string(), txn_id.clone()]),
13696 actions: vec![
13697 AlterTransactionAction {
13698 set_status_action: None,
13699 set_property_action: Some(Box::new(AlterTransactionSetProperty {
13700 key: Some("temp_key".to_string()),
13701 value: Some("temp_value".to_string()),
13702 mode: None,
13703 })),
13704 unset_property_action: None,
13705 },
13706 AlterTransactionAction {
13707 set_status_action: None,
13708 set_property_action: None,
13709 unset_property_action: Some(Box::new(AlterTransactionUnsetProperty {
13710 key: Some("temp_key".to_string()),
13711 mode: None,
13712 })),
13713 },
13714 ],
13715 ..Default::default()
13716 })
13717 .await
13718 .unwrap();
13719 assert_eq!(response.status, "SUCCEEDED");
13720 let props = response.properties.unwrap();
13721 assert!(!props.contains_key("temp_key"));
13722 }
13723
13724 #[tokio::test]
13725 async fn test_alter_transaction_invalid_status() {
13726 use lance_namespace::models::{
13727 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
13728 };
13729
13730 let (namespace, _temp_dir) = create_test_namespace().await;
13731 create_scalar_table(&namespace, "users").await;
13732 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
13733 .await
13734 .expect("create_scalar_index should return a transaction id");
13735
13736 let result = namespace
13737 .alter_transaction(AlterTransactionRequest {
13738 id: Some(vec!["users".to_string(), txn_id.clone()]),
13739 actions: vec![AlterTransactionAction {
13740 set_status_action: Some(Box::new(AlterTransactionSetStatus {
13741 status: Some("InvalidStatus".to_string()),
13742 })),
13743 set_property_action: None,
13744 unset_property_action: None,
13745 }],
13746 ..Default::default()
13747 })
13748 .await;
13749 assert!(result.is_err());
13750 }
13751
13752 #[tokio::test]
13753 async fn test_alter_transaction_not_found() {
13754 use lance_namespace::models::{
13755 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
13756 };
13757
13758 let (namespace, _temp_dir) = create_test_namespace().await;
13759 create_scalar_table(&namespace, "users").await;
13760
13761 let result = namespace
13763 .alter_transaction(AlterTransactionRequest {
13764 id: Some(vec!["users".to_string(), "non_existent_txn".to_string()]),
13765 actions: vec![AlterTransactionAction {
13766 set_status_action: Some(Box::new(AlterTransactionSetStatus {
13767 status: Some("Canceled".to_string()),
13768 })),
13769 set_property_action: None,
13770 unset_property_action: None,
13771 }],
13772 ..Default::default()
13773 })
13774 .await;
13775 assert!(result.is_err());
13776 }
13777
13778 #[tokio::test]
13779 async fn test_alter_transaction_missing_id() {
13780 use lance_namespace::models::{
13781 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
13782 };
13783
13784 let (namespace, _temp_dir) = create_test_namespace().await;
13785
13786 let result = namespace
13788 .alter_transaction(AlterTransactionRequest {
13789 id: None,
13790 actions: vec![AlterTransactionAction {
13791 set_status_action: Some(Box::new(AlterTransactionSetStatus {
13792 status: Some("Canceled".to_string()),
13793 })),
13794 set_property_action: None,
13795 unset_property_action: None,
13796 }],
13797 ..Default::default()
13798 })
13799 .await;
13800 assert!(result.is_err());
13801
13802 let result = namespace
13804 .alter_transaction(AlterTransactionRequest {
13805 id: Some(vec!["users".to_string()]),
13806 actions: vec![AlterTransactionAction {
13807 set_status_action: Some(Box::new(AlterTransactionSetStatus {
13808 status: Some("Canceled".to_string()),
13809 })),
13810 set_property_action: None,
13811 unset_property_action: None,
13812 }],
13813 ..Default::default()
13814 })
13815 .await;
13816 assert!(result.is_err());
13817 }
13818
13819 #[tokio::test]
13820 async fn test_alter_transaction_persists_changes() {
13821 use lance_namespace::models::{
13822 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
13823 AlterTransactionSetStatus, DescribeTransactionRequest,
13824 };
13825
13826 let (namespace, _temp_dir) = create_test_namespace().await;
13827 create_scalar_table(&namespace, "users").await;
13828 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
13829
13830 let txn_id = transaction_id.expect("scalar index should produce a transaction id");
13831
13832 namespace
13834 .alter_transaction(AlterTransactionRequest {
13835 id: Some(vec!["users".to_string(), txn_id.clone()]),
13836 actions: vec![
13837 AlterTransactionAction {
13838 set_status_action: Some(Box::new(AlterTransactionSetStatus {
13839 status: Some("Canceled".to_string()),
13840 })),
13841 set_property_action: None,
13842 unset_property_action: None,
13843 },
13844 AlterTransactionAction {
13845 set_status_action: None,
13846 set_property_action: Some(Box::new(AlterTransactionSetProperty {
13847 key: Some("owner".to_string()),
13848 value: Some("alice".to_string()),
13849 mode: None,
13850 })),
13851 unset_property_action: None,
13852 },
13853 ],
13854 ..Default::default()
13855 })
13856 .await
13857 .unwrap();
13858
13859 let describe_resp = namespace
13862 .describe_transaction(DescribeTransactionRequest {
13863 id: Some(vec!["users".to_string(), txn_id.clone()]),
13864 ..Default::default()
13865 })
13866 .await
13867 .unwrap();
13868 let props = describe_resp.properties.expect("properties should be set");
13869 assert_eq!(props.get("owner"), Some(&"alice".to_string()));
13870 assert!(!props.contains_key("_status"));
13874
13875 let follow_up = namespace
13876 .alter_transaction(AlterTransactionRequest {
13877 id: Some(vec!["users".to_string(), txn_id.clone()]),
13878 actions: vec![],
13879 ..Default::default()
13880 })
13881 .await
13882 .unwrap();
13883 assert_eq!(follow_up.status, "Canceled");
13884 let follow_up_props = follow_up.properties.unwrap();
13885 assert_eq!(follow_up_props.get("owner"), Some(&"alice".to_string()));
13886 }
13887}