1pub mod manifest;
10
11use arrow::array::Float32Array;
12use arrow::record_batch::RecordBatchIterator;
13use arrow_ipc::reader::StreamReader;
14use async_trait::async_trait;
15use bytes::Bytes;
16use futures::{StreamExt, TryStreamExt};
17use lance::dataset::builder::DatasetBuilder;
18use lance::dataset::scanner::Scanner;
19use lance::dataset::statistics::DatasetStatisticsExt;
20use lance::dataset::transaction::{Operation, Transaction};
21use lance::dataset::{
22 Dataset, MergeInsertBuilder, WhenMatched, WhenNotMatched, WhenNotMatchedBySource, WriteMode,
23 WriteParams,
24};
25use lance::index::{DatasetIndexExt, IndexParams, vector::VectorIndexParams};
26use lance::session::Session;
27use lance_index::scalar::{
28 BuiltinIndexType, FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams,
29};
30use lance_index::vector::{
31 bq::RQBuildParams, hnsw::builder::HnswBuildParams, ivf::IvfBuildParams, pq::PQBuildParams,
32 sq::builder::SQBuildParams,
33};
34use lance_index::{IndexType, is_system_index};
35use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
36use lance_linalg::distance::MetricType;
37use lance_table::io::commit::{ManifestNamingScheme, VERSIONS_DIR};
38use object_store::ObjectStoreExt;
39use object_store::path::Path;
40use object_store::{Error as ObjectStoreError, ObjectStore as OSObjectStore, PutMode, PutOptions};
41use std::collections::HashMap;
42use std::io::Cursor;
43use std::sync::{Arc, Mutex};
44
45use crate::context::DynamicContextProvider;
46use lance_namespace::models::{
47 AnalyzeTableQueryPlanRequest, BatchDeleteTableVersionsRequest,
48 BatchDeleteTableVersionsResponse, CountTableRowsRequest, CreateNamespaceRequest,
49 CreateNamespaceResponse, CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest,
50 CreateTableResponse, CreateTableScalarIndexResponse, CreateTableTagRequest,
51 CreateTableTagResponse, CreateTableVersionRequest, CreateTableVersionResponse,
52 DeclareTableRequest, DeclareTableResponse, DeleteTableTagRequest, DeleteTableTagResponse,
53 DescribeNamespaceRequest, DescribeNamespaceResponse, DescribeTableIndexStatsRequest,
54 DescribeTableIndexStatsResponse, DescribeTableRequest, DescribeTableResponse,
55 DescribeTableVersionRequest, DescribeTableVersionResponse, DescribeTransactionRequest,
56 DescribeTransactionResponse, DropNamespaceRequest, DropNamespaceResponse,
57 DropTableIndexRequest, DropTableIndexResponse, DropTableRequest, DropTableResponse,
58 ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary, GetTableStatsRequest,
59 GetTableStatsResponse, GetTableTagVersionRequest, GetTableTagVersionResponse, Identity,
60 IndexContent, InsertIntoTableRequest, InsertIntoTableResponse, ListNamespacesRequest,
61 ListNamespacesResponse, ListTableIndicesRequest, ListTableIndicesResponse,
62 ListTableTagsRequest, ListTableTagsResponse, ListTableVersionsRequest,
63 ListTableVersionsResponse, ListTablesRequest, ListTablesResponse, MergeInsertIntoTableRequest,
64 MergeInsertIntoTableResponse, NamespaceExistsRequest, QueryTableRequest,
65 QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest, RestoreTableResponse,
66 TableExistsRequest, TableVersion, TagContents as ModelTagContents,
67 UpdateTableSchemaMetadataRequest, UpdateTableSchemaMetadataResponse, UpdateTableTagRequest,
68 UpdateTableTagResponse,
69};
70
71use lance_core::{Error, Result};
72use lance_namespace::LanceNamespace;
73use lance_namespace::error::NamespaceError;
74use lance_namespace::schema::arrow_schema_to_json;
75
76use crate::credentials::{
77 CredentialVendor, create_credential_vendor_for_location, has_credential_vendor_config,
78};
79
80#[derive(Debug, Default)]
85pub struct OpsMetrics {
86 counters: Mutex<HashMap<String, u64>>,
87}
88
89impl OpsMetrics {
90 pub fn increment(&self, operation: &str) {
92 if let Ok(mut counters) = self.counters.lock() {
93 *counters.entry(operation.to_string()).or_insert(0) += 1;
94 }
95 }
96
97 pub fn retrieve(&self) -> HashMap<String, u64> {
99 self.counters.lock().map(|c| c.clone()).unwrap_or_default()
100 }
101
102 pub fn reset(&self) {
104 if let Ok(mut counters) = self.counters.lock() {
105 counters.clear();
106 }
107 }
108}
109
110pub(crate) struct TableStatus {
115 pub(crate) exists: bool,
117 pub(crate) is_deregistered: bool,
119 pub(crate) has_reserved_file: bool,
121}
122
123enum DirectoryIndexParams {
124 Scalar {
125 index_type: IndexType,
126 params: ScalarIndexParams,
127 },
128 Inverted(InvertedIndexParams),
129 Vector {
130 index_type: IndexType,
131 params: VectorIndexParams,
132 },
133}
134
135impl DirectoryIndexParams {
136 fn index_type(&self) -> IndexType {
137 match self {
138 Self::Scalar { index_type, .. } | Self::Vector { index_type, .. } => *index_type,
139 Self::Inverted(_) => IndexType::Inverted,
140 }
141 }
142
143 fn params(&self) -> &dyn IndexParams {
144 match self {
145 Self::Scalar { params, .. } => params,
146 Self::Inverted(params) => params,
147 Self::Vector { params, .. } => params,
148 }
149 }
150}
151
152#[derive(Clone)]
187pub struct DirectoryNamespaceBuilder {
188 root: String,
189 storage_options: Option<HashMap<String, String>>,
190 session: Option<Arc<Session>>,
191 manifest_enabled: bool,
192 dir_listing_enabled: bool,
193 inline_optimization_enabled: bool,
194 table_version_tracking_enabled: bool,
195 table_version_storage_enabled: bool,
198 dir_listing_to_manifest_migration_enabled: bool,
203 credential_vendor_properties: HashMap<String, String>,
204 context_provider: Option<Arc<dyn DynamicContextProvider>>,
205 commit_retries: Option<u32>,
206 vend_input_storage_options: bool,
209 vend_input_storage_options_refresh_interval_millis: Option<u64>,
214 ops_metrics_enabled: bool,
216}
217
218impl std::fmt::Debug for DirectoryNamespaceBuilder {
219 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
220 f.debug_struct("DirectoryNamespaceBuilder")
221 .field("root", &self.root)
222 .field("storage_options", &self.storage_options)
223 .field("manifest_enabled", &self.manifest_enabled)
224 .field("dir_listing_enabled", &self.dir_listing_enabled)
225 .field(
226 "inline_optimization_enabled",
227 &self.inline_optimization_enabled,
228 )
229 .field(
230 "table_version_tracking_enabled",
231 &self.table_version_tracking_enabled,
232 )
233 .field(
234 "table_version_storage_enabled",
235 &self.table_version_storage_enabled,
236 )
237 .field(
238 "dir_listing_to_manifest_migration_enabled",
239 &self.dir_listing_to_manifest_migration_enabled,
240 )
241 .field(
242 "context_provider",
243 &self.context_provider.as_ref().map(|_| "Some(...)"),
244 )
245 .field(
246 "vend_input_storage_options",
247 &self.vend_input_storage_options,
248 )
249 .field(
250 "vend_input_storage_options_refresh_interval_millis",
251 &self.vend_input_storage_options_refresh_interval_millis,
252 )
253 .field("ops_metrics_enabled", &self.ops_metrics_enabled)
254 .finish()
255 }
256}
257
258impl DirectoryNamespaceBuilder {
259 pub fn new(root: impl Into<String>) -> Self {
265 Self {
266 root: root.into().trim_end_matches('/').to_string(),
267 storage_options: None,
268 session: None,
269 manifest_enabled: true,
270 dir_listing_enabled: true, inline_optimization_enabled: true,
272 table_version_tracking_enabled: false, table_version_storage_enabled: false, dir_listing_to_manifest_migration_enabled: false, credential_vendor_properties: HashMap::new(),
276 context_provider: None,
277 commit_retries: None,
278 vend_input_storage_options: false,
279 vend_input_storage_options_refresh_interval_millis: None,
280 ops_metrics_enabled: false,
281 }
282 }
283
284 pub fn manifest_enabled(mut self, enabled: bool) -> Self {
289 self.manifest_enabled = enabled;
290 self
291 }
292
293 pub fn dir_listing_enabled(mut self, enabled: bool) -> Self {
298 self.dir_listing_enabled = enabled;
299 self
300 }
301
302 pub fn dir_listing_to_manifest_migration_enabled(mut self, enabled: bool) -> Self {
309 self.dir_listing_to_manifest_migration_enabled = enabled;
310 self
311 }
312
313 pub fn inline_optimization_enabled(mut self, enabled: bool) -> Self {
319 self.inline_optimization_enabled = enabled;
320 self
321 }
322
323 pub fn table_version_tracking_enabled(mut self, enabled: bool) -> Self {
331 self.table_version_tracking_enabled = enabled;
332 self
333 }
334
335 pub fn table_version_storage_enabled(mut self, enabled: bool) -> Self {
344 self.table_version_storage_enabled = enabled;
345 self
346 }
347
348 pub fn from_properties(
416 properties: HashMap<String, String>,
417 session: Option<Arc<Session>>,
418 ) -> Result<Self> {
419 let root = properties.get("root").cloned().ok_or_else(|| {
421 lance_core::Error::from(NamespaceError::InvalidInput {
422 message: "Missing required property 'root' for directory namespace".to_string(),
423 })
424 })?;
425
426 let storage_options: HashMap<String, String> = properties
428 .iter()
429 .filter_map(|(k, v)| {
430 k.strip_prefix("storage.")
431 .map(|key| (key.to_string(), v.clone()))
432 })
433 .collect();
434
435 let storage_options = if storage_options.is_empty() {
436 None
437 } else {
438 Some(storage_options)
439 };
440
441 let manifest_enabled = properties
443 .get("manifest_enabled")
444 .and_then(|v| v.parse::<bool>().ok())
445 .unwrap_or(true);
446
447 let dir_listing_enabled = properties
449 .get("dir_listing_enabled")
450 .and_then(|v| v.parse::<bool>().ok())
451 .unwrap_or(true);
452
453 let inline_optimization_enabled = properties
455 .get("inline_optimization_enabled")
456 .and_then(|v| v.parse::<bool>().ok())
457 .unwrap_or(true);
458
459 let table_version_tracking_enabled = properties
461 .get("table_version_tracking_enabled")
462 .and_then(|v| v.parse::<bool>().ok())
463 .unwrap_or(false);
464
465 let table_version_storage_enabled = properties
467 .get("table_version_storage_enabled")
468 .and_then(|v| v.parse::<bool>().ok())
469 .unwrap_or(false);
470
471 let dir_listing_to_manifest_migration_enabled = properties
473 .get("dir_listing_to_manifest_migration_enabled")
474 .and_then(|v| v.parse::<bool>().ok())
475 .unwrap_or(false);
476
477 let credential_vendor_properties: HashMap<String, String> = properties
481 .iter()
482 .filter_map(|(k, v)| {
483 k.strip_prefix("credential_vendor.")
484 .map(|key| (key.to_string(), v.clone()))
485 })
486 .collect();
487
488 let commit_retries = properties
489 .get("commit_retries")
490 .and_then(|v| v.parse::<u32>().ok());
491
492 let vend_input_storage_options = properties
494 .get("vend_input_storage_options")
495 .and_then(|v| v.parse::<bool>().ok())
496 .unwrap_or(false);
497
498 let vend_input_storage_options_refresh_interval_millis = properties
500 .get("vend_input_storage_options_refresh_interval_millis")
501 .and_then(|v| v.parse::<u64>().ok());
502
503 let ops_metrics_enabled = properties
505 .get("ops_metrics_enabled")
506 .and_then(|v| v.parse::<bool>().ok())
507 .unwrap_or(false);
508
509 Ok(Self {
510 root: root.trim_end_matches('/').to_string(),
511 storage_options,
512 session,
513 manifest_enabled,
514 dir_listing_enabled,
515 inline_optimization_enabled,
516 table_version_tracking_enabled,
517 table_version_storage_enabled,
518 dir_listing_to_manifest_migration_enabled,
519 credential_vendor_properties,
520 context_provider: None,
521 commit_retries,
522 vend_input_storage_options,
523 vend_input_storage_options_refresh_interval_millis,
524 ops_metrics_enabled,
525 })
526 }
527
528 pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
535 self.storage_options
536 .get_or_insert_with(HashMap::new)
537 .insert(key.into(), value.into());
538 self
539 }
540
541 pub fn storage_options(mut self, options: HashMap<String, String>) -> Self {
547 self.storage_options
548 .get_or_insert_with(HashMap::new)
549 .extend(options);
550 self
551 }
552
553 pub fn session(mut self, session: Arc<Session>) -> Self {
563 self.session = Some(session);
564 self
565 }
566
567 pub fn commit_retries(mut self, retries: u32) -> Self {
570 self.commit_retries = Some(retries);
571 self
572 }
573
574 pub fn credential_vendor_property(
602 mut self,
603 key: impl Into<String>,
604 value: impl Into<String>,
605 ) -> Self {
606 self.credential_vendor_properties
607 .insert(key.into(), value.into());
608 self
609 }
610
611 pub fn credential_vendor_properties(mut self, properties: HashMap<String, String>) -> Self {
619 self.credential_vendor_properties.extend(properties);
620 self
621 }
622
623 pub fn context_provider(mut self, provider: Arc<dyn DynamicContextProvider>) -> Self {
633 self.context_provider = Some(provider);
634 self
635 }
636
637 pub fn vend_input_storage_options(mut self, enabled: bool) -> Self {
646 self.vend_input_storage_options = enabled;
647 self
648 }
649
650 pub fn vend_input_storage_options_refresh_interval_millis(
662 mut self,
663 interval_millis: u64,
664 ) -> Self {
665 self.vend_input_storage_options_refresh_interval_millis = Some(interval_millis);
666 self
667 }
668
669 pub fn ops_metrics_enabled(mut self, enabled: bool) -> Self {
677 self.ops_metrics_enabled = enabled;
678 self
679 }
680
681 pub async fn build(self) -> Result<DirectoryNamespace> {
694 if self.table_version_storage_enabled && !self.manifest_enabled {
696 return Err(NamespaceError::InvalidInput {
697 message: "table_version_storage_enabled requires manifest_enabled=true".to_string(),
698 }
699 .into());
700 }
701
702 let (object_store, base_path) =
703 Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?;
704
705 let manifest_ns = if self.manifest_enabled {
706 match manifest::ManifestNamespace::from_directory(
707 self.root.clone(),
708 self.storage_options.clone(),
709 self.session.clone(),
710 object_store.clone(),
711 base_path.clone(),
712 self.dir_listing_enabled,
713 self.inline_optimization_enabled,
714 self.commit_retries,
715 self.table_version_storage_enabled,
716 )
717 .await
718 {
719 Ok(ns) => Some(Arc::new(ns)),
720 Err(e) => {
721 log::warn!(
723 "Failed to initialize manifest namespace, falling back to directory listing only: {}",
724 e
725 );
726 None
727 }
728 }
729 } else {
730 None
731 };
732
733 let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties)
735 {
736 create_credential_vendor_for_location(&self.root, &self.credential_vendor_properties)
737 .await?
738 .map(Arc::from)
739 } else {
740 None
741 };
742
743 let ops_metrics = if self.ops_metrics_enabled {
744 Some(Arc::new(OpsMetrics::default()))
745 } else {
746 None
747 };
748
749 Ok(DirectoryNamespace {
750 root: self.root,
751 storage_options: self.storage_options,
752 session: self.session,
753 object_store,
754 base_path,
755 manifest_ns,
756 dir_listing_enabled: self.dir_listing_enabled,
757 dir_listing_to_manifest_migration_enabled: self
758 .dir_listing_to_manifest_migration_enabled,
759 table_version_tracking_enabled: self.table_version_tracking_enabled,
760 table_version_storage_enabled: self.table_version_storage_enabled,
761 credential_vendor,
762 context_provider: self.context_provider,
763 vend_input_storage_options: self.vend_input_storage_options,
764 vend_input_storage_options_refresh_interval_millis: self
765 .vend_input_storage_options_refresh_interval_millis,
766 ops_metrics,
767 })
768 }
769
770 async fn initialize_object_store(
772 root: &str,
773 storage_options: &Option<HashMap<String, String>>,
774 session: &Option<Arc<Session>>,
775 ) -> Result<(Arc<ObjectStore>, Path)> {
776 let accessor = storage_options.clone().map(|opts| {
778 Arc::new(lance_io::object_store::StorageOptionsAccessor::with_static_options(opts))
779 });
780 let params = ObjectStoreParams {
781 storage_options_accessor: accessor,
782 ..Default::default()
783 };
784
785 let registry = if let Some(session) = session {
787 session.store_registry()
788 } else {
789 Arc::new(ObjectStoreRegistry::default())
790 };
791
792 let (object_store, base_path) = ObjectStore::from_uri_and_params(registry, root, ¶ms)
794 .await
795 .map_err(|e| {
796 lance_core::Error::from(NamespaceError::Internal {
797 message: format!("Failed to create object store: {:?}", e),
798 })
799 })?;
800
801 Ok((object_store, base_path))
802 }
803}
804
805pub struct DirectoryNamespace {
829 root: String,
830 storage_options: Option<HashMap<String, String>>,
831 session: Option<Arc<Session>>,
832 object_store: Arc<ObjectStore>,
833 base_path: Path,
834 manifest_ns: Option<Arc<manifest::ManifestNamespace>>,
835 dir_listing_enabled: bool,
836 dir_listing_to_manifest_migration_enabled: bool,
840 table_version_tracking_enabled: bool,
843 table_version_storage_enabled: bool,
845 credential_vendor: Option<Arc<dyn CredentialVendor>>,
848 #[allow(dead_code)]
851 context_provider: Option<Arc<dyn DynamicContextProvider>>,
852 vend_input_storage_options: bool,
854 vend_input_storage_options_refresh_interval_millis: Option<u64>,
857 ops_metrics: Option<Arc<OpsMetrics>>,
859}
860
861impl std::fmt::Debug for DirectoryNamespace {
862 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
863 write!(f, "{}", self.namespace_id())
864 }
865}
866
867impl std::fmt::Display for DirectoryNamespace {
868 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
869 write!(f, "{}", self.namespace_id())
870 }
871}
872
873struct TableDeleteEntry {
876 table_id: Option<Vec<String>>,
877 ranges: Vec<(i64, i64)>,
878}
879
880impl DirectoryNamespace {
881 fn apply_pagination(
894 names: &mut Vec<String>,
895 page_token: Option<String>,
896 limit: Option<i32>,
897 ) -> Option<String> {
898 names.sort();
900
901 if let Some(start_after) = page_token {
903 if let Some(index) = names
904 .iter()
905 .position(|name| name.as_str() > start_after.as_str())
906 {
907 names.drain(0..index);
908 } else {
909 names.clear();
910 }
911 }
912
913 if let Some(limit) = limit
915 && limit >= 0
916 {
917 let limit = limit as usize;
918 if names.len() > limit {
919 let next_page_token = if limit > 0 {
920 Some(names[limit - 1].clone())
921 } else {
922 None
923 };
924 names.truncate(limit);
925 return next_page_token;
926 }
927 }
928
929 None
930 }
931
932 async fn list_directory_tables(&self) -> Result<Vec<String>> {
934 let mut tables = Vec::new();
935 let entries = self
936 .object_store
937 .read_dir(self.base_path.clone())
938 .await
939 .map_err(|e| {
940 lance_core::Error::from(NamespaceError::Internal {
941 message: format!("Failed to list directory: {:?}", e),
942 })
943 })?;
944
945 for entry in entries {
946 let path = entry.trim_end_matches('/');
947 if !path.ends_with(".lance") {
948 continue;
949 }
950
951 let table_name = &path[..path.len() - 6];
952
953 let status = self.check_table_status(table_name).await;
955 if status.is_deregistered {
956 continue;
957 }
958
959 tables.push(table_name.to_string());
960 }
961
962 Ok(tables)
963 }
964
965 fn validate_root_namespace_id(id: &Option<Vec<String>>) -> Result<()> {
967 if let Some(id) = id
968 && !id.is_empty()
969 {
970 return Err(NamespaceError::Unsupported {
971 message: format!(
972 "Directory namespace only supports root namespace operations, but got namespace ID: {:?}. Expected empty ID.",
973 id
974 ),
975 }
976 .into());
977 }
978 Ok(())
979 }
980
981 fn table_name_from_id(id: &Option<Vec<String>>) -> Result<String> {
983 let id = id.as_ref().ok_or_else(|| {
984 lance_core::Error::from(NamespaceError::InvalidInput {
985 message: "Directory namespace table ID cannot be empty".to_string(),
986 })
987 })?;
988
989 if id.len() != 1 {
990 return Err(NamespaceError::Unsupported {
991 message: format!(
992 "Multi-level table IDs are only supported when manifest mode is enabled, but got: {:?}",
993 id
994 ),
995 }
996 .into());
997 }
998
999 Ok(id[0].clone())
1000 }
1001
1002 fn format_table_id(table_id: &[String]) -> String {
1003 format!(
1004 "table id '{}'",
1005 manifest::ManifestNamespace::str_object_id(table_id)
1006 )
1007 }
1008
1009 fn format_table_id_from_request(id: &Option<Vec<String>>) -> String {
1010 id.as_ref()
1011 .map(|table_id| Self::format_table_id(table_id))
1012 .unwrap_or_else(|| "table id '<unknown>'".to_string())
1013 }
1014
1015 async fn resolve_table_location(&self, id: &Option<Vec<String>>) -> Result<String> {
1016 let mut describe_req = DescribeTableRequest::new();
1017 describe_req.id = id.clone();
1018 describe_req.load_detailed_metadata = Some(false);
1019
1020 let describe_resp = self.describe_table_impl(describe_req).await?;
1022
1023 describe_resp.location.ok_or_else(|| {
1024 lance_core::Error::from(NamespaceError::TableNotFound {
1025 message: format!("Table location not found for: {:?}", id),
1026 })
1027 })
1028 }
1029
1030 fn map_tag_error(err: lance_core::Error, tag: &str, table_uri: &str) -> lance_core::Error {
1033 match err {
1034 lance_core::Error::RefNotFound { .. } => NamespaceError::TableTagNotFound {
1035 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1036 }
1037 .into(),
1038 lance_core::Error::RefConflict { .. } => NamespaceError::TableTagAlreadyExists {
1039 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1040 }
1041 .into(),
1042 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1043 message: format!("invalid tag '{}': {}", tag, message),
1044 }
1045 .into(),
1046 lance_core::Error::VersionNotFound { message } => {
1047 NamespaceError::TableVersionNotFound {
1048 message: format!(
1049 "version referenced by tag '{}' not found for table at '{}': {}",
1050 tag, table_uri, message
1051 ),
1052 }
1053 .into()
1054 }
1055 other => NamespaceError::Internal {
1056 message: format!(
1057 "tag operation failed for tag '{}' on table at '{}': {}",
1058 tag, table_uri, other
1059 ),
1060 }
1061 .into(),
1062 }
1063 }
1064
1065 async fn table_has_actual_manifests(&self, table_name: &str) -> Result<bool> {
1066 manifest::ManifestNamespace::path_has_actual_manifests(
1067 &self.object_store,
1068 &self.table_path(table_name),
1069 )
1070 .await
1071 }
1072
1073 async fn filter_declared_tables(
1074 &self,
1075 tables: Vec<String>,
1076 include_declared: bool,
1077 ) -> Result<Vec<String>> {
1078 if include_declared {
1079 return Ok(tables);
1080 }
1081
1082 let mut stream = futures::stream::iter(tables.into_iter().map(|table_name| async move {
1083 if self.table_has_actual_manifests(&table_name).await? {
1087 Ok::<Option<String>, Error>(Some(table_name))
1088 } else {
1089 Ok::<Option<String>, Error>(None)
1090 }
1091 }))
1092 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1093
1094 let mut filtered = Vec::new();
1095 while let Some(result) = stream.next().await {
1096 if let Some(table_name) = result? {
1097 filtered.push(table_name);
1098 }
1099 }
1100 Ok(filtered)
1101 }
1102
1103 fn ipc_reader_from_request_data(
1104 request_data: &Bytes,
1105 operation: &str,
1106 ) -> Result<(
1107 Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1108 usize,
1109 )> {
1110 if request_data.is_empty() {
1111 return Err(NamespaceError::InvalidInput {
1112 message: format!(
1113 "Request data (Arrow IPC stream) is required for {}",
1114 operation
1115 ),
1116 }
1117 .into());
1118 }
1119
1120 let cursor = Cursor::new(request_data.as_ref());
1121 let stream_reader =
1122 StreamReader::try_new(cursor, None).map_err(|e| NamespaceError::InvalidInput {
1123 message: format!("Invalid Arrow IPC stream: {}", e),
1124 })?;
1125 let arrow_schema = stream_reader.schema();
1126
1127 let mut num_rows = 0usize;
1128 let mut batches = Vec::new();
1129 for batch_result in stream_reader {
1130 let batch = batch_result.map_err(|e| NamespaceError::Internal {
1131 message: format!("Failed to read batch from IPC stream: {}", e),
1132 })?;
1133 num_rows += batch.num_rows();
1134 batches.push(batch);
1135 }
1136
1137 let reader: Box<dyn arrow::record_batch::RecordBatchReader + Send> = if batches.is_empty() {
1138 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
1139 Box::new(RecordBatchIterator::new(vec![Ok(batch)], arrow_schema))
1140 } else {
1141 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
1142 Box::new(RecordBatchIterator::new(batch_results, arrow_schema))
1143 };
1144
1145 Ok((reader, num_rows))
1146 }
1147
1148 async fn table_uri_has_actual_manifests(&self, table_uri: &str) -> Result<bool> {
1149 let table_path = self.object_store_path_from_uri(table_uri)?;
1150 manifest::ManifestNamespace::path_has_actual_manifests(&self.object_store, &table_path)
1151 .await
1152 }
1153
1154 fn object_store_path_from_uri(&self, uri: &str) -> Result<Path> {
1155 let registry = self
1156 .session
1157 .as_ref()
1158 .map(|session| session.store_registry())
1159 .unwrap_or_else(|| Arc::new(ObjectStoreRegistry::default()));
1160 ObjectStore::extract_path_from_uri(registry, uri)
1161 }
1162
1163 fn validate_dir_only_properties(
1164 properties: Option<&HashMap<String, String>>,
1165 operation: &str,
1166 ) -> Result<()> {
1167 if properties.is_some_and(|properties| !properties.is_empty()) {
1171 return Err(NamespaceError::Unsupported {
1172 message: format!(
1173 "{} with non-empty table properties requires manifest_enabled=true",
1174 operation
1175 ),
1176 }
1177 .into());
1178 }
1179 Ok(())
1180 }
1181
1182 async fn write_reader_to_table(
1183 &self,
1184 table_uri: &str,
1185 reader: Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1186 mode: WriteMode,
1187 extra_storage_options: Option<HashMap<String, String>>,
1188 ) -> Result<Dataset> {
1189 let mut merged_storage_options = self.storage_options.clone().unwrap_or_default();
1192 if let Some(extra_storage_options) = extra_storage_options {
1193 merged_storage_options.extend(extra_storage_options);
1194 }
1195 let store_params = (!merged_storage_options.is_empty()).then(|| ObjectStoreParams {
1196 storage_options_accessor: Some(Arc::new(
1197 lance_io::object_store::StorageOptionsAccessor::with_static_options(
1198 merged_storage_options,
1199 ),
1200 )),
1201 ..Default::default()
1202 });
1203
1204 let write_params = WriteParams {
1205 mode,
1206 store_params,
1207 session: self.session.clone(),
1208 ..Default::default()
1209 };
1210
1211 let dataset = Dataset::write(reader, table_uri, Some(write_params))
1212 .await
1213 .map_err(|e| NamespaceError::Internal {
1214 message: format!("Failed to write table at '{}': {}", table_uri, e),
1215 })?;
1216
1217 Ok(dataset)
1218 }
1219
1220 async fn list_table_versions_from_storage(
1221 &self,
1222 table_uri: &str,
1223 descending: bool,
1224 limit: Option<i32>,
1225 ) -> Result<Vec<TableVersion>> {
1226 let table_path = self.object_store_path_from_uri(table_uri)?;
1227 let versions_dir = table_path.clone().join(VERSIONS_DIR);
1228 let manifest_metas: Vec<_> = self
1229 .object_store
1230 .read_dir_all(&versions_dir, None)
1231 .try_collect()
1232 .await
1233 .map_err(|e| {
1234 lance_core::Error::from(NamespaceError::Internal {
1235 message: format!(
1236 "Failed to list manifest files for table at '{}': {}",
1237 table_uri, e
1238 ),
1239 })
1240 })?;
1241
1242 let is_v2_naming = manifest_metas
1243 .first()
1244 .is_some_and(|meta| meta.location.filename().is_some_and(|f| f.len() == 29));
1245
1246 let mut table_versions: Vec<TableVersion> = manifest_metas
1247 .into_iter()
1248 .filter_map(|meta| {
1249 let filename = meta.location.filename()?;
1250 let version_str = filename.strip_suffix(".manifest")?;
1251 if version_str.starts_with('d') {
1252 return None;
1253 }
1254 let file_version: u64 = version_str.parse().ok()?;
1255
1256 let actual_version = if file_version > u64::MAX / 2 {
1257 u64::MAX - file_version
1258 } else {
1259 file_version
1260 };
1261
1262 Some(TableVersion {
1263 version: actual_version as i64,
1264 manifest_path: meta.location.to_string(),
1265 manifest_size: Some(meta.size as i64),
1266 e_tag: meta.e_tag,
1267 timestamp_millis: Some(meta.last_modified.timestamp_millis()),
1268 metadata: None,
1269 })
1270 })
1271 .collect();
1272
1273 let list_is_ordered = self.object_store.list_is_lexically_ordered;
1274
1275 let needs_sort = if list_is_ordered {
1276 if is_v2_naming {
1277 !descending
1278 } else {
1279 descending
1280 }
1281 } else {
1282 true
1283 };
1284
1285 if needs_sort {
1286 if descending {
1287 table_versions.sort_by(|a, b| b.version.cmp(&a.version));
1288 } else {
1289 table_versions.sort_by(|a, b| a.version.cmp(&b.version));
1290 }
1291 }
1292
1293 if let Some(limit) = limit {
1294 table_versions.truncate(limit as usize);
1295 }
1296
1297 Ok(table_versions)
1298 }
1299
1300 async fn describe_table_impl(
1304 &self,
1305 request: DescribeTableRequest,
1306 ) -> Result<DescribeTableResponse> {
1307 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
1308 let skip_manifest_for_root = self.dir_listing_enabled
1309 && is_root_level
1310 && !self.dir_listing_to_manifest_migration_enabled;
1311 if let Some(ref manifest_ns) = self.manifest_ns
1312 && !skip_manifest_for_root
1313 {
1314 match manifest_ns.describe_table(request.clone()).await {
1315 Ok(mut response) => {
1316 if let Some(ref table_uri) = response.table_uri {
1317 let vend = request.vend_credentials.unwrap_or(true);
1319 let identity = request.identity.as_deref();
1320 response.storage_options = self
1321 .get_storage_options_for_table(table_uri, vend, identity)
1322 .await?;
1323 }
1324 if self.table_version_tracking_enabled {
1326 response.managed_versioning = Some(true);
1327 }
1328 return Ok(response);
1329 }
1330 Err(_) if self.dir_listing_enabled && is_root_level => {
1331 }
1333 Err(e) => return Err(e),
1334 }
1335 }
1336
1337 let table_name = Self::table_name_from_id(&request.id)?;
1338 let table_id = Self::format_table_id_from_request(&request.id);
1339 let table_uri = self.table_full_uri(&table_name);
1340
1341 let status = self.check_table_status(&table_name).await;
1343
1344 if !status.exists {
1345 return Err(NamespaceError::TableNotFound {
1346 message: table_id.clone(),
1347 }
1348 .into());
1349 }
1350
1351 if status.is_deregistered {
1352 return Err(NamespaceError::TableNotFound {
1353 message: format!("Table is deregistered: {}", table_id),
1354 }
1355 .into());
1356 }
1357
1358 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
1359 let should_check_declared =
1360 load_detailed_metadata || request.check_declared.unwrap_or(false);
1361 let vend_credentials = request.vend_credentials.unwrap_or(true);
1363 let identity = request.identity.as_deref();
1364 let is_only_declared = if should_check_declared {
1365 if status.has_reserved_file {
1366 Some(!self.table_has_actual_manifests(&table_name).await?)
1367 } else {
1368 Some(false)
1369 }
1370 } else {
1371 None
1372 };
1373
1374 if !load_detailed_metadata {
1375 let storage_options = self
1376 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1377 .await?;
1378 return Ok(DescribeTableResponse {
1379 table: Some(table_name),
1380 namespace: request.id.as_ref().map(|id| {
1381 if id.len() > 1 {
1382 id[..id.len() - 1].to_vec()
1383 } else {
1384 vec![]
1385 }
1386 }),
1387 location: Some(table_uri.clone()),
1388 table_uri: Some(table_uri),
1389 storage_options,
1390 is_only_declared,
1391 managed_versioning: if self.table_version_tracking_enabled {
1392 Some(true)
1393 } else {
1394 None
1395 },
1396 ..Default::default()
1397 });
1398 }
1399
1400 if is_only_declared == Some(true) {
1401 let storage_options = self
1402 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1403 .await?;
1404 return Ok(DescribeTableResponse {
1405 table: Some(table_name),
1406 namespace: request.id.as_ref().map(|id| {
1407 if id.len() > 1 {
1408 id[..id.len() - 1].to_vec()
1409 } else {
1410 vec![]
1411 }
1412 }),
1413 location: Some(table_uri.clone()),
1414 table_uri: Some(table_uri),
1415 storage_options,
1416 is_only_declared,
1417 managed_versioning: if self.table_version_tracking_enabled {
1418 Some(true)
1419 } else {
1420 None
1421 },
1422 ..Default::default()
1423 });
1424 }
1425
1426 let mut builder = DatasetBuilder::from_uri(&table_uri);
1429 if let Some(opts) = &self.storage_options {
1430 builder = builder.with_storage_options(opts.clone());
1431 }
1432 if let Some(sess) = &self.session {
1433 builder = builder.with_session(sess.clone());
1434 }
1435 match builder.load().await {
1436 Ok(mut dataset) => {
1437 if let Some(requested_version) = request.version {
1439 dataset = dataset
1440 .checkout_version(requested_version as u64)
1441 .await
1442 .map_err(|e| {
1443 lance_core::Error::from(NamespaceError::TableVersionNotFound {
1444 message: format!(
1445 "Version {} not found for table '{}': {}",
1446 requested_version, table_name, e
1447 ),
1448 })
1449 })?;
1450 }
1451
1452 let version_info = dataset.version();
1453 let lance_schema = dataset.schema();
1454 let arrow_schema: arrow_schema::Schema = lance_schema.into();
1455 let json_schema = arrow_schema_to_json(&arrow_schema)?;
1456 let storage_options = self
1457 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1458 .await?;
1459
1460 let metadata: std::collections::HashMap<String, String> =
1462 version_info.metadata.into_iter().collect();
1463
1464 Ok(DescribeTableResponse {
1465 table: Some(table_name),
1466 namespace: request.id.as_ref().map(|id| {
1467 if id.len() > 1 {
1468 id[..id.len() - 1].to_vec()
1469 } else {
1470 vec![]
1471 }
1472 }),
1473 version: Some(version_info.version as i64),
1474 location: Some(table_uri.clone()),
1475 table_uri: Some(table_uri),
1476 schema: Some(Box::new(json_schema)),
1477 storage_options,
1478 metadata: Some(metadata),
1479 is_only_declared,
1480 managed_versioning: if self.table_version_tracking_enabled {
1481 Some(true)
1482 } else {
1483 None
1484 },
1485 ..Default::default()
1486 })
1487 }
1488 Err(err) => {
1489 if manifest::ManifestNamespace::is_not_found_load_error(&err)
1490 && is_only_declared == Some(true)
1491 {
1492 let storage_options = self
1493 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
1494 .await?;
1495 Ok(DescribeTableResponse {
1496 table: Some(table_name),
1497 namespace: request.id.as_ref().map(|id| {
1498 if id.len() > 1 {
1499 id[..id.len() - 1].to_vec()
1500 } else {
1501 vec![]
1502 }
1503 }),
1504 location: Some(table_uri.clone()),
1505 table_uri: Some(table_uri),
1506 storage_options,
1507 is_only_declared,
1508 managed_versioning: if self.table_version_tracking_enabled {
1509 Some(true)
1510 } else {
1511 None
1512 },
1513 ..Default::default()
1514 })
1515 } else {
1516 Err(NamespaceError::Internal {
1517 message: format!(
1518 "Table directory exists but cannot load dataset {}: {:?}",
1519 table_name, err
1520 ),
1521 }
1522 .into())
1523 }
1524 }
1525 }
1526 }
1527
1528 async fn load_dataset(
1529 &self,
1530 table_uri: &str,
1531 version: Option<i64>,
1532 operation: &str,
1533 ) -> Result<Dataset> {
1534 if let Some(version) = version
1535 && version < 0
1536 {
1537 return Err(NamespaceError::InvalidInput {
1538 message: format!(
1539 "Table version for {} must be non-negative, got {}",
1540 operation, version
1541 ),
1542 }
1543 .into());
1544 }
1545
1546 let mut builder = DatasetBuilder::from_uri(table_uri);
1547 if let Some(opts) = &self.storage_options {
1548 builder = builder.with_storage_options(opts.clone());
1549 }
1550 if let Some(sess) = &self.session {
1551 builder = builder.with_session(sess.clone());
1552 }
1553
1554 let dataset = builder.load().await.map_err(|e| {
1555 lance_core::Error::from(NamespaceError::TableNotFound {
1556 message: format!(
1557 "Failed to open table at '{}' for {}: {}",
1558 table_uri, operation, e
1559 ),
1560 })
1561 })?;
1562
1563 if let Some(version) = version {
1564 return dataset.checkout_version(version as u64).await.map_err(|e| {
1565 lance_core::Error::from(NamespaceError::TableVersionNotFound {
1566 message: format!(
1567 "Failed to checkout version {} for table at '{}' during {}: {}",
1568 version, table_uri, operation, e
1569 ),
1570 })
1571 });
1572 }
1573
1574 Ok(dataset)
1575 }
1576
1577 fn parse_index_type(index_type: &str) -> Result<IndexType> {
1578 match index_type.trim().to_ascii_uppercase().as_str() {
1579 "SCALAR" | "BTREE" => Ok(IndexType::BTree),
1580 "BITMAP" => Ok(IndexType::Bitmap),
1581 "LABEL_LIST" | "LABELLIST" => Ok(IndexType::LabelList),
1582 "INVERTED" | "FTS" => Ok(IndexType::Inverted),
1583 "NGRAM" => Ok(IndexType::NGram),
1584 "ZONEMAP" | "ZONE_MAP" => Ok(IndexType::ZoneMap),
1585 "BLOOMFILTER" | "BLOOM_FILTER" => Ok(IndexType::BloomFilter),
1586 "RTREE" | "R_TREE" => Ok(IndexType::RTree),
1587 "VECTOR" | "IVF_PQ" => Ok(IndexType::IvfPq),
1588 "IVF_FLAT" => Ok(IndexType::IvfFlat),
1589 "IVF_SQ" => Ok(IndexType::IvfSq),
1590 "IVF_RQ" => Ok(IndexType::IvfRq),
1591 "IVF_HNSW_FLAT" => Ok(IndexType::IvfHnswFlat),
1592 "IVF_HNSW_SQ" => Ok(IndexType::IvfHnswSq),
1593 "IVF_HNSW_PQ" => Ok(IndexType::IvfHnswPq),
1594 other => Err(NamespaceError::InvalidInput {
1595 message: format!("Unsupported index_type '{}'", other),
1596 }
1597 .into()),
1598 }
1599 }
1600
1601 fn parse_metric_type(distance_type: Option<&str>) -> Result<MetricType> {
1602 let distance_type = distance_type.unwrap_or("l2");
1603 MetricType::try_from(distance_type).map_err(|e| {
1604 lance_core::Error::from(NamespaceError::InvalidInput {
1605 message: format!(
1606 "Unsupported distance_type '{}' for vector index: {}",
1607 distance_type, e
1608 ),
1609 })
1610 })
1611 }
1612
1613 fn build_index_params(request: &CreateTableIndexRequest) -> Result<DirectoryIndexParams> {
1614 let index_type = Self::parse_index_type(&request.index_type)?;
1615 Ok(match index_type {
1616 IndexType::BTree => DirectoryIndexParams::Scalar {
1617 index_type,
1618 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
1619 },
1620 IndexType::Bitmap => DirectoryIndexParams::Scalar {
1621 index_type,
1622 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
1623 },
1624 IndexType::LabelList => DirectoryIndexParams::Scalar {
1625 index_type,
1626 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
1627 },
1628 IndexType::NGram => DirectoryIndexParams::Scalar {
1629 index_type,
1630 params: ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
1631 },
1632 IndexType::ZoneMap => DirectoryIndexParams::Scalar {
1633 index_type,
1634 params: ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
1635 },
1636 IndexType::BloomFilter => DirectoryIndexParams::Scalar {
1637 index_type,
1638 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter),
1639 },
1640 IndexType::RTree => DirectoryIndexParams::Scalar {
1641 index_type,
1642 params: ScalarIndexParams::for_builtin(BuiltinIndexType::RTree),
1643 },
1644 IndexType::Inverted => {
1645 let mut params = InvertedIndexParams::default();
1646 if let Some(with_position) = request.with_position {
1647 params = params.with_position(with_position);
1648 }
1649 if let Some(base_tokenizer) = &request.base_tokenizer {
1650 params = params.base_tokenizer(base_tokenizer.clone());
1651 }
1652 if let Some(language) = &request.language {
1653 params = params.language(language)?;
1654 }
1655 if let Some(max_token_length) = request.max_token_length {
1656 if max_token_length < 0 {
1657 return Err(NamespaceError::InvalidInput {
1658 message: format!(
1659 "FTS max_token_length must be non-negative, got {}",
1660 max_token_length
1661 ),
1662 }
1663 .into());
1664 }
1665 params = params.max_token_length(Some(max_token_length as usize));
1666 }
1667 if let Some(lower_case) = request.lower_case {
1668 params = params.lower_case(lower_case);
1669 }
1670 if let Some(stem) = request.stem {
1671 params = params.stem(stem);
1672 }
1673 if let Some(remove_stop_words) = request.remove_stop_words {
1674 params = params.remove_stop_words(remove_stop_words);
1675 }
1676 if let Some(ascii_folding) = request.ascii_folding {
1677 params = params.ascii_folding(ascii_folding);
1678 }
1679 DirectoryIndexParams::Inverted(params)
1680 }
1681 IndexType::IvfFlat => DirectoryIndexParams::Vector {
1682 index_type,
1683 params: VectorIndexParams::with_ivf_flat_params(
1684 Self::parse_metric_type(request.distance_type.as_deref())?,
1685 IvfBuildParams::default(),
1686 ),
1687 },
1688 IndexType::IvfPq => DirectoryIndexParams::Vector {
1689 index_type,
1690 params: VectorIndexParams::with_ivf_pq_params(
1691 Self::parse_metric_type(request.distance_type.as_deref())?,
1692 IvfBuildParams::default(),
1693 PQBuildParams::default(),
1694 ),
1695 },
1696 IndexType::IvfSq => DirectoryIndexParams::Vector {
1697 index_type,
1698 params: VectorIndexParams::with_ivf_sq_params(
1699 Self::parse_metric_type(request.distance_type.as_deref())?,
1700 IvfBuildParams::default(),
1701 SQBuildParams::default(),
1702 ),
1703 },
1704 IndexType::IvfRq => DirectoryIndexParams::Vector {
1705 index_type,
1706 params: VectorIndexParams::with_ivf_rq_params(
1707 Self::parse_metric_type(request.distance_type.as_deref())?,
1708 IvfBuildParams::default(),
1709 RQBuildParams::default(),
1710 ),
1711 },
1712 IndexType::IvfHnswFlat => DirectoryIndexParams::Vector {
1713 index_type,
1714 params: VectorIndexParams::ivf_hnsw(
1715 Self::parse_metric_type(request.distance_type.as_deref())?,
1716 IvfBuildParams::default(),
1717 HnswBuildParams::default(),
1718 ),
1719 },
1720 IndexType::IvfHnswSq => DirectoryIndexParams::Vector {
1721 index_type,
1722 params: VectorIndexParams::with_ivf_hnsw_sq_params(
1723 Self::parse_metric_type(request.distance_type.as_deref())?,
1724 IvfBuildParams::default(),
1725 HnswBuildParams::default(),
1726 SQBuildParams::default(),
1727 ),
1728 },
1729 IndexType::IvfHnswPq => DirectoryIndexParams::Vector {
1730 index_type,
1731 params: VectorIndexParams::with_ivf_hnsw_pq_params(
1732 Self::parse_metric_type(request.distance_type.as_deref())?,
1733 IvfBuildParams::default(),
1734 HnswBuildParams::default(),
1735 PQBuildParams::default(),
1736 ),
1737 },
1738 other => {
1739 return Err(NamespaceError::InvalidInput {
1740 message: format!("Unsupported index type for namespace API: {}", other),
1741 }
1742 .into());
1743 }
1744 })
1745 }
1746
1747 fn paginate_indices(
1748 indices: &mut Vec<IndexContent>,
1749 page_token: Option<String>,
1750 limit: Option<i32>,
1751 ) -> Option<String> {
1752 indices.sort_by(|a, b| a.index_name.cmp(&b.index_name));
1753
1754 if let Some(start_after) = page_token {
1755 if let Some(index) = indices
1756 .iter()
1757 .position(|index| index.index_name.as_str() > start_after.as_str())
1758 {
1759 indices.drain(0..index);
1760 } else {
1761 indices.clear();
1762 }
1763 }
1764
1765 let mut next_page_token = None;
1766 if let Some(limit) = limit
1767 && limit >= 0
1768 {
1769 let limit = limit as usize;
1770 if limit > 0 && indices.len() > limit {
1771 next_page_token = Some(indices[limit - 1].index_name.clone());
1772 }
1773 indices.truncate(limit);
1774 }
1775 if indices.is_empty() {
1776 None
1777 } else {
1778 next_page_token
1779 }
1780 }
1781
1782 fn transaction_operation_name(transaction: &Transaction) -> String {
1783 match &transaction.operation {
1784 Operation::CreateIndex {
1785 new_indices,
1786 removed_indices,
1787 } if new_indices.is_empty() && !removed_indices.is_empty() => "DropIndex".to_string(),
1788 _ => transaction.operation.to_string(),
1789 }
1790 }
1791
1792 fn transaction_response(
1793 version: u64,
1794 transaction: &Transaction,
1795 ) -> DescribeTransactionResponse {
1796 let mut properties = transaction
1797 .transaction_properties
1798 .as_ref()
1799 .map(|properties| (**properties).clone())
1800 .unwrap_or_default();
1801 properties.insert("uuid".to_string(), transaction.uuid.clone());
1802 properties.insert("version".to_string(), version.to_string());
1803 properties.insert(
1804 "read_version".to_string(),
1805 transaction.read_version.to_string(),
1806 );
1807 properties.insert(
1808 "operation".to_string(),
1809 Self::transaction_operation_name(transaction),
1810 );
1811 if let Some(tag) = &transaction.tag {
1812 properties.insert("tag".to_string(), tag.clone());
1813 }
1814
1815 DescribeTransactionResponse {
1816 status: "SUCCEEDED".to_string(),
1817 properties: Some(properties),
1818 }
1819 }
1820
1821 fn describe_table_index_stats_response(
1822 stats: &serde_json::Value,
1823 ) -> DescribeTableIndexStatsResponse {
1824 let get_i64 = |key: &str| {
1825 stats.get(key).and_then(|value| {
1826 value
1827 .as_i64()
1828 .or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
1829 })
1830 };
1831
1832 DescribeTableIndexStatsResponse {
1833 distance_type: stats
1834 .get("distance_type")
1835 .and_then(|value| value.as_str())
1836 .map(str::to_string),
1837 index_type: stats
1838 .get("index_type")
1839 .and_then(|value| value.as_str())
1840 .map(str::to_string),
1841 num_indexed_rows: get_i64("num_indexed_rows"),
1842 num_unindexed_rows: get_i64("num_unindexed_rows"),
1843 num_indices: get_i64("num_indices").and_then(|value| i32::try_from(value).ok()),
1844 }
1845 }
1846
1847 async fn find_transaction(&self, dataset: &Dataset, id: &str) -> Result<(u64, Transaction)> {
1852 if let Ok(version) = id.parse::<u64>() {
1853 let transaction = dataset
1854 .read_transaction_by_version(version)
1855 .await
1856 .map_err(|e| {
1857 lance_core::Error::from(NamespaceError::TransactionNotFound {
1858 message: format!(
1859 "Failed to read transaction for version {}: {}",
1860 version, e
1861 ),
1862 })
1863 })?
1864 .ok_or_else(|| {
1865 lance_core::Error::from(NamespaceError::TransactionNotFound {
1866 message: format!("version {}", version),
1867 })
1868 })?;
1869 return Ok((version, transaction));
1870 }
1871
1872 let versions = dataset.versions().await.map_err(|e| {
1873 lance_core::Error::from(NamespaceError::Internal {
1874 message: format!(
1875 "Failed to list table versions while resolving transaction '{}': {}",
1876 id, e
1877 ),
1878 })
1879 })?;
1880
1881 for version in versions.into_iter().rev() {
1882 if let Some(transaction) = dataset
1883 .read_transaction_by_version(version.version)
1884 .await
1885 .map_err(|e| {
1886 lance_core::Error::from(NamespaceError::Internal {
1887 message: format!(
1888 "Failed to read transaction for version {} while resolving '{}': {}",
1889 version.version, id, e
1890 ),
1891 })
1892 })?
1893 && transaction.uuid == id
1894 {
1895 return Ok((version.version, transaction));
1896 }
1897 }
1898
1899 Err(NamespaceError::TransactionNotFound {
1900 message: id.to_string(),
1901 }
1902 .into())
1903 }
1904
1905 fn table_full_uri(&self, table_name: &str) -> String {
1906 format!("{}/{}.lance", &self.root, table_name)
1907 }
1908
1909 fn table_path(&self, table_name: &str) -> Path {
1911 self.base_path
1912 .clone()
1913 .join(format!("{}.lance", table_name).as_str())
1914 }
1915
1916 fn table_reserved_file_path(&self, table_name: &str) -> Path {
1918 self.base_path
1919 .clone()
1920 .join(format!("{}.lance", table_name).as_str())
1921 .join(".lance-reserved")
1922 }
1923
1924 fn table_deregistered_file_path(&self, table_name: &str) -> Path {
1926 self.base_path
1927 .clone()
1928 .join(format!("{}.lance", table_name).as_str())
1929 .join(".lance-deregistered")
1930 }
1931
1932 pub(crate) async fn check_table_status(&self, table_name: &str) -> TableStatus {
1938 let table_path = self.table_path(table_name);
1939 match self.object_store.read_dir(table_path).await {
1940 Ok(entries) => {
1941 let exists = !entries.is_empty();
1942 let is_deregistered = entries.iter().any(|e| e.ends_with(".lance-deregistered"));
1943 let has_reserved_file = entries.iter().any(|e| e.ends_with(".lance-reserved"));
1944 TableStatus {
1945 exists,
1946 is_deregistered,
1947 has_reserved_file,
1948 }
1949 }
1950 Err(_) => TableStatus {
1951 exists: false,
1952 is_deregistered: false,
1953 has_reserved_file: false,
1954 },
1955 }
1956 }
1957
1958 async fn put_marker_file_atomic(
1959 &self,
1960 path: &Path,
1961 file_description: &str,
1962 ) -> std::result::Result<(), String> {
1963 let put_opts = PutOptions {
1964 mode: PutMode::Create,
1965 ..Default::default()
1966 };
1967
1968 match self
1969 .object_store
1970 .inner
1971 .put_opts(path, bytes::Bytes::new().into(), put_opts)
1972 .await
1973 {
1974 Ok(_) => Ok(()),
1975 Err(ObjectStoreError::AlreadyExists { .. })
1976 | Err(ObjectStoreError::Precondition { .. }) => {
1977 Err(format!("{} already exists", file_description))
1978 }
1979 Err(e) => Err(format!("Failed to create {}: {:?}", file_description, e)),
1980 }
1981 }
1982
1983 async fn get_storage_options_for_table(
2003 &self,
2004 table_uri: &str,
2005 vend_credentials: bool,
2006 identity: Option<&Identity>,
2007 ) -> Result<Option<HashMap<String, String>>> {
2008 if vend_credentials && let Some(ref vendor) = self.credential_vendor {
2009 let vended = vendor.vend_credentials(table_uri, identity).await?;
2010 return Ok(Some(vended.storage_options));
2011 }
2012 if self.vend_input_storage_options {
2015 let mut options = self.storage_options.clone().unwrap_or_default();
2016 if let Some(refresh_interval_millis) =
2018 self.vend_input_storage_options_refresh_interval_millis
2019 {
2020 let now_millis = std::time::SystemTime::now()
2021 .duration_since(std::time::UNIX_EPOCH)
2022 .unwrap()
2023 .as_millis() as u64;
2024 let expires_at_millis = now_millis + refresh_interval_millis;
2025 options.insert(
2026 "expires_at_millis".to_string(),
2027 expires_at_millis.to_string(),
2028 );
2029 }
2030 return Ok(Some(options));
2031 }
2032 Ok(None)
2035 }
2036
2037 pub async fn migrate(&self) -> Result<usize> {
2090 let Some(ref manifest_ns) = self.manifest_ns else {
2092 return Ok(0); };
2094
2095 let manifest_locations = manifest_ns.list_manifest_table_locations().await?;
2097
2098 let dir_tables = self
2101 .filter_declared_tables(self.list_directory_tables().await?, false)
2102 .await?;
2103
2104 let mut migrated_count = 0;
2109 for table_name in dir_tables {
2110 let dir_name = format!("{}.lance", table_name);
2112 if !manifest_locations.contains(&dir_name) {
2113 manifest_ns.register_table(&table_name, dir_name).await?;
2114 migrated_count += 1;
2115 }
2116 }
2117
2118 Ok(migrated_count)
2119 }
2120
2121 async fn delete_physical_version_files(
2130 &self,
2131 table_entries: &[TableDeleteEntry],
2132 best_effort: bool,
2133 ) -> Result<i64> {
2134 let mut deleted_count = 0i64;
2135 for te in table_entries {
2136 let table_uri = self.resolve_table_location(&te.table_id).await?;
2137 let table_path = self.object_store_path_from_uri(&table_uri)?;
2138 let versions_dir_path = table_path.clone().join(VERSIONS_DIR);
2139
2140 for (start, end) in &te.ranges {
2141 for version in *start..=*end {
2142 let version_path = versions_dir_path
2143 .clone()
2144 .join(format!("{}.manifest", version as u64));
2145 match self.object_store.inner.delete(&version_path).await {
2146 Ok(_) => {
2147 deleted_count += 1;
2148 }
2149 Err(object_store::Error::NotFound { .. }) => {}
2150 Err(e) => {
2151 if best_effort {
2152 log::warn!(
2153 "Failed to delete manifest file for version {} of table {:?}: {:?}",
2154 version,
2155 te.table_id,
2156 e
2157 );
2158 } else {
2159 return Err(NamespaceError::Internal {
2160 message: format!(
2161 "Failed to delete version {} for table at '{}': {}",
2162 version, table_uri, e
2163 ),
2164 }
2165 .into());
2166 }
2167 }
2168 }
2169 }
2170 }
2171 }
2172 Ok(deleted_count)
2173 }
2174
2175 #[allow(clippy::too_many_arguments)]
2180 fn apply_query_params_to_scanner(
2181 scanner: &mut Scanner,
2182 filter: Option<&str>,
2183 columns: Option<&QueryTableRequestColumns>,
2184 vector_column: Option<&str>,
2185 vector: &QueryTableRequestVector,
2186 k: i32,
2187 offset: Option<i32>,
2188 prefilter: Option<bool>,
2189 bypass_vector_index: Option<bool>,
2190 nprobes: Option<i32>,
2191 ef: Option<i32>,
2192 refine_factor: Option<i32>,
2193 distance_type: Option<&str>,
2194 fast_search_flag: Option<bool>,
2195 with_row_id: Option<bool>,
2196 lower_bound: Option<f32>,
2197 upper_bound: Option<f32>,
2198 operation: &str,
2199 ) -> Result<()> {
2200 if let Some(pf) = prefilter {
2202 scanner.prefilter(pf);
2203 }
2204
2205 if let Some(filter) = filter {
2206 scanner.filter(filter).map_err(|e| {
2207 Error::invalid_input_source(
2208 format!("Invalid filter expression for {}: {}", operation, e).into(),
2209 )
2210 })?;
2211 }
2212
2213 if let Some(cols) = columns {
2214 if let Some(ref names) = cols.column_names {
2215 scanner.project(names.as_slice()).map_err(|e| {
2216 Error::invalid_input_source(
2217 format!("Invalid column projection for {}: {}", operation, e).into(),
2218 )
2219 })?;
2220 } else if let Some(ref aliases) = cols.column_aliases {
2221 let pairs: Vec<(&str, &str)> = aliases
2223 .iter()
2224 .map(|(alias, src)| (alias.as_str(), src.as_str()))
2225 .collect();
2226 scanner.project_with_transform(&pairs).map_err(|e| {
2227 Error::invalid_input_source(
2228 format!("Invalid column aliases for {}: {}", operation, e).into(),
2229 )
2230 })?;
2231 }
2232 }
2233
2234 let query_vec: Option<Vec<f32>> = vector
2236 .single_vector
2237 .as_ref()
2238 .filter(|v| !v.is_empty())
2239 .cloned()
2240 .or_else(|| {
2241 vector
2242 .multi_vector
2243 .as_ref()
2244 .and_then(|mv| mv.first())
2245 .filter(|v| !v.is_empty())
2246 .cloned()
2247 });
2248
2249 if let Some(q_vec) = query_vec {
2250 let col = vector_column.unwrap_or("vector");
2251 let q = Arc::new(Float32Array::from(q_vec));
2252 scanner
2253 .nearest(col, q.as_ref(), k.max(1) as usize)
2254 .map_err(|e| {
2255 Error::invalid_input_source(
2256 format!("Invalid vector query for {}: {}", operation, e).into(),
2257 )
2258 })?;
2259
2260 if let Some(n) = nprobes {
2262 scanner.nprobes(n.max(1) as usize);
2263 }
2264 if let Some(e) = ef {
2265 scanner.ef(e.max(1) as usize);
2266 }
2267 if let Some(rf) = refine_factor {
2268 scanner.refine(rf.max(0) as u32);
2269 }
2270 if let Some(true) = bypass_vector_index {
2272 scanner.use_index(false);
2273 }
2274 if let Some(true) = fast_search_flag {
2275 scanner.fast_search();
2276 }
2277 if lower_bound.is_some() || upper_bound.is_some() {
2278 scanner.distance_range(lower_bound, upper_bound);
2279 }
2280 if let Some(dt) = distance_type {
2281 let metric = Self::parse_metric_type(Some(dt))?;
2282 scanner.distance_metric(metric);
2283 }
2284 if let Some(off) = offset.filter(|&o| o > 0) {
2286 scanner.limit(None, Some(off as i64)).map_err(|e| {
2287 Error::invalid_input_source(
2288 format!("Invalid offset for {}: {}", operation, e).into(),
2289 )
2290 })?;
2291 }
2292 } else {
2293 let limit = if k > 0 { Some(k as i64) } else { None };
2295 scanner
2296 .limit(limit, offset.map(|o| o as i64))
2297 .map_err(|e| {
2298 Error::invalid_input_source(
2299 format!("Invalid limit/offset for {}: {}", operation, e).into(),
2300 )
2301 })?;
2302 }
2303
2304 if let Some(true) = with_row_id {
2305 scanner.with_row_id();
2306 }
2307
2308 Ok(())
2309 }
2310
2311 pub fn retrieve_ops_metrics(&self) -> HashMap<String, u64> {
2318 self.ops_metrics
2319 .as_ref()
2320 .map(|m| m.retrieve())
2321 .unwrap_or_default()
2322 }
2323
2324 pub fn reset_ops_metrics(&self) {
2328 if let Some(ref metrics) = self.ops_metrics {
2329 metrics.reset();
2330 }
2331 }
2332
2333 fn record_op(&self, operation: &str) {
2335 if let Some(ref metrics) = self.ops_metrics {
2336 metrics.increment(operation);
2337 }
2338 }
2339}
2340
2341#[async_trait]
2342impl LanceNamespace for DirectoryNamespace {
2343 async fn list_namespaces(
2344 &self,
2345 request: ListNamespacesRequest,
2346 ) -> Result<ListNamespacesResponse> {
2347 self.record_op("list_namespaces");
2348 if let Some(ref manifest_ns) = self.manifest_ns {
2349 return manifest_ns.list_namespaces(request).await;
2350 }
2351
2352 Self::validate_root_namespace_id(&request.id)?;
2353 Ok(ListNamespacesResponse::new(vec![]))
2354 }
2355
2356 async fn describe_namespace(
2357 &self,
2358 request: DescribeNamespaceRequest,
2359 ) -> Result<DescribeNamespaceResponse> {
2360 self.record_op("describe_namespace");
2361 if let Some(ref manifest_ns) = self.manifest_ns {
2362 return manifest_ns.describe_namespace(request).await;
2363 }
2364
2365 Self::validate_root_namespace_id(&request.id)?;
2366 #[allow(clippy::needless_update)]
2367 Ok(DescribeNamespaceResponse {
2368 properties: Some(HashMap::new()),
2369 ..Default::default()
2370 })
2371 }
2372
2373 async fn create_namespace(
2374 &self,
2375 request: CreateNamespaceRequest,
2376 ) -> Result<CreateNamespaceResponse> {
2377 self.record_op("create_namespace");
2378 if let Some(ref manifest_ns) = self.manifest_ns {
2379 return manifest_ns.create_namespace(request).await;
2380 }
2381
2382 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
2383 return Err(NamespaceError::NamespaceAlreadyExists {
2384 message: "root namespace".to_string(),
2385 }
2386 .into());
2387 }
2388
2389 Err(NamespaceError::Unsupported {
2390 message: "Child namespaces are only supported when manifest mode is enabled"
2391 .to_string(),
2392 }
2393 .into())
2394 }
2395
2396 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
2397 self.record_op("drop_namespace");
2398 if let Some(ref manifest_ns) = self.manifest_ns {
2399 return manifest_ns.drop_namespace(request).await;
2400 }
2401
2402 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
2403 return Err(NamespaceError::InvalidInput {
2404 message: "Root namespace cannot be dropped".to_string(),
2405 }
2406 .into());
2407 }
2408
2409 Err(NamespaceError::Unsupported {
2410 message: "Child namespaces are only supported when manifest mode is enabled"
2411 .to_string(),
2412 }
2413 .into())
2414 }
2415
2416 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
2417 self.record_op("namespace_exists");
2418 if let Some(ref manifest_ns) = self.manifest_ns {
2419 return manifest_ns.namespace_exists(request).await;
2420 }
2421
2422 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
2423 return Ok(());
2424 }
2425
2426 Err(NamespaceError::NamespaceNotFound {
2427 message: "Child namespaces are only supported when manifest mode is enabled"
2428 .to_string(),
2429 }
2430 .into())
2431 }
2432
2433 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
2434 self.record_op("list_tables");
2435 let namespace_id = request.id.as_ref().ok_or_else(|| {
2437 lance_core::Error::from(NamespaceError::InvalidInput {
2438 message: "Namespace ID is required".to_string(),
2439 })
2440 })?;
2441
2442 if !namespace_id.is_empty() {
2444 if let Some(ref manifest_ns) = self.manifest_ns {
2445 return manifest_ns.list_tables(request).await;
2446 }
2447 return Err(NamespaceError::Unsupported {
2448 message: "Child namespaces are only supported when manifest mode is enabled"
2449 .to_string(),
2450 }
2451 .into());
2452 }
2453
2454 if let Some(ref manifest_ns) = self.manifest_ns
2456 && !self.dir_listing_enabled
2457 {
2458 return manifest_ns.list_tables(request).await;
2459 }
2460
2461 let mut tables = if self.manifest_ns.is_some()
2464 && self.dir_listing_enabled
2465 && self.dir_listing_to_manifest_migration_enabled
2466 {
2467 let manifest_locations = if let Some(ref manifest_ns) = self.manifest_ns {
2469 manifest_ns.list_manifest_table_locations().await?
2470 } else {
2471 std::collections::HashSet::new()
2472 };
2473
2474 let mut manifest_request = request.clone();
2476 manifest_request.limit = None;
2477 manifest_request.page_token = None;
2478 let manifest_tables = if let Some(ref manifest_ns) = self.manifest_ns {
2479 let manifest_response = manifest_ns.list_tables(manifest_request).await?;
2480 manifest_response.tables
2481 } else {
2482 vec![]
2483 };
2484
2485 let mut all_tables: Vec<String> = manifest_tables;
2488 let dir_tables = self.list_directory_tables().await?;
2489 for table_name in dir_tables {
2490 let full_location = format!("{}/{}.lance", self.root, table_name);
2493 let relative_location = format!("{}.lance", table_name);
2494 if !manifest_locations.contains(&full_location)
2495 && !manifest_locations.contains(&relative_location)
2496 {
2497 all_tables.push(table_name);
2498 }
2499 }
2500
2501 all_tables
2502 } else {
2503 self.list_directory_tables().await?
2504 };
2505
2506 tables = self
2507 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
2508 .await?;
2509
2510 let next_page_token =
2512 Self::apply_pagination(&mut tables, request.page_token, request.limit);
2513 let mut response = ListTablesResponse::new(tables);
2514 response.page_token = next_page_token;
2515 Ok(response)
2516 }
2517
2518 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
2519 self.record_op("describe_table");
2520 self.describe_table_impl(request).await
2521 }
2522
2523 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
2524 self.record_op("table_exists");
2525 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
2526 let skip_manifest_for_root = self.dir_listing_enabled
2527 && is_root_level
2528 && !self.dir_listing_to_manifest_migration_enabled;
2529 if let Some(ref manifest_ns) = self.manifest_ns
2530 && !skip_manifest_for_root
2531 {
2532 match manifest_ns.table_exists(request.clone()).await {
2533 Ok(()) => return Ok(()),
2534 Err(_) if self.dir_listing_enabled && is_root_level => {
2535 }
2537 Err(e) => return Err(e),
2538 }
2539 }
2540
2541 let table_name = Self::table_name_from_id(&request.id)?;
2542 let table_id = Self::format_table_id_from_request(&request.id);
2543
2544 let status = self.check_table_status(&table_name).await;
2546
2547 if !status.exists {
2548 return Err(NamespaceError::TableNotFound {
2549 message: table_id.clone(),
2550 }
2551 .into());
2552 }
2553
2554 if status.is_deregistered {
2555 return Err(NamespaceError::TableNotFound {
2556 message: format!("Table is deregistered: {}", table_id),
2557 }
2558 .into());
2559 }
2560
2561 Ok(())
2562 }
2563
2564 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
2565 self.record_op("drop_table");
2566 if let Some(ref manifest_ns) = self.manifest_ns {
2567 return manifest_ns.drop_table(request).await;
2568 }
2569
2570 let table_name = Self::table_name_from_id(&request.id)?;
2571 let table_uri = self.table_full_uri(&table_name);
2572 let table_path = self.table_path(&table_name);
2573
2574 self.object_store
2575 .remove_dir_all(table_path)
2576 .await
2577 .map_err(|e| {
2578 lance_core::Error::from(NamespaceError::Internal {
2579 message: format!("Failed to drop table {}: {:?}", table_name, e),
2580 })
2581 })?;
2582
2583 Ok(DropTableResponse {
2584 id: request.id,
2585 location: Some(table_uri),
2586 ..Default::default()
2587 })
2588 }
2589
2590 async fn create_table(
2591 &self,
2592 request: CreateTableRequest,
2593 request_data: Bytes,
2594 ) -> Result<CreateTableResponse> {
2595 self.record_op("create_table");
2596 if let Some(ref manifest_ns) = self.manifest_ns {
2597 return manifest_ns.create_table(request, request_data).await;
2598 }
2599
2600 Self::validate_dir_only_properties(request.properties.as_ref(), "create_table")?;
2601
2602 let table_name = Self::table_name_from_id(&request.id)?;
2603 let table_uri = self.table_full_uri(&table_name);
2604 let status = self.check_table_status(&table_name).await;
2605 let (reader, _num_rows) =
2606 Self::ipc_reader_from_request_data(&request_data, "create_table")?;
2607
2608 if status.exists && self.table_has_actual_manifests(&table_name).await? {
2609 return Err(NamespaceError::TableAlreadyExists {
2610 message: table_name,
2611 }
2612 .into());
2613 }
2614
2615 let write_result = self
2616 .write_reader_to_table(
2617 &table_uri,
2618 reader,
2619 WriteMode::Create,
2620 request.storage_options.clone(),
2621 )
2622 .await;
2623 if let Err(err) = write_result {
2624 if self.table_uri_has_actual_manifests(&table_uri).await? {
2625 return Err(NamespaceError::TableAlreadyExists {
2626 message: table_name,
2627 }
2628 .into());
2629 }
2630 return Err(err);
2631 }
2632 Ok(CreateTableResponse {
2633 version: Some(1),
2634 location: Some(table_uri),
2635 storage_options: self.storage_options.clone(),
2636 properties: request.properties,
2637 ..Default::default()
2638 })
2639 }
2640
2641 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
2642 self.record_op("declare_table");
2643 if let Some(ref manifest_ns) = self.manifest_ns {
2644 let mut response = manifest_ns.declare_table(request.clone()).await?;
2645 if let Some(ref location) = response.location {
2646 let vend = request.vend_credentials.unwrap_or(true);
2648 let identity = request.identity.as_deref();
2649 response.storage_options = self
2650 .get_storage_options_for_table(location, vend, identity)
2651 .await?;
2652 }
2653 if self.table_version_tracking_enabled {
2655 response.managed_versioning = Some(true);
2656 }
2657 return Ok(response);
2658 }
2659
2660 Self::validate_dir_only_properties(request.properties.as_ref(), "declare_table")?;
2661
2662 let table_name = Self::table_name_from_id(&request.id)?;
2663 let table_uri = self.table_full_uri(&table_name);
2664
2665 if let Some(location) = &request.location {
2667 let location = location.trim_end_matches('/');
2668 if location != table_uri {
2669 return Err(NamespaceError::InvalidInput {
2670 message: format!(
2671 "Cannot declare table {} at location {}, must be at location {}",
2672 table_name, location, table_uri
2673 ),
2674 }
2675 .into());
2676 }
2677 }
2678
2679 let status = self.check_table_status(&table_name).await;
2683 if status.exists && !status.has_reserved_file {
2684 return Err(NamespaceError::TableAlreadyExists {
2686 message: table_name.to_string(),
2687 }
2688 .into());
2689 }
2690
2691 let reserved_file_path = self.table_reserved_file_path(&table_name);
2695
2696 self.put_marker_file_atomic(&reserved_file_path, &format!("table {}", table_name))
2697 .await
2698 .map_err(|e| {
2699 if e.contains("already exists") {
2700 lance_core::Error::from(NamespaceError::TableAlreadyExists {
2701 message: table_name.to_string(),
2702 })
2703 } else {
2704 lance_core::Error::from(NamespaceError::Internal { message: e })
2705 }
2706 })?;
2707
2708 let vend_credentials = request.vend_credentials.unwrap_or(true);
2710 let identity = request.identity.as_deref();
2711 let storage_options = self
2712 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2713 .await?;
2714
2715 Ok(DeclareTableResponse {
2716 location: Some(table_uri),
2717 storage_options,
2718 properties: request.properties,
2719 managed_versioning: if self.table_version_tracking_enabled {
2720 Some(true)
2721 } else {
2722 None
2723 },
2724 ..Default::default()
2725 })
2726 }
2727
2728 async fn register_table(
2729 &self,
2730 request: lance_namespace::models::RegisterTableRequest,
2731 ) -> Result<lance_namespace::models::RegisterTableResponse> {
2732 self.record_op("register_table");
2733 if let Some(ref manifest_ns) = self.manifest_ns {
2735 return LanceNamespace::register_table(manifest_ns.as_ref(), request).await;
2736 }
2737
2738 Err(NamespaceError::Unsupported {
2740 message: "register_table is only supported when manifest mode is enabled".to_string(),
2741 }
2742 .into())
2743 }
2744
2745 async fn deregister_table(
2746 &self,
2747 request: lance_namespace::models::DeregisterTableRequest,
2748 ) -> Result<lance_namespace::models::DeregisterTableResponse> {
2749 self.record_op("deregister_table");
2750 if let Some(ref manifest_ns) = self.manifest_ns {
2752 return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await;
2753 }
2754
2755 let table_name = Self::table_name_from_id(&request.id)?;
2757 let table_uri = self.table_full_uri(&table_name);
2758
2759 let status = self.check_table_status(&table_name).await;
2762
2763 if !status.exists {
2764 return Err(NamespaceError::TableNotFound {
2765 message: table_name.to_string(),
2766 }
2767 .into());
2768 }
2769
2770 if status.is_deregistered {
2771 return Err(NamespaceError::TableNotFound {
2772 message: format!("Table is already deregistered: {}", table_name),
2773 }
2774 .into());
2775 }
2776
2777 let deregistered_path = self.table_deregistered_file_path(&table_name);
2783 self.put_marker_file_atomic(
2784 &deregistered_path,
2785 &format!("deregistration marker for table {}", table_name),
2786 )
2787 .await
2788 .map_err(|e| {
2789 if e.contains("already exists") {
2790 lance_core::Error::from(NamespaceError::InvalidTableState {
2791 message: format!("Table is already deregistered: {}", table_name),
2792 })
2793 } else {
2794 lance_core::Error::from(NamespaceError::Internal { message: e })
2795 }
2796 })?;
2797
2798 Ok(lance_namespace::models::DeregisterTableResponse {
2799 id: request.id,
2800 location: Some(table_uri),
2801 ..Default::default()
2802 })
2803 }
2804
2805 async fn list_table_versions(
2806 &self,
2807 request: ListTableVersionsRequest,
2808 ) -> Result<ListTableVersionsResponse> {
2809 self.record_op("list_table_versions");
2810 if self.table_version_storage_enabled
2812 && let Some(ref manifest_ns) = self.manifest_ns
2813 {
2814 let table_id = request.id.clone().unwrap_or_default();
2815 let want_descending = request.descending == Some(true);
2816 return manifest_ns
2817 .list_table_versions(&table_id, want_descending, request.limit)
2818 .await;
2819 }
2820
2821 let table_uri = self.resolve_table_location(&request.id).await?;
2823 let want_descending = request.descending == Some(true);
2824 let table_versions = self
2825 .list_table_versions_from_storage(&table_uri, want_descending, request.limit)
2826 .await?;
2827
2828 Ok(ListTableVersionsResponse {
2829 versions: table_versions,
2830 page_token: None,
2831 })
2832 }
2833
2834 async fn create_table_version(
2835 &self,
2836 request: CreateTableVersionRequest,
2837 ) -> Result<CreateTableVersionResponse> {
2838 self.record_op("create_table_version");
2839 let table_uri = self.resolve_table_location(&request.id).await?;
2840
2841 let staging_manifest_path = &request.manifest_path;
2842 let version = request.version as u64;
2843
2844 let table_path = self.object_store_path_from_uri(&table_uri)?;
2845
2846 let naming_scheme = match request.naming_scheme.as_deref() {
2848 Some("V1") => ManifestNamingScheme::V1,
2849 _ => ManifestNamingScheme::V2,
2850 };
2851
2852 let final_path = naming_scheme.manifest_path(&table_path, version);
2854
2855 let staging_path = Path::parse(staging_manifest_path).map_err(|e| {
2856 lance_core::Error::from(NamespaceError::InvalidInput {
2857 message: format!(
2858 "Invalid staging manifest path '{}': {}",
2859 staging_manifest_path, e
2860 ),
2861 })
2862 })?;
2863
2864 let copy_result = match self
2865 .object_store
2866 .inner
2867 .copy_if_not_exists(&staging_path, &final_path)
2868 .await
2869 {
2870 Ok(()) => Ok(()),
2871 Err(ObjectStoreError::NotImplemented { .. })
2872 | Err(ObjectStoreError::NotSupported { .. }) => {
2873 let manifest_data = self
2874 .object_store
2875 .inner
2876 .get(&staging_path)
2877 .await
2878 .map_err(|e| {
2879 lance_core::Error::from(NamespaceError::Internal {
2880 message: format!(
2881 "Failed to read staging manifest at '{}': {}",
2882 staging_manifest_path, e
2883 ),
2884 })
2885 })?
2886 .bytes()
2887 .await
2888 .map_err(|e| {
2889 lance_core::Error::from(NamespaceError::Internal {
2890 message: format!(
2891 "Failed to read staging manifest bytes at '{}': {}",
2892 staging_manifest_path, e
2893 ),
2894 })
2895 })?;
2896 self.object_store
2897 .inner
2898 .put_opts(
2899 &final_path,
2900 manifest_data.into(),
2901 PutOptions {
2902 mode: PutMode::Create,
2903 ..Default::default()
2904 },
2905 )
2906 .await
2907 .map(|_| ())
2908 }
2909 Err(e) => Err(e),
2910 };
2911
2912 match copy_result {
2913 Ok(()) => {}
2914 Err(ObjectStoreError::AlreadyExists { .. })
2915 | Err(ObjectStoreError::Precondition { .. }) => {
2916 return Err(lance_core::Error::from(
2917 NamespaceError::ConcurrentModification {
2918 message: format!(
2919 "Version {} already exists for table at '{}'",
2920 version, table_uri
2921 ),
2922 },
2923 ));
2924 }
2925 Err(e) => {
2926 return Err(lance_core::Error::from(NamespaceError::Internal {
2927 message: format!(
2928 "Failed to create version {} for table at '{}': {}",
2929 version, table_uri, e
2930 ),
2931 }));
2932 }
2933 }
2934
2935 let final_meta = self
2936 .object_store
2937 .inner
2938 .head(&final_path)
2939 .await
2940 .map_err(|e| {
2941 lance_core::Error::from(NamespaceError::Internal {
2942 message: format!(
2943 "Failed to stat created version {} for table at '{}': {}",
2944 version, table_uri, e
2945 ),
2946 })
2947 })?;
2948 let manifest_size = final_meta.size as i64;
2949
2950 if let Err(e) = self.object_store.inner.delete(&staging_path).await {
2952 log::warn!(
2953 "Failed to delete staging manifest at '{}': {:?}",
2954 staging_path,
2955 e
2956 );
2957 }
2958
2959 if self.table_version_storage_enabled
2961 && let Some(ref manifest_ns) = self.manifest_ns
2962 {
2963 let table_id_str =
2964 manifest::ManifestNamespace::str_object_id(&request.id.clone().unwrap_or_default());
2965 let object_id =
2966 manifest::ManifestNamespace::build_version_object_id(&table_id_str, version as i64);
2967 let metadata_json = serde_json::json!({
2968 "manifest_path": final_path.to_string(),
2969 "manifest_size": manifest_size,
2970 "e_tag": final_meta.e_tag,
2971 "naming_scheme": request.naming_scheme.as_deref().unwrap_or("V2"),
2972 })
2973 .to_string();
2974
2975 if let Err(e) = manifest_ns
2976 .insert_into_manifest_with_metadata(
2977 vec![manifest::ManifestEntry {
2978 object_id,
2979 object_type: manifest::ObjectType::TableVersion,
2980 location: None,
2981 metadata: Some(metadata_json),
2982 }],
2983 None,
2984 )
2985 .await
2986 {
2987 log::warn!(
2988 "Failed to record table version in __manifest (best-effort): {:?}",
2989 e
2990 );
2991 }
2992 }
2993
2994 Ok(CreateTableVersionResponse {
2995 transaction_id: None,
2996 version: Some(Box::new(TableVersion {
2997 version: version as i64,
2998 manifest_path: final_path.to_string(),
2999 manifest_size: Some(manifest_size),
3000 e_tag: final_meta.e_tag,
3001 timestamp_millis: None,
3002 metadata: None,
3003 })),
3004 })
3005 }
3006
3007 async fn describe_table_version(
3008 &self,
3009 request: DescribeTableVersionRequest,
3010 ) -> Result<DescribeTableVersionResponse> {
3011 self.record_op("describe_table_version");
3012 if self.table_version_storage_enabled
3015 && let (Some(manifest_ns), Some(version)) = (&self.manifest_ns, request.version)
3016 {
3017 let table_id = request.id.clone().unwrap_or_default();
3018 return manifest_ns.describe_table_version(&table_id, version).await;
3019 }
3020
3021 let table_uri = self.resolve_table_location(&request.id).await?;
3023 let versions = self
3024 .list_table_versions_from_storage(&table_uri, true, None)
3025 .await?;
3026 let table_version = if let Some(requested_version) = request.version {
3027 versions
3028 .into_iter()
3029 .find(|version| version.version == requested_version)
3030 .ok_or_else(|| {
3031 lance_core::Error::from(NamespaceError::TableVersionNotFound {
3032 message: format!(
3033 "version {} for table {}",
3034 requested_version,
3035 Self::format_table_id_from_request(&request.id)
3036 ),
3037 })
3038 })?
3039 } else {
3040 versions.into_iter().next().ok_or_else(|| {
3041 lance_core::Error::from(NamespaceError::TableVersionNotFound {
3042 message: format!(
3043 "latest version for table {}",
3044 Self::format_table_id_from_request(&request.id)
3045 ),
3046 })
3047 })?
3048 };
3049
3050 Ok(DescribeTableVersionResponse {
3051 version: Box::new(table_version),
3052 })
3053 }
3054
3055 async fn batch_delete_table_versions(
3056 &self,
3057 request: BatchDeleteTableVersionsRequest,
3058 ) -> Result<BatchDeleteTableVersionsResponse> {
3059 self.record_op("batch_delete_table_versions");
3060 let ranges: Vec<(i64, i64)> = request
3063 .ranges
3064 .iter()
3065 .map(|r| {
3066 let start = r.start_version;
3067 let end = if r.end_version > 0 {
3068 r.end_version
3069 } else {
3070 start
3071 };
3072 (start, end)
3073 })
3074 .collect();
3075 let table_entries = vec![TableDeleteEntry {
3076 table_id: request.id.clone(),
3077 ranges,
3078 }];
3079
3080 let mut total_deleted_count = 0i64;
3081
3082 if self.table_version_storage_enabled
3083 && let Some(ref manifest_ns) = self.manifest_ns
3084 {
3085 let mut all_object_ids: Vec<String> = Vec::new();
3092 for te in &table_entries {
3093 let table_id_str = manifest::ManifestNamespace::str_object_id(
3094 &te.table_id.clone().unwrap_or_default(),
3095 );
3096 for (start, end) in &te.ranges {
3097 for version in *start..=*end {
3098 let object_id = manifest::ManifestNamespace::build_version_object_id(
3099 &table_id_str,
3100 version,
3101 );
3102 all_object_ids.push(object_id);
3103 }
3104 }
3105 }
3106
3107 if !all_object_ids.is_empty() {
3108 total_deleted_count = manifest_ns
3109 .batch_delete_table_versions_by_object_ids(&all_object_ids)
3110 .await?;
3111 }
3112
3113 let _ = self
3118 .delete_physical_version_files(&table_entries, true)
3119 .await;
3120
3121 return Ok(BatchDeleteTableVersionsResponse {
3122 deleted_count: Some(total_deleted_count),
3123 transaction_id: None,
3124 });
3125 }
3126
3127 total_deleted_count = self
3129 .delete_physical_version_files(&table_entries, false)
3130 .await?;
3131
3132 Ok(BatchDeleteTableVersionsResponse {
3133 deleted_count: Some(total_deleted_count),
3134 transaction_id: None,
3135 })
3136 }
3137
3138 async fn create_table_index(
3139 &self,
3140 request: CreateTableIndexRequest,
3141 ) -> Result<CreateTableIndexResponse> {
3142 self.record_op("create_table_index");
3143 let table_uri = self.resolve_table_location(&request.id).await?;
3144 let mut dataset = self
3145 .load_dataset(&table_uri, None, "create_table_index")
3146 .await?;
3147 let index_request = Self::build_index_params(&request)?;
3148
3149 dataset
3150 .create_index(
3151 &[request.column.as_str()],
3152 index_request.index_type(),
3153 request.name.clone(),
3154 index_request.params(),
3155 false,
3156 )
3157 .await
3158 .map_err(|e| {
3159 let err_msg = format!("{}", e);
3160 let ns_err = if err_msg.contains("already exists") {
3161 NamespaceError::TableIndexAlreadyExists {
3162 message: format!(
3163 "Index '{}' already exists on table '{}': {:?}",
3164 request.name.as_deref().unwrap_or("<auto-generated>"),
3165 table_uri,
3166 e
3167 ),
3168 }
3169 } else if err_msg.contains("not found") || err_msg.contains("does not exist") {
3170 NamespaceError::TableColumnNotFound {
3171 message: format!(
3172 "Column '{}' not found for table '{}': {:?}",
3173 request.column, table_uri, e
3174 ),
3175 }
3176 } else {
3177 NamespaceError::Internal {
3178 message: format!(
3179 "Failed to create {} index '{}' on column '{}' for table '{}': {:?}",
3180 request.index_type,
3181 request.name.as_deref().unwrap_or("<auto-generated>"),
3182 request.column,
3183 table_uri,
3184 e
3185 ),
3186 }
3187 };
3188 lance_core::Error::from(ns_err)
3189 })?;
3190
3191 let transaction_id = dataset
3192 .read_transaction()
3193 .await
3194 .map_err(|e| {
3195 lance_core::Error::from(NamespaceError::Internal {
3196 message: format!(
3197 "Failed to read committed transaction after creating index on '{}': {}",
3198 table_uri, e
3199 ),
3200 })
3201 })?
3202 .map(|transaction| transaction.uuid);
3203
3204 Ok(CreateTableIndexResponse { transaction_id })
3205 }
3206
3207 async fn list_table_indices(
3208 &self,
3209 request: ListTableIndicesRequest,
3210 ) -> Result<ListTableIndicesResponse> {
3211 self.record_op("list_table_indices");
3212 let table_uri = self.resolve_table_location(&request.id).await?;
3213 let dataset = self
3214 .load_dataset(&table_uri, request.version, "list_table_indices")
3215 .await?;
3216 let mut indices = dataset
3217 .describe_indices(None)
3218 .await
3219 .map_err(|e| {
3220 lance_core::Error::from(NamespaceError::Internal {
3221 message: format!("Failed to describe table indices for '{}': {:?}", table_uri, e),
3222 })
3223 })?
3224 .into_iter()
3225 .filter(|description| {
3226 description
3227 .metadata()
3228 .first()
3229 .map(|metadata| !is_system_index(metadata))
3230 .unwrap_or(false)
3231 })
3232 .map(|description| {
3233 let columns = description
3234 .field_ids()
3235 .iter()
3236 .map(|field_id| {
3237 dataset
3238 .schema()
3239 .field_path(i32::try_from(*field_id).map_err(|e| {
3240 lance_core::Error::from(NamespaceError::Internal {
3241 message: format!(
3242 "Field id {} does not fit in i32 for table '{}': {}",
3243 field_id, table_uri, e
3244 ),
3245 })
3246 })?)
3247 .map_err(|e| {
3248 lance_core::Error::from(NamespaceError::Internal {
3249 message: format!(
3250 "Failed to resolve field path for field_id {} in table '{}': {}",
3251 field_id, table_uri, e
3252 ),
3253 })
3254 })
3255 })
3256 .collect::<Result<Vec<_>>>()?;
3257
3258 Ok(IndexContent {
3259 index_name: description.name().to_string(),
3260 index_uuid: description.metadata()[0].uuid.to_string(),
3261 columns,
3262 status: "SUCCEEDED".to_string(),
3263 })
3264 })
3265 .collect::<Result<Vec<_>>>()?;
3266
3267 let page_token = Self::paginate_indices(&mut indices, request.page_token, request.limit);
3268 Ok(ListTableIndicesResponse {
3269 indexes: indices,
3270 page_token,
3271 })
3272 }
3273
3274 async fn describe_table_index_stats(
3275 &self,
3276 request: DescribeTableIndexStatsRequest,
3277 ) -> Result<DescribeTableIndexStatsResponse> {
3278 self.record_op("describe_table_index_stats");
3279 let table_uri = self.resolve_table_location(&request.id).await?;
3280 let dataset = self
3281 .load_dataset(&table_uri, request.version, "describe_table_index_stats")
3282 .await?;
3283 let index_name = request.index_name.as_deref().ok_or_else(|| {
3284 lance_core::Error::from(NamespaceError::InvalidInput {
3285 message: "Index name is required for describe_table_index_stats".to_string(),
3286 })
3287 })?;
3288 let metadatas = dataset
3289 .load_indices_by_name(index_name)
3290 .await
3291 .map_err(|e| {
3292 lance_core::Error::from(NamespaceError::TableIndexNotFound {
3293 message: format!(
3294 "Failed to load index '{}' metadata for table '{}': {}",
3295 index_name, table_uri, e
3296 ),
3297 })
3298 })?;
3299 if metadatas.first().is_some_and(is_system_index) {
3300 return Err(NamespaceError::Unsupported {
3301 message: format!("System index '{}' is not exposed by this API", index_name),
3302 }
3303 .into());
3304 }
3305
3306 let stats = <Dataset as DatasetIndexExt>::index_statistics(&dataset, index_name)
3307 .await
3308 .map_err(|e| {
3309 lance_core::Error::from(NamespaceError::TableIndexNotFound {
3310 message: format!(
3311 "Failed to describe index statistics for '{}' on table '{}': {}",
3312 index_name, table_uri, e
3313 ),
3314 })
3315 })?;
3316 let stats: serde_json::Value = serde_json::from_str(&stats).map_err(|e| {
3317 lance_core::Error::from(NamespaceError::Internal {
3318 message: format!(
3319 "Failed to parse index statistics for '{}' on table '{}': {}",
3320 index_name, table_uri, e
3321 ),
3322 })
3323 })?;
3324
3325 Ok(Self::describe_table_index_stats_response(&stats))
3326 }
3327
3328 async fn describe_transaction(
3329 &self,
3330 request: DescribeTransactionRequest,
3331 ) -> Result<DescribeTransactionResponse> {
3332 self.record_op("describe_transaction");
3333 let mut request_id = request.id.ok_or_else(|| {
3334 lance_core::Error::from(NamespaceError::InvalidInput {
3335 message: "Transaction id must include table id and transaction identifier"
3336 .to_string(),
3337 })
3338 })?;
3339 if request_id.len() < 2 {
3340 return Err(NamespaceError::InvalidInput {
3341 message: format!(
3342 "Transaction request id must include table id and transaction identifier, got {:?}",
3343 request_id
3344 ),
3345 }
3346 .into());
3347 }
3348
3349 let id = request_id.pop().expect("request_id len checked above");
3350 let table_id = Some(request_id);
3351 let table_uri = self.resolve_table_location(&table_id).await?;
3352 let dataset = self
3353 .load_dataset(&table_uri, None, "describe_transaction")
3354 .await?;
3355 let (version, transaction) = self.find_transaction(&dataset, &id).await?;
3356
3357 Ok(Self::transaction_response(version, &transaction))
3358 }
3359
3360 async fn create_table_scalar_index(
3361 &self,
3362 request: CreateTableIndexRequest,
3363 ) -> Result<CreateTableScalarIndexResponse> {
3364 self.record_op("create_table_scalar_index");
3365 let index_type = Self::parse_index_type(&request.index_type)?;
3366 if !index_type.is_scalar() {
3367 return Err(NamespaceError::InvalidInput {
3368 message: format!(
3369 "create_table_scalar_index only supports scalar index types, got {}",
3370 request.index_type
3371 ),
3372 }
3373 .into());
3374 }
3375
3376 let response = self.create_table_index(request).await?;
3377 Ok(CreateTableScalarIndexResponse {
3378 transaction_id: response.transaction_id,
3379 })
3380 }
3381
3382 async fn drop_table_index(
3383 &self,
3384 request: DropTableIndexRequest,
3385 ) -> Result<DropTableIndexResponse> {
3386 self.record_op("drop_table_index");
3387 let table_uri = self.resolve_table_location(&request.id).await?;
3388 let index_name = request.index_name.as_deref().ok_or_else(|| {
3389 lance_core::Error::from(NamespaceError::InvalidInput {
3390 message: "Index name is required for drop_table_index".to_string(),
3391 })
3392 })?;
3393 let mut dataset = self
3394 .load_dataset(&table_uri, None, "drop_table_index")
3395 .await?;
3396 let metadatas = dataset
3397 .load_indices_by_name(index_name)
3398 .await
3399 .map_err(|e| {
3400 lance_core::Error::from(NamespaceError::TableIndexNotFound {
3401 message: format!(
3402 "Failed to load index '{}' before dropping it from table '{}': {}",
3403 index_name, table_uri, e
3404 ),
3405 })
3406 })?;
3407 if metadatas.first().is_some_and(is_system_index) {
3408 return Err(NamespaceError::Unsupported {
3409 message: format!(
3410 "System index '{}' cannot be dropped via this API",
3411 index_name
3412 ),
3413 }
3414 .into());
3415 }
3416
3417 dataset.drop_index(index_name).await.map_err(|e| {
3418 lance_core::Error::from(NamespaceError::TableIndexNotFound {
3419 message: format!(
3420 "Failed to drop index '{}' from table '{}': {}",
3421 index_name, table_uri, e
3422 ),
3423 })
3424 })?;
3425
3426 let transaction_id = dataset
3427 .read_transaction()
3428 .await
3429 .map_err(|e| {
3430 lance_core::Error::from(NamespaceError::Internal {
3431 message: format!(
3432 "Failed to read committed transaction after dropping index '{}' from '{}': {}",
3433 index_name, table_uri, e
3434 ),
3435 })
3436 })?
3437 .map(|transaction| transaction.uuid);
3438
3439 Ok(DropTableIndexResponse { transaction_id })
3440 }
3441
3442 async fn list_all_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
3443 let mut tables = self.list_directory_tables().await?;
3446 tables = self
3447 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
3448 .await?;
3449 Self::apply_pagination(&mut tables, request.page_token, request.limit);
3450 Ok(ListTablesResponse::new(tables))
3451 }
3452
3453 async fn restore_table(&self, request: RestoreTableRequest) -> Result<RestoreTableResponse> {
3454 let version = request.version;
3455 if version < 0 {
3456 return Err(Error::invalid_input_source(
3457 format!(
3458 "Table version for restore_table must be non-negative, got {}",
3459 version
3460 )
3461 .into(),
3462 ));
3463 }
3464
3465 let table_uri = self.resolve_table_location(&request.id).await?;
3466 let mut dataset = self.load_dataset(&table_uri, None, "restore_table").await?;
3467
3468 dataset = dataset
3469 .checkout_version(version as u64)
3470 .await
3471 .map_err(|e| {
3472 Error::namespace_source(
3473 format!(
3474 "Failed to checkout version {} for restore at '{}': {}",
3475 version, table_uri, e
3476 )
3477 .into(),
3478 )
3479 })?;
3480
3481 dataset.restore().await.map_err(|e| {
3482 Error::namespace_source(
3483 format!(
3484 "Failed to restore table at '{}' to version {}: {}",
3485 table_uri, version, e
3486 )
3487 .into(),
3488 )
3489 })?;
3490
3491 let transaction_id = dataset
3492 .read_transaction()
3493 .await
3494 .map_err(|e| {
3495 Error::namespace_source(
3496 format!(
3497 "Failed to read transaction after restoring '{}': {}",
3498 table_uri, e
3499 )
3500 .into(),
3501 )
3502 })?
3503 .map(|t| t.uuid);
3504
3505 Ok(RestoreTableResponse { transaction_id })
3506 }
3507
3508 async fn update_table_schema_metadata(
3509 &self,
3510 request: UpdateTableSchemaMetadataRequest,
3511 ) -> Result<UpdateTableSchemaMetadataResponse> {
3512 let table_uri = self.resolve_table_location(&request.id).await?;
3513 let mut dataset = self
3514 .load_dataset(&table_uri, None, "update_table_schema_metadata")
3515 .await?;
3516
3517 let new_metadata = request.metadata.unwrap_or_default();
3518 let updated_metadata = dataset
3519 .update_schema_metadata(new_metadata.iter().map(|(k, v)| (k.as_str(), v.as_str())))
3520 .await
3521 .map_err(|e| {
3522 Error::namespace_source(
3523 format!(
3524 "Failed to update schema metadata for table at '{}': {}",
3525 table_uri, e
3526 )
3527 .into(),
3528 )
3529 })?;
3530
3531 let transaction_id = dataset
3532 .read_transaction()
3533 .await
3534 .map_err(|e| {
3535 Error::namespace_source(
3536 format!(
3537 "Failed to read transaction after updating metadata for '{}': {}",
3538 table_uri, e
3539 )
3540 .into(),
3541 )
3542 })?
3543 .map(|t| t.uuid);
3544
3545 Ok(UpdateTableSchemaMetadataResponse {
3546 metadata: Some(updated_metadata),
3547 transaction_id,
3548 })
3549 }
3550
3551 async fn get_table_stats(
3552 &self,
3553 request: GetTableStatsRequest,
3554 ) -> Result<GetTableStatsResponse> {
3555 let table_uri = self.resolve_table_location(&request.id).await?;
3556 let dataset = Arc::new(
3557 self.load_dataset(&table_uri, None, "get_table_stats")
3558 .await?,
3559 );
3560
3561 let data_stats = dataset.calculate_data_stats().await.map_err(|e| {
3563 Error::namespace_source(
3564 format!(
3565 "Failed to calculate data statistics for table at '{}': {}",
3566 table_uri, e
3567 )
3568 .into(),
3569 )
3570 })?;
3571 let total_bytes: i64 = data_stats
3572 .fields
3573 .iter()
3574 .map(|f| f.bytes_on_disk as i64)
3575 .sum();
3576
3577 let fragment_row_futures: Vec<_> = dataset
3579 .get_fragments()
3580 .into_iter()
3581 .map(|f| async move { f.physical_rows().await })
3582 .collect();
3583 let fragment_row_results = futures::future::join_all(fragment_row_futures).await;
3584 let mut fragment_row_counts: Vec<i64> = fragment_row_results
3585 .into_iter()
3586 .filter_map(|r| r.ok())
3587 .map(|r| r as i64)
3588 .collect();
3589
3590 let num_fragments = fragment_row_counts.len() as i64;
3591 let num_rows: i64 = fragment_row_counts.iter().sum();
3592
3593 const SMALL_FRAGMENT_THRESHOLD: i64 = 1024 * 1024;
3596 let num_small_fragments = fragment_row_counts
3597 .iter()
3598 .filter(|&&r| r < SMALL_FRAGMENT_THRESHOLD)
3599 .count() as i64;
3600
3601 fragment_row_counts.sort_unstable();
3603 let lengths = if fragment_row_counts.is_empty() {
3604 FragmentSummary::new(0, 0, 0, 0, 0, 0, 0)
3605 } else {
3606 let len = fragment_row_counts.len();
3607 let min = fragment_row_counts[0];
3608 let max = fragment_row_counts[len - 1];
3609 let mean = num_rows / num_fragments;
3610 let pct = |p: f64| fragment_row_counts[((len - 1) as f64 * p) as usize];
3611 FragmentSummary::new(min, max, mean, pct(0.25), pct(0.50), pct(0.75), pct(0.99))
3612 };
3613
3614 let indices = dataset.load_indices().await.map_err(|e| {
3616 Error::namespace_source(
3617 format!("Failed to load indices for table at '{}': {}", table_uri, e).into(),
3618 )
3619 })?;
3620 let num_indices = indices.iter().filter(|m| !is_system_index(m)).count() as i64;
3621
3622 let fragment_stats = FragmentStats::new(num_fragments, num_small_fragments, lengths);
3623 Ok(GetTableStatsResponse::new(
3624 total_bytes,
3625 num_rows,
3626 num_indices,
3627 fragment_stats,
3628 ))
3629 }
3630
3631 async fn explain_table_query_plan(
3632 &self,
3633 request: ExplainTableQueryPlanRequest,
3634 ) -> Result<String> {
3635 let table_uri = self.resolve_table_location(&request.id).await?;
3636 let dataset = self
3637 .load_dataset(
3638 &table_uri,
3639 request.query.version,
3640 "explain_table_query_plan",
3641 )
3642 .await?;
3643 let verbose = request.verbose.unwrap_or(false);
3644
3645 let mut scanner = dataset.scan();
3646 Self::apply_query_params_to_scanner(
3647 &mut scanner,
3648 request.query.filter.as_deref(),
3649 request.query.columns.as_deref(),
3650 request.query.vector_column.as_deref(),
3651 &request.query.vector,
3652 request.query.k,
3653 request.query.offset,
3654 request.query.prefilter,
3655 request.query.bypass_vector_index,
3656 request.query.nprobes,
3657 request.query.ef,
3658 request.query.refine_factor,
3659 request.query.distance_type.as_deref(),
3660 request.query.fast_search,
3661 request.query.with_row_id,
3662 request.query.lower_bound,
3663 request.query.upper_bound,
3664 "explain_table_query_plan",
3665 )?;
3666
3667 scanner.explain_plan(verbose).await.map_err(|e| {
3668 Error::namespace_source(
3669 format!(
3670 "Failed to explain query plan for table at '{}': {}",
3671 table_uri, e
3672 )
3673 .into(),
3674 )
3675 })
3676 }
3677
3678 async fn analyze_table_query_plan(
3679 &self,
3680 request: AnalyzeTableQueryPlanRequest,
3681 ) -> Result<String> {
3682 let table_uri = self.resolve_table_location(&request.id).await?;
3683 let dataset = self
3684 .load_dataset(&table_uri, request.version, "analyze_table_query_plan")
3685 .await?;
3686
3687 let mut scanner = dataset.scan();
3688 Self::apply_query_params_to_scanner(
3689 &mut scanner,
3690 request.filter.as_deref(),
3691 request.columns.as_deref(),
3692 request.vector_column.as_deref(),
3693 &request.vector,
3694 request.k,
3695 request.offset,
3696 request.prefilter,
3697 request.bypass_vector_index,
3698 request.nprobes,
3699 request.ef,
3700 request.refine_factor,
3701 request.distance_type.as_deref(),
3702 request.fast_search,
3703 request.with_row_id,
3704 request.lower_bound,
3705 request.upper_bound,
3706 "analyze_table_query_plan",
3707 )?;
3708
3709 scanner.analyze_plan().await.map_err(|e| {
3710 Error::namespace_source(
3711 format!(
3712 "Failed to analyze query plan for table at '{}': {}",
3713 table_uri, e
3714 )
3715 .into(),
3716 )
3717 })
3718 }
3719
3720 async fn count_table_rows(&self, request: CountTableRowsRequest) -> Result<i64> {
3721 self.record_op("count_table_rows");
3722 let table_uri = self.resolve_table_location(&request.id).await?;
3723 let dataset = self
3724 .load_dataset(&table_uri, request.version, "count_table_rows")
3725 .await?;
3726
3727 let count =
3728 dataset
3729 .count_rows(request.predicate)
3730 .await
3731 .map_err(|e| NamespaceError::Internal {
3732 message: format!("Failed to count rows for table at '{}': {:?}", table_uri, e),
3733 })?;
3734
3735 Ok(count as i64)
3736 }
3737
3738 async fn insert_into_table(
3739 &self,
3740 request: InsertIntoTableRequest,
3741 request_data: Bytes,
3742 ) -> Result<InsertIntoTableResponse> {
3743 self.record_op("insert_into_table");
3744 let table_uri = self.resolve_table_location(&request.id).await?;
3745 let (reader, _num_rows) =
3746 Self::ipc_reader_from_request_data(&request_data, "insert_into_table")?;
3747
3748 let mode = match request.mode.as_deref() {
3749 Some(m) if m.eq_ignore_ascii_case("overwrite") => WriteMode::Overwrite,
3750 Some(m) if m.eq_ignore_ascii_case("append") => WriteMode::Append,
3751 None => WriteMode::Append,
3752 Some(m) => {
3753 return Err(lance_namespace::error::NamespaceError::InvalidInput {
3754 message: format!(
3755 "Unsupported write mode '{}'. Supported modes are: 'append', 'overwrite'",
3756 m
3757 ),
3758 }
3759 .into());
3760 }
3761 };
3762
3763 if !self.table_uri_has_actual_manifests(&table_uri).await? {
3764 self.write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
3765 .await?;
3766 } else {
3767 self.write_reader_to_table(&table_uri, reader, mode, None)
3768 .await?;
3769 }
3770
3771 Ok(InsertIntoTableResponse {
3772 transaction_id: None,
3773 })
3774 }
3775
3776 async fn merge_insert_into_table(
3777 &self,
3778 request: MergeInsertIntoTableRequest,
3779 request_data: Bytes,
3780 ) -> Result<MergeInsertIntoTableResponse> {
3781 self.record_op("merge_insert_into_table");
3782 let table_uri = self.resolve_table_location(&request.id).await?;
3783 let on = request.on.as_ref().ok_or_else(|| {
3784 lance_core::Error::from(NamespaceError::InvalidInput {
3785 message: "'on' field is required for merge_insert_into_table".to_string(),
3786 })
3787 })?;
3788
3789 let table_has_manifests = self.table_uri_has_actual_manifests(&table_uri).await?;
3790 let (reader, num_rows) =
3791 Self::ipc_reader_from_request_data(&request_data, "merge_insert_into_table")?;
3792
3793 if !table_has_manifests {
3794 let dataset = self
3795 .write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
3796 .await?;
3797 let version = dataset.version().version as i64;
3798 return Ok(MergeInsertIntoTableResponse {
3799 transaction_id: None,
3800 num_updated_rows: Some(0),
3801 num_inserted_rows: Some(num_rows as i64),
3802 num_deleted_rows: Some(0),
3803 version: Some(version),
3804 });
3805 }
3806
3807 let dataset = Arc::new(
3808 self.load_dataset(&table_uri, None, "merge_insert_into_table")
3809 .await?,
3810 );
3811
3812 let mut merge_builder = MergeInsertBuilder::try_new(dataset.clone(), vec![on.clone()])
3813 .map_err(|e| {
3814 lance_core::Error::from(NamespaceError::InvalidInput {
3815 message: format!("Failed to create merge_insert_into_table builder: {}", e),
3816 })
3817 })?;
3818
3819 if let Some(filter) = request.when_matched_update_all_filt.as_deref() {
3820 let behavior = WhenMatched::update_if(dataset.as_ref(), filter).map_err(|e| {
3821 lance_core::Error::from(NamespaceError::InvalidInput {
3822 message: format!(
3823 "Invalid when_matched_update_all_filt for merge_insert_into_table: {}",
3824 e
3825 ),
3826 })
3827 })?;
3828 merge_builder.when_matched(behavior);
3829 } else if request.when_matched_update_all.unwrap_or(false) {
3830 merge_builder.when_matched(WhenMatched::UpdateAll);
3831 }
3832
3833 if matches!(request.when_not_matched_insert_all, Some(false)) {
3834 merge_builder.when_not_matched(WhenNotMatched::DoNothing);
3835 } else {
3836 merge_builder.when_not_matched(WhenNotMatched::InsertAll);
3837 }
3838
3839 if let Some(filter) = request.when_not_matched_by_source_delete_filt.as_deref() {
3840 let behavior = WhenNotMatchedBySource::delete_if(dataset.as_ref(), filter).map_err(|e| {
3841 lance_core::Error::from(NamespaceError::InvalidInput {
3842 message: format!(
3843 "Invalid when_not_matched_by_source_delete_filt for merge_insert_into_table: {}",
3844 e
3845 ),
3846 })
3847 })?;
3848 merge_builder.when_not_matched_by_source(behavior);
3849 } else if request.when_not_matched_by_source_delete.unwrap_or(false) {
3850 merge_builder.when_not_matched_by_source(WhenNotMatchedBySource::Delete);
3851 }
3852
3853 if let Some(use_index) = request.use_index {
3854 merge_builder.use_index(use_index);
3855 }
3856
3857 let (dataset, stats) = merge_builder
3858 .try_build()
3859 .map_err(|e| {
3860 lance_core::Error::from(NamespaceError::InvalidInput {
3861 message: format!("Failed to build merge_insert_into_table job: {}", e),
3862 })
3863 })?
3864 .execute_reader(reader)
3865 .await
3866 .map_err(|e| NamespaceError::Internal {
3867 message: format!(
3868 "Failed to merge_insert_into_table at '{}': {}",
3869 table_uri, e
3870 ),
3871 })?;
3872
3873 Ok(MergeInsertIntoTableResponse {
3874 transaction_id: None,
3875 num_updated_rows: Some(stats.num_updated_rows as i64),
3876 num_inserted_rows: Some(stats.num_inserted_rows as i64),
3877 num_deleted_rows: Some(stats.num_deleted_rows as i64),
3878 version: Some(dataset.version().version as i64),
3879 })
3880 }
3881
3882 async fn query_table(&self, request: QueryTableRequest) -> Result<Bytes> {
3883 use arrow::ipc::writer::FileWriter;
3884
3885 self.record_op("query_table");
3886 let table_uri = self.resolve_table_location(&request.id).await?;
3887 let dataset = self
3888 .load_dataset(&table_uri, request.version, "query_table")
3889 .await?;
3890
3891 let mut scanner = dataset.scan();
3893
3894 let has_vector_query = request
3897 .vector
3898 .single_vector
3899 .as_ref()
3900 .map(|sv| !sv.is_empty())
3901 .unwrap_or(false)
3902 || request
3903 .vector
3904 .multi_vector
3905 .as_ref()
3906 .map(|mv| !mv.is_empty())
3907 .unwrap_or(false);
3908
3909 if let Some(prefilter) = request.prefilter {
3911 scanner.prefilter(prefilter);
3912 }
3913
3914 if has_vector_query {
3916 let vector_column = request.vector_column.as_deref().unwrap_or("vector");
3917
3918 let query_vector: Vec<f32> = request
3920 .vector
3921 .single_vector
3922 .clone()
3923 .or_else(|| {
3924 request
3925 .vector
3926 .multi_vector
3927 .as_ref()
3928 .and_then(|mv| mv.first().cloned())
3929 })
3930 .unwrap_or_default();
3931
3932 if !query_vector.is_empty() {
3933 let k = if request.k > 0 {
3934 request.k as usize
3935 } else {
3936 10
3937 };
3938 let query_array = Float32Array::from(query_vector);
3939 scanner
3940 .nearest(vector_column, &query_array, k)
3941 .map_err(|e| NamespaceError::InvalidInput {
3942 message: format!("Invalid vector search: {:?}", e),
3943 })?;
3944
3945 if let Some(ref distance_type) = request.distance_type {
3947 let metric = match distance_type.to_lowercase().as_str() {
3948 "l2" | "euclidean" => MetricType::L2,
3949 "cosine" => MetricType::Cosine,
3950 "dot" | "inner_product" => MetricType::Dot,
3951 "hamming" => MetricType::Hamming,
3952 _ => {
3953 return Err(NamespaceError::InvalidInput {
3954 message: format!("Unknown distance type: {}", distance_type),
3955 }
3956 .into());
3957 }
3958 };
3959 scanner.distance_metric(metric);
3960 }
3961
3962 if let Some(nprobes) = request.nprobes {
3964 scanner.minimum_nprobes(nprobes as usize);
3965 }
3966
3967 if let Some(ef) = request.ef {
3969 scanner.ef(ef as usize);
3970 }
3971
3972 if let Some(refine_factor) = request.refine_factor {
3974 scanner.refine(refine_factor as u32);
3975 }
3976
3977 if request.lower_bound.is_some() || request.upper_bound.is_some() {
3979 scanner.distance_range(request.lower_bound, request.upper_bound);
3980 }
3981
3982 if let Some(bypass) = request.bypass_vector_index {
3984 scanner.use_index(!bypass);
3985 }
3986
3987 if request.fast_search == Some(true) {
3989 scanner.fast_search();
3990 }
3991 }
3992 }
3993
3994 if let Some(ref fts_query) = request.full_text_query {
3996 if let Some(ref string_query) = fts_query.string_query {
3998 let mut fts = FullTextSearchQuery::new(string_query.query.clone());
3999
4000 if let Some(ref columns) = string_query.columns
4002 && !columns.is_empty()
4003 {
4004 fts = fts
4005 .with_columns(columns)
4006 .map_err(|e| NamespaceError::InvalidInput {
4007 message: format!("Invalid FTS columns: {:?}", e),
4008 })?;
4009 }
4010
4011 scanner
4012 .full_text_search(fts)
4013 .map_err(|e| NamespaceError::InvalidInput {
4014 message: format!("Invalid full text search: {:?}", e),
4015 })?;
4016 }
4017 }
4020
4021 if let Some(ref columns) = request.columns {
4023 if let Some(ref column_names) = columns.column_names
4024 && !column_names.is_empty()
4025 {
4026 scanner
4027 .project(column_names)
4028 .map_err(|e| NamespaceError::InvalidInput {
4029 message: format!("Invalid column projection: {:?}", e),
4030 })?;
4031 } else if let Some(ref column_aliases) = columns.column_aliases
4032 && !column_aliases.is_empty()
4033 {
4034 let transform_pairs: Vec<(String, String)> = column_aliases
4036 .iter()
4037 .map(|(alias, sql)| (alias.clone(), sql.clone()))
4038 .collect();
4039 scanner
4040 .project_with_transform(
4041 &transform_pairs
4042 .iter()
4043 .map(|(a, s)| (a.as_str(), s.as_str()))
4044 .collect::<Vec<_>>(),
4045 )
4046 .map_err(|e| NamespaceError::InvalidInput {
4047 message: format!("Invalid column alias expression: {:?}", e),
4048 })?;
4049 }
4050 }
4051
4052 if let Some(ref filter) = request.filter
4054 && !filter.is_empty()
4055 {
4056 scanner
4057 .filter(filter)
4058 .map_err(|e| NamespaceError::InvalidInput {
4059 message: format!("Invalid filter expression: {:?}", e),
4060 })?;
4061 }
4062
4063 if request.with_row_id == Some(true) {
4065 scanner.with_row_id();
4066 }
4067
4068 if !has_vector_query && request.k > 0 {
4072 let offset = request.offset.map(|o| o as i64);
4073 scanner.limit(Some(request.k as i64), offset).map_err(|e| {
4074 NamespaceError::InvalidInput {
4075 message: format!("Invalid limit/offset: {:?}", e),
4076 }
4077 })?;
4078 } else if has_vector_query && request.offset.is_some() {
4079 let offset = request.offset.map(|o| o as i64);
4081 scanner
4082 .limit(None, offset)
4083 .map_err(|e| NamespaceError::InvalidInput {
4084 message: format!("Invalid offset: {:?}", e),
4085 })?;
4086 }
4087
4088 let batch = scanner
4090 .try_into_batch()
4091 .await
4092 .map_err(|e| NamespaceError::Internal {
4093 message: format!("Failed to execute query: {:?}", e),
4094 })?;
4095
4096 let schema = batch.schema();
4098 let mut buffer = Vec::new();
4099 {
4100 let mut writer = FileWriter::try_new(&mut buffer, &schema).map_err(|e| {
4101 NamespaceError::Internal {
4102 message: format!("Failed to create IPC writer: {:?}", e),
4103 }
4104 })?;
4105 writer.write(&batch).map_err(|e| NamespaceError::Internal {
4106 message: format!("Failed to write batch to IPC: {:?}", e),
4107 })?;
4108 writer.finish().map_err(|e| NamespaceError::Internal {
4109 message: format!("Failed to finish IPC writer: {:?}", e),
4110 })?;
4111 }
4112
4113 Ok(Bytes::from(buffer))
4114 }
4115
4116 async fn list_table_tags(
4117 &self,
4118 request: ListTableTagsRequest,
4119 ) -> Result<ListTableTagsResponse> {
4120 self.record_op("list_table_tags");
4121 let table_uri = self.resolve_table_location(&request.id).await?;
4122 let dataset = self
4123 .load_dataset(&table_uri, None, "list_table_tags")
4124 .await?;
4125
4126 let raw_tags = dataset.tags().list().await.map_err(|e| {
4127 lance_core::Error::from(NamespaceError::Internal {
4128 message: format!("Failed to list tags for table at '{}': {}", table_uri, e),
4129 })
4130 })?;
4131
4132 let tags = raw_tags
4133 .into_iter()
4134 .map(|(name, contents)| {
4135 let mut tag_model =
4136 ModelTagContents::new(contents.version as i64, contents.manifest_size as i64);
4137 tag_model.branch = contents.branch;
4138 (name, tag_model)
4139 })
4140 .collect();
4141
4142 Ok(ListTableTagsResponse {
4143 tags,
4144 page_token: None,
4145 })
4146 }
4147
4148 async fn get_table_tag_version(
4149 &self,
4150 request: GetTableTagVersionRequest,
4151 ) -> Result<GetTableTagVersionResponse> {
4152 self.record_op("get_table_tag_version");
4153 if request.tag.is_empty() {
4154 return Err(NamespaceError::InvalidInput {
4155 message: "tag name must not be empty for get_table_tag_version".to_string(),
4156 }
4157 .into());
4158 }
4159
4160 let table_uri = self.resolve_table_location(&request.id).await?;
4161 let dataset = self
4162 .load_dataset(&table_uri, None, "get_table_tag_version")
4163 .await?;
4164
4165 let version = dataset
4166 .tags()
4167 .get_version(&request.tag)
4168 .await
4169 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
4170
4171 Ok(GetTableTagVersionResponse {
4172 version: version as i64,
4173 })
4174 }
4175
4176 async fn create_table_tag(
4177 &self,
4178 request: CreateTableTagRequest,
4179 ) -> Result<CreateTableTagResponse> {
4180 self.record_op("create_table_tag");
4181 if request.tag.is_empty() {
4182 return Err(NamespaceError::InvalidInput {
4183 message: "tag name must not be empty for create_table_tag".to_string(),
4184 }
4185 .into());
4186 }
4187 if request.version <= 0 {
4188 return Err(NamespaceError::InvalidInput {
4189 message: format!(
4190 "tag version must be a positive integer, got {} for create_table_tag",
4191 request.version
4192 ),
4193 }
4194 .into());
4195 }
4196
4197 let table_uri = self.resolve_table_location(&request.id).await?;
4198 let dataset = self
4199 .load_dataset(&table_uri, None, "create_table_tag")
4200 .await?;
4201
4202 dataset
4203 .tags()
4204 .create(&request.tag, request.version as u64)
4205 .await
4206 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
4207
4208 Ok(CreateTableTagResponse {
4209 transaction_id: None,
4210 })
4211 }
4212
4213 async fn delete_table_tag(
4214 &self,
4215 request: DeleteTableTagRequest,
4216 ) -> Result<DeleteTableTagResponse> {
4217 self.record_op("delete_table_tag");
4218 if request.tag.is_empty() {
4219 return Err(NamespaceError::InvalidInput {
4220 message: "tag name must not be empty for delete_table_tag".to_string(),
4221 }
4222 .into());
4223 }
4224
4225 let table_uri = self.resolve_table_location(&request.id).await?;
4226 let dataset = self
4227 .load_dataset(&table_uri, None, "delete_table_tag")
4228 .await?;
4229
4230 dataset
4231 .tags()
4232 .delete(&request.tag)
4233 .await
4234 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
4235
4236 Ok(DeleteTableTagResponse {
4237 transaction_id: None,
4238 })
4239 }
4240
4241 async fn update_table_tag(
4242 &self,
4243 request: UpdateTableTagRequest,
4244 ) -> Result<UpdateTableTagResponse> {
4245 self.record_op("update_table_tag");
4246 if request.tag.is_empty() {
4247 return Err(NamespaceError::InvalidInput {
4248 message: "tag name must not be empty for update_table_tag".to_string(),
4249 }
4250 .into());
4251 }
4252 if request.version <= 0 {
4253 return Err(NamespaceError::InvalidInput {
4254 message: format!(
4255 "tag version must be a positive integer, got {} for update_table_tag",
4256 request.version
4257 ),
4258 }
4259 .into());
4260 }
4261
4262 let table_uri = self.resolve_table_location(&request.id).await?;
4263 let dataset = self
4264 .load_dataset(&table_uri, None, "update_table_tag")
4265 .await?;
4266
4267 dataset
4268 .tags()
4269 .update(&request.tag, request.version as u64)
4270 .await
4271 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
4272
4273 Ok(UpdateTableTagResponse {
4274 transaction_id: None,
4275 })
4276 }
4277
4278 fn namespace_id(&self) -> String {
4279 format!("DirectoryNamespace {{ root: {:?} }}", self.root)
4280 }
4281}
4282
4283#[cfg(test)]
4284mod tests {
4285 use super::*;
4286 use arrow_ipc::reader::{FileReader, StreamReader};
4287 use lance::dataset::Dataset;
4288 use lance::index::DatasetIndexExt;
4289 use lance_core::utils::tempfile::{TempStdDir, TempStrDir};
4290 use lance_core::utils::testing::CountingObjectStore;
4291 use lance_io::object_store::{providers::local::FileStoreProvider, uri_to_url};
4292 use lance_namespace::models::{
4293 CreateTableRequest, JsonArrowDataType, JsonArrowField, JsonArrowSchema, ListTablesRequest,
4294 QueryTableRequestColumns,
4295 };
4296 use lance_namespace::schema::convert_json_arrow_schema;
4297 use std::io::Cursor;
4298 use std::sync::{
4299 Arc,
4300 atomic::{AtomicUsize, Ordering},
4301 };
4302 use url::Url;
4303
4304 fn assert_plan_contains_all(plan: &str, expected_fragments: &[&str], context: &str) {
4305 for expected_fragment in expected_fragments {
4306 assert!(
4307 plan.contains(expected_fragment),
4308 "{}. Missing fragment: '{}'. Plan:\n{}",
4309 context,
4310 expected_fragment,
4311 plan
4312 );
4313 }
4314 }
4315
4316 async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) {
4318 let temp_dir = TempStdDir::default();
4319
4320 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
4321 .build()
4322 .await
4323 .unwrap();
4324 (namespace, temp_dir)
4325 }
4326
4327 #[derive(Debug)]
4328 struct CountingFileStoreProvider {
4329 listing_count: Arc<AtomicUsize>,
4330 }
4331
4332 #[async_trait]
4333 impl lance_io::object_store::ObjectStoreProvider for CountingFileStoreProvider {
4334 async fn new_store(
4335 &self,
4336 base_path: Url,
4337 params: &ObjectStoreParams,
4338 ) -> Result<ObjectStore> {
4339 let provider = FileStoreProvider;
4340 let mut store = provider.new_store(base_path, params).await?;
4341 store.inner = Arc::new(CountingObjectStore::new(
4342 store.inner.clone(),
4343 self.listing_count.clone(),
4344 ));
4345 Ok(store)
4346 }
4347
4348 fn extract_path(&self, url: &Url) -> Result<Path> {
4349 let provider = FileStoreProvider;
4350 provider.extract_path(url)
4351 }
4352
4353 fn calculate_object_store_prefix(
4354 &self,
4355 url: &Url,
4356 storage_options: Option<&HashMap<String, String>>,
4357 ) -> Result<String> {
4358 let provider = FileStoreProvider;
4359 provider.calculate_object_store_prefix(url, storage_options)
4360 }
4361 }
4362
4363 fn file_object_store_uri(path: &str) -> String {
4364 let file_url = uri_to_url(path).unwrap();
4365 let mut url = Url::parse("file-object-store:///").unwrap();
4366 url.set_path(file_url.path());
4367 url.to_string()
4368 }
4369
4370 fn build_listing_counting_session(listing_count: Arc<AtomicUsize>) -> Arc<Session> {
4371 let registry = Arc::new(ObjectStoreRegistry::default());
4372 registry.insert(
4373 "file-object-store",
4374 Arc::new(CountingFileStoreProvider { listing_count }),
4375 );
4376 Arc::new(Session::new(0, 0, registry))
4377 }
4378
4379 fn create_test_ipc_data(schema: &JsonArrowSchema) -> Vec<u8> {
4381 use arrow::ipc::writer::StreamWriter;
4382
4383 let arrow_schema = convert_json_arrow_schema(schema).unwrap();
4384 let arrow_schema = Arc::new(arrow_schema);
4385 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
4386 let mut buffer = Vec::new();
4387 {
4388 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
4389 writer.write(&batch).unwrap();
4390 writer.finish().unwrap();
4391 }
4392 buffer
4393 }
4394
4395 fn create_ipc_data_from_batches(
4396 schema: Arc<arrow_schema::Schema>,
4397 batches: Vec<arrow::record_batch::RecordBatch>,
4398 ) -> Vec<u8> {
4399 use arrow::ipc::writer::StreamWriter;
4400
4401 let mut buffer = Vec::new();
4402 {
4403 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
4404 for batch in &batches {
4405 writer.write(batch).unwrap();
4406 }
4407 writer.finish().unwrap();
4408 }
4409 buffer
4410 }
4411
4412 fn create_non_empty_test_ipc_data() -> Vec<u8> {
4413 use arrow::array::{Int32Array, StringArray};
4414 use arrow::record_batch::RecordBatch;
4415
4416 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
4417 let batch = RecordBatch::try_new(
4418 schema.clone(),
4419 vec![
4420 Arc::new(Int32Array::from(vec![1, 2])),
4421 Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])),
4422 ],
4423 )
4424 .unwrap();
4425 create_ipc_data_from_batches(schema, vec![batch])
4426 }
4427
4428 fn create_single_row_test_ipc_data() -> Vec<u8> {
4429 use arrow::array::{Int32Array, StringArray};
4430 use arrow::record_batch::RecordBatch;
4431
4432 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
4433 let batch = RecordBatch::try_new(
4434 schema.clone(),
4435 vec![
4436 Arc::new(Int32Array::from(vec![10])),
4437 Arc::new(StringArray::from(vec![Some("carol")])),
4438 ],
4439 )
4440 .unwrap();
4441 create_ipc_data_from_batches(schema, vec![batch])
4442 }
4443
4444 fn create_test_schema() -> JsonArrowSchema {
4446 let int_type = JsonArrowDataType::new("int32".to_string());
4447 let string_type = JsonArrowDataType::new("utf8".to_string());
4448
4449 let id_field = JsonArrowField {
4450 name: "id".to_string(),
4451 r#type: Box::new(int_type),
4452 nullable: false,
4453 metadata: None,
4454 };
4455
4456 let name_field = JsonArrowField {
4457 name: "name".to_string(),
4458 r#type: Box::new(string_type),
4459 nullable: true,
4460 metadata: None,
4461 };
4462
4463 JsonArrowSchema {
4464 fields: vec![id_field, name_field],
4465 metadata: None,
4466 }
4467 }
4468
4469 fn create_scalar_table_ipc_data() -> Vec<u8> {
4470 use arrow::array::{Int32Array, StringArray};
4471 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
4472
4473 let schema = Arc::new(ArrowSchema::new(vec![
4474 Field::new("id", DataType::Int32, false),
4475 Field::new("name", DataType::Utf8, true),
4476 ]));
4477 let batch = arrow::record_batch::RecordBatch::try_new(
4478 schema.clone(),
4479 vec![
4480 Arc::new(Int32Array::from(vec![1, 2, 3])),
4481 Arc::new(StringArray::from(vec!["alice", "bob", "cory"])),
4482 ],
4483 )
4484 .unwrap();
4485 create_ipc_data_from_batches(schema, vec![batch])
4486 }
4487
4488 fn create_vector_table_ipc_data() -> Vec<u8> {
4489 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
4490 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
4491
4492 let schema = Arc::new(ArrowSchema::new(vec![
4493 Field::new("id", DataType::Int32, false),
4494 Field::new(
4495 "vector",
4496 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
4497 true,
4498 ),
4499 ]));
4500 let vector_field = Arc::new(Field::new("item", DataType::Float32, true));
4501 let vectors = FixedSizeListArray::try_new(
4502 vector_field,
4503 2,
4504 Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6])),
4505 None,
4506 )
4507 .unwrap();
4508 let batch = arrow::record_batch::RecordBatch::try_new(
4509 schema.clone(),
4510 vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(vectors)],
4511 )
4512 .unwrap();
4513 create_ipc_data_from_batches(schema, vec![batch])
4514 }
4515
4516 async fn create_scalar_table(namespace: &DirectoryNamespace, table_name: &str) {
4517 let mut create_table_request = CreateTableRequest::new();
4518 create_table_request.id = Some(vec![table_name.to_string()]);
4519 namespace
4520 .create_table(
4521 create_table_request,
4522 Bytes::from(create_scalar_table_ipc_data()),
4523 )
4524 .await
4525 .unwrap();
4526 }
4527
4528 async fn create_vector_table(namespace: &DirectoryNamespace, table_name: &str) {
4529 let mut create_table_request = CreateTableRequest::new();
4530 create_table_request.id = Some(vec![table_name.to_string()]);
4531 namespace
4532 .create_table(
4533 create_table_request,
4534 Bytes::from(create_vector_table_ipc_data()),
4535 )
4536 .await
4537 .unwrap();
4538 }
4539
4540 async fn open_dataset(namespace: &DirectoryNamespace, table_name: &str) -> Dataset {
4541 let mut describe_request = DescribeTableRequest::new();
4542 describe_request.id = Some(vec![table_name.to_string()]);
4543 let table_uri = namespace
4544 .describe_table(describe_request)
4545 .await
4546 .unwrap()
4547 .location
4548 .expect("table location should exist");
4549 Dataset::open(&table_uri).await.unwrap()
4550 }
4551
4552 async fn create_scalar_index(
4553 namespace: &DirectoryNamespace,
4554 table_name: &str,
4555 index_name: &str,
4556 ) -> Option<String> {
4557 use lance_namespace::models::CreateTableIndexRequest;
4558
4559 let mut create_index_request =
4560 CreateTableIndexRequest::new("id".to_string(), "BTREE".to_string());
4561 create_index_request.id = Some(vec![table_name.to_string()]);
4562 create_index_request.name = Some(index_name.to_string());
4563 namespace
4564 .create_table_scalar_index(create_index_request)
4565 .await
4566 .unwrap()
4567 .transaction_id
4568 }
4569
4570 #[tokio::test]
4571 async fn test_create_table() {
4572 let (namespace, _temp_dir) = create_test_namespace().await;
4573
4574 let schema = create_test_schema();
4576 let ipc_data = create_test_ipc_data(&schema);
4577
4578 let mut request = CreateTableRequest::new();
4579 request.id = Some(vec!["test_table".to_string()]);
4580
4581 let response = namespace
4582 .create_table(request, bytes::Bytes::from(ipc_data))
4583 .await
4584 .unwrap();
4585
4586 assert!(response.location.is_some());
4587 assert!(response.location.unwrap().ends_with("test_table.lance"));
4588 assert_eq!(response.version, Some(1));
4589 }
4590
4591 #[tokio::test]
4592 async fn test_create_table_without_data() {
4593 let (namespace, _temp_dir) = create_test_namespace().await;
4594
4595 let mut request = CreateTableRequest::new();
4596 request.id = Some(vec!["test_table".to_string()]);
4597
4598 let result = namespace.create_table(request, bytes::Bytes::new()).await;
4599 assert!(result.is_err());
4600 assert!(
4601 result
4602 .unwrap_err()
4603 .to_string()
4604 .contains("Arrow IPC stream) is required")
4605 );
4606 }
4607
4608 #[tokio::test]
4609 async fn test_create_table_with_invalid_id() {
4610 let (namespace, _temp_dir) = create_test_namespace().await;
4611
4612 let schema = create_test_schema();
4614 let ipc_data = create_test_ipc_data(&schema);
4615
4616 let mut request = CreateTableRequest::new();
4618 request.id = Some(vec![]);
4619
4620 let result = namespace
4621 .create_table(request, bytes::Bytes::from(ipc_data.clone()))
4622 .await;
4623 assert!(result.is_err());
4624
4625 let mut create_ns_req = CreateNamespaceRequest::new();
4628 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
4629 namespace.create_namespace(create_ns_req).await.unwrap();
4630
4631 let mut request = CreateTableRequest::new();
4633 request.id = Some(vec!["test_namespace".to_string(), "table".to_string()]);
4634
4635 let result = namespace
4636 .create_table(request, bytes::Bytes::from(ipc_data))
4637 .await;
4638 assert!(
4640 result.is_ok(),
4641 "Multi-level table IDs should work with manifest enabled"
4642 );
4643 }
4644
4645 #[tokio::test]
4646 async fn test_list_tables() {
4647 let (namespace, _temp_dir) = create_test_namespace().await;
4648
4649 let mut request = ListTablesRequest::new();
4651 request.id = Some(vec![]);
4652 let response = namespace.list_tables(request).await.unwrap();
4653 assert_eq!(response.tables.len(), 0);
4654
4655 let schema = create_test_schema();
4657 let ipc_data = create_test_ipc_data(&schema);
4658
4659 let mut create_request = CreateTableRequest::new();
4661 create_request.id = Some(vec!["table1".to_string()]);
4662 namespace
4663 .create_table(create_request, bytes::Bytes::from(ipc_data.clone()))
4664 .await
4665 .unwrap();
4666
4667 let mut create_request = CreateTableRequest::new();
4669 create_request.id = Some(vec!["table2".to_string()]);
4670 namespace
4671 .create_table(create_request, bytes::Bytes::from(ipc_data))
4672 .await
4673 .unwrap();
4674
4675 let mut request = ListTablesRequest::new();
4677 request.id = Some(vec![]);
4678 let response = namespace.list_tables(request).await.unwrap();
4679 let tables = response.tables;
4680 assert_eq!(tables.len(), 2);
4681 assert!(tables.contains(&"table1".to_string()));
4682 assert!(tables.contains(&"table2".to_string()));
4683 }
4684
4685 #[tokio::test]
4686 async fn test_list_tables_pagination() {
4687 let (namespace, _temp_dir) = create_test_namespace().await;
4688
4689 let schema = create_test_schema();
4690 let ipc_data = create_test_ipc_data(&schema);
4691
4692 for name in ["alpha", "bravo", "charlie"] {
4693 let mut req = CreateTableRequest::new();
4694 req.id = Some(vec![name.to_string()]);
4695 namespace
4696 .create_table(req, bytes::Bytes::from(ipc_data.clone()))
4697 .await
4698 .unwrap();
4699 }
4700
4701 let first_page = namespace
4703 .list_tables(ListTablesRequest {
4704 id: Some(vec![]),
4705 limit: Some(2),
4706 ..Default::default()
4707 })
4708 .await
4709 .unwrap();
4710
4711 assert_eq!(first_page.tables, vec!["alpha", "bravo"]);
4712 assert_eq!(first_page.page_token.as_deref(), Some("bravo"));
4713
4714 let second_page = namespace
4716 .list_tables(ListTablesRequest {
4717 id: Some(vec![]),
4718 limit: Some(2),
4719 page_token: first_page.page_token.clone(),
4720 ..Default::default()
4721 })
4722 .await
4723 .unwrap();
4724
4725 assert_eq!(second_page.tables, vec!["charlie"]);
4726 assert!(second_page.page_token.is_none());
4727 }
4728
4729 #[tokio::test]
4730 async fn test_list_tables_pagination_limit_zero() {
4731 let (namespace, _temp_dir) = create_test_namespace().await;
4732
4733 let schema = create_test_schema();
4734 let ipc_data = create_test_ipc_data(&schema);
4735
4736 let mut req = CreateTableRequest::new();
4737 req.id = Some(vec!["alpha".to_string()]);
4738 namespace
4739 .create_table(req, bytes::Bytes::from(ipc_data))
4740 .await
4741 .unwrap();
4742
4743 let response = namespace
4744 .list_tables(ListTablesRequest {
4745 id: Some(vec![]),
4746 limit: Some(0),
4747 ..Default::default()
4748 })
4749 .await
4750 .unwrap();
4751
4752 assert!(response.tables.is_empty());
4753 assert!(response.page_token.is_none());
4754 }
4755
4756 #[tokio::test]
4757 async fn test_list_tables_with_namespace_id() {
4758 let (namespace, _temp_dir) = create_test_namespace().await;
4759
4760 let mut create_ns_req = CreateNamespaceRequest::new();
4762 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
4763 namespace.create_namespace(create_ns_req).await.unwrap();
4764
4765 let mut request = ListTablesRequest::new();
4767 request.id = Some(vec!["test_namespace".to_string()]);
4768
4769 let result = namespace.list_tables(request).await;
4770 assert!(
4772 result.is_ok(),
4773 "list_tables should work with child namespace when manifest is enabled"
4774 );
4775 let response = result.unwrap();
4776 assert_eq!(
4777 response.tables.len(),
4778 0,
4779 "Namespace should have no tables yet"
4780 );
4781 }
4782
4783 #[tokio::test]
4784 async fn test_create_scalar_index() {
4785 let (namespace, _temp_dir) = create_test_namespace().await;
4786 create_scalar_table(&namespace, "users").await;
4787
4788 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
4789 let dataset = open_dataset(&namespace, "users").await;
4790 let expected_transaction_id = dataset
4791 .read_transaction()
4792 .await
4793 .unwrap()
4794 .map(|transaction| transaction.uuid);
4795 assert_eq!(transaction_id, expected_transaction_id);
4796 let indices = dataset.load_indices().await.unwrap();
4797 assert!(indices.iter().any(|index| index.name == "users_id_idx"));
4798 }
4799
4800 #[tokio::test]
4801 async fn test_create_vector_index() {
4802 use lance_namespace::models::CreateTableIndexRequest;
4803
4804 let (namespace, _temp_dir) = create_test_namespace().await;
4805 create_vector_table(&namespace, "vectors").await;
4806
4807 let mut create_index_request =
4808 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
4809 create_index_request.id = Some(vec!["vectors".to_string()]);
4810 create_index_request.name = Some("vector_idx".to_string());
4811 create_index_request.distance_type = Some("l2".to_string());
4812 let transaction_id = namespace
4813 .create_table_index(create_index_request)
4814 .await
4815 .unwrap()
4816 .transaction_id;
4817
4818 let dataset = open_dataset(&namespace, "vectors").await;
4819 let expected_transaction_id = dataset
4820 .read_transaction()
4821 .await
4822 .unwrap()
4823 .map(|transaction| transaction.uuid);
4824 assert_eq!(transaction_id, expected_transaction_id);
4825 let indices = dataset.load_indices().await.unwrap();
4826 assert!(indices.iter().any(|index| index.name == "vector_idx"));
4827 }
4828
4829 #[tokio::test]
4830 async fn test_list_table_indices() {
4831 use lance_namespace::models::ListTableIndicesRequest;
4832
4833 let (namespace, _temp_dir) = create_test_namespace().await;
4834 create_scalar_table(&namespace, "users").await;
4835 create_scalar_index(&namespace, "users", "a_idx").await;
4836 create_scalar_index(&namespace, "users", "b_idx").await;
4837 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
4838
4839 let response = namespace
4840 .list_table_indices(ListTableIndicesRequest {
4841 id: Some(vec!["users".to_string()]),
4842 ..Default::default()
4843 })
4844 .await
4845 .unwrap();
4846
4847 assert_eq!(response.indexes.len(), 3);
4848 assert_eq!(response.indexes[0].index_name, "a_idx");
4849 assert_eq!(response.indexes[1].index_name, "b_idx");
4850 assert_eq!(response.indexes[2].index_name, "users_id_idx");
4851 assert!(response.page_token.is_none());
4852 let users_id_idx = response
4853 .indexes
4854 .iter()
4855 .find(|index| index.index_name == "users_id_idx")
4856 .unwrap();
4857 assert_eq!(users_id_idx.columns, vec!["id"]);
4858 assert_eq!(users_id_idx.status, "SUCCEEDED");
4859
4860 let dataset = open_dataset(&namespace, "users").await;
4861 let expected_transaction_id = dataset
4862 .read_transaction()
4863 .await
4864 .unwrap()
4865 .map(|transaction| transaction.uuid);
4866 assert_eq!(transaction_id, expected_transaction_id);
4867 let indices = dataset.load_indices().await.unwrap();
4868 assert_eq!(
4869 indices
4870 .iter()
4871 .filter(|index| index.name == "users_id_idx")
4872 .count(),
4873 1
4874 );
4875
4876 let first_page = namespace
4877 .list_table_indices(ListTableIndicesRequest {
4878 id: Some(vec!["users".to_string()]),
4879 limit: Some(2),
4880 ..Default::default()
4881 })
4882 .await
4883 .unwrap();
4884
4885 assert_eq!(first_page.indexes.len(), 2);
4886 assert_eq!(first_page.indexes[0].index_name, "a_idx");
4887 assert_eq!(first_page.indexes[1].index_name, "b_idx");
4888 assert_eq!(first_page.page_token.as_deref(), Some("b_idx"));
4889
4890 let second_page = namespace
4891 .list_table_indices(ListTableIndicesRequest {
4892 id: Some(vec!["users".to_string()]),
4893 page_token: first_page.page_token.clone(),
4894 limit: Some(2),
4895 ..Default::default()
4896 })
4897 .await
4898 .unwrap();
4899
4900 assert_eq!(second_page.indexes.len(), 1);
4901 assert_eq!(second_page.indexes[0].index_name, "users_id_idx");
4902 assert!(second_page.page_token.is_none());
4903 }
4904
4905 #[tokio::test]
4906 async fn test_describe_table_index_stats() {
4907 use lance_namespace::models::DescribeTableIndexStatsRequest;
4908
4909 let (namespace, _temp_dir) = create_test_namespace().await;
4910 create_scalar_table(&namespace, "users").await;
4911 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
4912
4913 let response = namespace
4914 .describe_table_index_stats(DescribeTableIndexStatsRequest {
4915 id: Some(vec!["users".to_string()]),
4916 index_name: Some("users_id_idx".to_string()),
4917 ..Default::default()
4918 })
4919 .await
4920 .unwrap();
4921 assert_eq!(response.index_type, Some("BTree".to_string()));
4922 assert_eq!(response.num_indices, Some(1));
4923 assert_eq!(response.num_indexed_rows, Some(3));
4924 assert_eq!(response.num_unindexed_rows, Some(0));
4925
4926 let dataset = open_dataset(&namespace, "users").await;
4927 let expected_transaction_id = dataset
4928 .read_transaction()
4929 .await
4930 .unwrap()
4931 .map(|transaction| transaction.uuid);
4932 assert_eq!(transaction_id, expected_transaction_id);
4933 let stats: serde_json::Value =
4934 serde_json::from_str(&dataset.index_statistics("users_id_idx").await.unwrap()).unwrap();
4935 assert_eq!(stats["index_type"], "BTree");
4936 assert_eq!(stats["num_indices"], 1);
4937 assert_eq!(stats["num_indexed_rows"], 3);
4938 assert_eq!(stats["num_unindexed_rows"], 0);
4939 }
4940
4941 #[tokio::test]
4942 async fn test_describe_transaction() {
4943 use lance_namespace::models::DescribeTransactionRequest;
4944
4945 let (namespace, _temp_dir) = create_test_namespace().await;
4946 create_scalar_table(&namespace, "users").await;
4947 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
4948 let dataset = open_dataset(&namespace, "users").await;
4949 let latest_transaction = dataset.read_transaction().await.unwrap();
4950 assert_eq!(
4951 transaction_id,
4952 latest_transaction
4953 .as_ref()
4954 .map(|transaction| transaction.uuid.clone())
4955 );
4956
4957 if let Some(transaction_id) = transaction_id {
4958 let response = namespace
4959 .describe_transaction(DescribeTransactionRequest {
4960 id: Some(vec!["users".to_string(), transaction_id.clone()]),
4961 ..Default::default()
4962 })
4963 .await
4964 .unwrap();
4965 assert_eq!(response.status, "SUCCEEDED");
4966 assert_eq!(
4967 response
4968 .properties
4969 .as_ref()
4970 .and_then(|props| props.get("operation")),
4971 Some(&"CreateIndex".to_string())
4972 );
4973 assert_eq!(
4974 response
4975 .properties
4976 .as_ref()
4977 .and_then(|props| props.get("uuid")),
4978 Some(&transaction_id)
4979 );
4980 } else {
4981 assert!(latest_transaction.is_none());
4982 }
4983 }
4984
4985 #[tokio::test]
4986 async fn test_drop_table_index() {
4987 use lance_namespace::models::{DropTableIndexRequest, ListTableIndicesRequest};
4988
4989 let (namespace, _temp_dir) = create_test_namespace().await;
4990 create_scalar_table(&namespace, "users").await;
4991 let create_transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
4992
4993 let drop_transaction_id = namespace
4994 .drop_table_index(DropTableIndexRequest {
4995 id: Some(vec!["users".to_string()]),
4996 index_name: Some("users_id_idx".to_string()),
4997 ..Default::default()
4998 })
4999 .await
5000 .unwrap()
5001 .transaction_id;
5002
5003 let dataset = open_dataset(&namespace, "users").await;
5004 let previous_dataset = dataset
5005 .checkout_version(dataset.version().version - 1)
5006 .await
5007 .unwrap();
5008 let previous_transaction_id = previous_dataset
5009 .read_transaction()
5010 .await
5011 .unwrap()
5012 .map(|transaction| transaction.uuid);
5013 assert_eq!(create_transaction_id, previous_transaction_id);
5014 let expected_drop_transaction_id = dataset
5015 .read_transaction()
5016 .await
5017 .unwrap()
5018 .map(|transaction| transaction.uuid);
5019 assert_eq!(drop_transaction_id, expected_drop_transaction_id);
5020 let indices = dataset.load_indices().await.unwrap();
5021 assert!(!indices.iter().any(|index| index.name == "users_id_idx"));
5022
5023 let list_response = namespace
5024 .list_table_indices(ListTableIndicesRequest {
5025 id: Some(vec!["users".to_string()]),
5026 ..Default::default()
5027 })
5028 .await
5029 .unwrap();
5030 assert!(list_response.indexes.is_empty());
5031 }
5032
5033 #[tokio::test]
5034 async fn test_describe_table() {
5035 let (namespace, _temp_dir) = create_test_namespace().await;
5036
5037 let schema = create_test_schema();
5039 let ipc_data = create_test_ipc_data(&schema);
5040
5041 let mut create_request = CreateTableRequest::new();
5042 create_request.id = Some(vec!["test_table".to_string()]);
5043 namespace
5044 .create_table(create_request, bytes::Bytes::from(ipc_data))
5045 .await
5046 .unwrap();
5047
5048 let mut request = DescribeTableRequest::new();
5050 request.id = Some(vec!["test_table".to_string()]);
5051 let response = namespace.describe_table(request).await.unwrap();
5052
5053 assert!(response.location.is_some());
5054 assert!(response.location.unwrap().ends_with("test_table.lance"));
5055 }
5056
5057 #[tokio::test]
5058 async fn test_describe_nonexistent_table() {
5059 let (namespace, _temp_dir) = create_test_namespace().await;
5060
5061 let mut request = DescribeTableRequest::new();
5062 request.id = Some(vec!["nonexistent".to_string()]);
5063
5064 let result = namespace.describe_table(request).await;
5065 assert!(result.is_err());
5066 assert!(result.unwrap_err().to_string().contains("Table not found"));
5067 }
5068
5069 #[tokio::test]
5070 async fn test_table_exists() {
5071 let (namespace, _temp_dir) = create_test_namespace().await;
5072
5073 let schema = create_test_schema();
5075 let ipc_data = create_test_ipc_data(&schema);
5076
5077 let mut create_request = CreateTableRequest::new();
5078 create_request.id = Some(vec!["existing_table".to_string()]);
5079 namespace
5080 .create_table(create_request, bytes::Bytes::from(ipc_data))
5081 .await
5082 .unwrap();
5083
5084 let mut request = TableExistsRequest::new();
5086 request.id = Some(vec!["existing_table".to_string()]);
5087 let result = namespace.table_exists(request).await;
5088 assert!(result.is_ok());
5089
5090 let mut request = TableExistsRequest::new();
5092 request.id = Some(vec!["nonexistent".to_string()]);
5093 let result = namespace.table_exists(request).await;
5094 assert!(result.is_err());
5095 assert!(result.unwrap_err().to_string().contains("Table not found"));
5096 }
5097
5098 #[tokio::test]
5099 async fn test_drop_table() {
5100 let (namespace, _temp_dir) = create_test_namespace().await;
5101
5102 let schema = create_test_schema();
5104 let ipc_data = create_test_ipc_data(&schema);
5105
5106 let mut create_request = CreateTableRequest::new();
5107 create_request.id = Some(vec!["table_to_drop".to_string()]);
5108 namespace
5109 .create_table(create_request, bytes::Bytes::from(ipc_data))
5110 .await
5111 .unwrap();
5112
5113 let mut exists_request = TableExistsRequest::new();
5115 exists_request.id = Some(vec!["table_to_drop".to_string()]);
5116 assert!(namespace.table_exists(exists_request.clone()).await.is_ok());
5117
5118 let mut drop_request = DropTableRequest::new();
5120 drop_request.id = Some(vec!["table_to_drop".to_string()]);
5121 let response = namespace.drop_table(drop_request).await.unwrap();
5122 assert!(response.location.is_some());
5123
5124 assert!(namespace.table_exists(exists_request).await.is_err());
5126 }
5127
5128 #[tokio::test]
5129 async fn test_drop_nonexistent_table() {
5130 let (namespace, _temp_dir) = create_test_namespace().await;
5131
5132 let mut request = DropTableRequest::new();
5133 request.id = Some(vec!["nonexistent".to_string()]);
5134
5135 let result = namespace.drop_table(request).await;
5137 let _ = result;
5140 }
5141
5142 #[tokio::test]
5143 async fn test_root_namespace_operations() {
5144 let (namespace, _temp_dir) = create_test_namespace().await;
5145
5146 let mut request = ListNamespacesRequest::new();
5148 request.id = Some(vec![]);
5149 let result = namespace.list_namespaces(request).await;
5150 assert!(result.is_ok());
5151 assert_eq!(result.unwrap().namespaces.len(), 0);
5152
5153 let mut request = DescribeNamespaceRequest::new();
5155 request.id = Some(vec![]);
5156 let result = namespace.describe_namespace(request).await;
5157 assert!(result.is_ok());
5158
5159 let mut request = NamespaceExistsRequest::new();
5161 request.id = Some(vec![]);
5162 let result = namespace.namespace_exists(request).await;
5163 assert!(result.is_ok());
5164
5165 let mut request = CreateNamespaceRequest::new();
5167 request.id = Some(vec![]);
5168 let result = namespace.create_namespace(request).await;
5169 assert!(result.is_err());
5170 assert!(result.unwrap_err().to_string().contains("already exists"));
5171
5172 let mut request = DropNamespaceRequest::new();
5174 request.id = Some(vec![]);
5175 let result = namespace.drop_namespace(request).await;
5176 assert!(result.is_err());
5177 assert!(
5178 result
5179 .unwrap_err()
5180 .to_string()
5181 .contains("cannot be dropped")
5182 );
5183 }
5184
5185 #[tokio::test]
5186 async fn test_non_root_namespace_operations() {
5187 let (namespace, _temp_dir) = create_test_namespace().await;
5188
5189 let mut request = CreateNamespaceRequest::new();
5192 request.id = Some(vec!["child".to_string()]);
5193 let result = namespace.create_namespace(request).await;
5194 assert!(
5195 result.is_ok(),
5196 "Child namespace creation should succeed with manifest enabled"
5197 );
5198
5199 let mut request = NamespaceExistsRequest::new();
5201 request.id = Some(vec!["child".to_string()]);
5202 let result = namespace.namespace_exists(request).await;
5203 assert!(
5204 result.is_ok(),
5205 "Child namespace should exist after creation"
5206 );
5207
5208 let mut request = DropNamespaceRequest::new();
5210 request.id = Some(vec!["child".to_string()]);
5211 let result = namespace.drop_namespace(request).await;
5212 assert!(
5213 result.is_ok(),
5214 "Child namespace drop should succeed with manifest enabled"
5215 );
5216
5217 let mut request = NamespaceExistsRequest::new();
5219 request.id = Some(vec!["child".to_string()]);
5220 let result = namespace.namespace_exists(request).await;
5221 assert!(
5222 result.is_err(),
5223 "Child namespace should not exist after drop"
5224 );
5225 }
5226
5227 #[tokio::test]
5228 async fn test_config_custom_root() {
5229 let temp_dir = TempStdDir::default();
5230 let custom_path = temp_dir.join("custom");
5231 std::fs::create_dir(&custom_path).unwrap();
5232
5233 let namespace = DirectoryNamespaceBuilder::new(custom_path.to_string_lossy().to_string())
5234 .build()
5235 .await
5236 .unwrap();
5237
5238 let schema = create_test_schema();
5240 let ipc_data = create_test_ipc_data(&schema);
5241
5242 let mut request = CreateTableRequest::new();
5244 request.id = Some(vec!["test_table".to_string()]);
5245
5246 let response = namespace
5247 .create_table(request, bytes::Bytes::from(ipc_data))
5248 .await
5249 .unwrap();
5250
5251 assert!(response.location.unwrap().contains("custom"));
5252 }
5253
5254 #[tokio::test]
5255 async fn test_config_storage_options() {
5256 let temp_dir = TempStdDir::default();
5257
5258 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
5259 .storage_option("option1", "value1")
5260 .storage_option("option2", "value2")
5261 .build()
5262 .await
5263 .unwrap();
5264
5265 let schema = create_test_schema();
5267 let ipc_data = create_test_ipc_data(&schema);
5268
5269 let mut request = CreateTableRequest::new();
5271 request.id = Some(vec!["test_table".to_string()]);
5272
5273 let response = namespace
5274 .create_table(request, bytes::Bytes::from(ipc_data))
5275 .await
5276 .unwrap();
5277
5278 let storage_options = response.storage_options.unwrap();
5279 assert_eq!(storage_options.get("option1"), Some(&"value1".to_string()));
5280 assert_eq!(storage_options.get("option2"), Some(&"value2".to_string()));
5281 }
5282
5283 #[tokio::test]
5287 async fn test_no_storage_options_without_vendor() {
5288 use lance_namespace::models::DeclareTableRequest;
5289
5290 let temp_dir = TempStdDir::default();
5291
5292 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
5294 .manifest_enabled(false)
5295 .storage_option("aws_access_key_id", "AKID")
5296 .storage_option("aws_secret_access_key", "SECRET")
5297 .storage_option("region", "us-east-1")
5298 .build()
5299 .await
5300 .unwrap();
5301
5302 let schema = create_test_schema();
5303 let ipc_data = create_test_ipc_data(&schema);
5304
5305 let mut create_req = CreateTableRequest::new();
5307 create_req.id = Some(vec!["t1".to_string()]);
5308 namespace
5309 .create_table(create_req, bytes::Bytes::from(ipc_data))
5310 .await
5311 .unwrap();
5312
5313 let mut desc_req = DescribeTableRequest::new();
5315 desc_req.id = Some(vec!["t1".to_string()]);
5316 let resp = namespace.describe_table(desc_req).await.unwrap();
5317 assert!(resp.storage_options.is_none());
5318
5319 let mut decl_req = DeclareTableRequest::new();
5321 decl_req.id = Some(vec!["t2".to_string()]);
5322 let resp = namespace.declare_table(decl_req).await.unwrap();
5323 assert!(resp.storage_options.is_none());
5324 }
5325
5326 #[tokio::test]
5328 async fn test_no_storage_options_without_vendor_manifest() {
5329 let temp_dir = TempStdDir::default();
5330
5331 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
5332 .storage_option("aws_access_key_id", "AKID")
5333 .storage_option("aws_secret_access_key", "SECRET")
5334 .storage_option("region", "us-east-1")
5335 .build()
5336 .await
5337 .unwrap();
5338
5339 let schema = create_test_schema();
5340 let ipc_data = create_test_ipc_data(&schema);
5341
5342 let mut create_req = CreateTableRequest::new();
5343 create_req.id = Some(vec!["t1".to_string()]);
5344 namespace
5345 .create_table(create_req, bytes::Bytes::from(ipc_data))
5346 .await
5347 .unwrap();
5348
5349 let mut desc_req = DescribeTableRequest::new();
5351 desc_req.id = Some(vec!["t1".to_string()]);
5352 let resp = namespace.describe_table(desc_req).await.unwrap();
5353 assert!(resp.storage_options.is_none());
5354 }
5355
5356 #[tokio::test]
5357 async fn test_from_properties_manifest_enabled() {
5358 let temp_dir = TempStdDir::default();
5359
5360 let mut properties = HashMap::new();
5361 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
5362 properties.insert("manifest_enabled".to_string(), "true".to_string());
5363 properties.insert("dir_listing_enabled".to_string(), "false".to_string());
5364
5365 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
5366 assert!(builder.manifest_enabled);
5367 assert!(!builder.dir_listing_enabled);
5368
5369 let namespace = builder.build().await.unwrap();
5370
5371 let schema = create_test_schema();
5373 let ipc_data = create_test_ipc_data(&schema);
5374
5375 let mut request = CreateTableRequest::new();
5377 request.id = Some(vec!["test_table".to_string()]);
5378
5379 let response = namespace
5380 .create_table(request, bytes::Bytes::from(ipc_data))
5381 .await
5382 .unwrap();
5383
5384 assert!(response.location.is_some());
5385 }
5386
5387 #[tokio::test]
5388 async fn test_from_properties_dir_listing_enabled() {
5389 let temp_dir = TempStdDir::default();
5390
5391 let mut properties = HashMap::new();
5392 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
5393 properties.insert("manifest_enabled".to_string(), "false".to_string());
5394 properties.insert("dir_listing_enabled".to_string(), "true".to_string());
5395
5396 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
5397 assert!(!builder.manifest_enabled);
5398 assert!(builder.dir_listing_enabled);
5399
5400 let namespace = builder.build().await.unwrap();
5401
5402 let schema = create_test_schema();
5404 let ipc_data = create_test_ipc_data(&schema);
5405
5406 let mut request = CreateTableRequest::new();
5408 request.id = Some(vec!["test_table".to_string()]);
5409
5410 let response = namespace
5411 .create_table(request, bytes::Bytes::from(ipc_data))
5412 .await
5413 .unwrap();
5414
5415 assert!(response.location.is_some());
5416 }
5417
5418 #[tokio::test]
5419 async fn test_from_properties_defaults() {
5420 let temp_dir = TempStdDir::default();
5421
5422 let mut properties = HashMap::new();
5423 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
5424
5425 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
5426 assert!(builder.manifest_enabled);
5428 assert!(builder.dir_listing_enabled);
5429 }
5430
5431 #[tokio::test]
5432 async fn test_from_properties_with_storage_options() {
5433 let temp_dir = TempStdDir::default();
5434
5435 let mut properties = HashMap::new();
5436 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
5437 properties.insert("manifest_enabled".to_string(), "true".to_string());
5438 properties.insert("storage.region".to_string(), "us-west-2".to_string());
5439 properties.insert("storage.bucket".to_string(), "my-bucket".to_string());
5440
5441 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
5442 assert!(builder.manifest_enabled);
5443 assert!(builder.storage_options.is_some());
5444
5445 let storage_options = builder.storage_options.unwrap();
5446 assert_eq!(
5447 storage_options.get("region"),
5448 Some(&"us-west-2".to_string())
5449 );
5450 assert_eq!(
5451 storage_options.get("bucket"),
5452 Some(&"my-bucket".to_string())
5453 );
5454 }
5455
5456 #[tokio::test]
5457 async fn test_various_arrow_types() {
5458 let (namespace, _temp_dir) = create_test_namespace().await;
5459
5460 let fields = vec![
5462 JsonArrowField {
5463 name: "bool_col".to_string(),
5464 r#type: Box::new(JsonArrowDataType::new("bool".to_string())),
5465 nullable: true,
5466 metadata: None,
5467 },
5468 JsonArrowField {
5469 name: "int8_col".to_string(),
5470 r#type: Box::new(JsonArrowDataType::new("int8".to_string())),
5471 nullable: true,
5472 metadata: None,
5473 },
5474 JsonArrowField {
5475 name: "float64_col".to_string(),
5476 r#type: Box::new(JsonArrowDataType::new("float64".to_string())),
5477 nullable: true,
5478 metadata: None,
5479 },
5480 JsonArrowField {
5481 name: "binary_col".to_string(),
5482 r#type: Box::new(JsonArrowDataType::new("binary".to_string())),
5483 nullable: true,
5484 metadata: None,
5485 },
5486 ];
5487
5488 let schema = JsonArrowSchema {
5489 fields,
5490 metadata: None,
5491 };
5492
5493 let ipc_data = create_test_ipc_data(&schema);
5495
5496 let mut request = CreateTableRequest::new();
5497 request.id = Some(vec!["complex_table".to_string()]);
5498
5499 let response = namespace
5500 .create_table(request, bytes::Bytes::from(ipc_data))
5501 .await
5502 .unwrap();
5503
5504 assert!(response.location.is_some());
5505 }
5506
5507 #[tokio::test]
5508 async fn test_connect_dir() {
5509 let temp_dir = TempStdDir::default();
5510
5511 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
5512 .build()
5513 .await
5514 .unwrap();
5515
5516 let mut request = ListTablesRequest::new();
5518 request.id = Some(vec![]);
5519 let response = namespace.list_tables(request).await.unwrap();
5520 assert_eq!(response.tables.len(), 0);
5521 }
5522
5523 #[tokio::test]
5524 async fn test_create_table_with_ipc_data() {
5525 use arrow::array::{Int32Array, StringArray};
5526 use arrow::ipc::writer::StreamWriter;
5527
5528 let (namespace, _temp_dir) = create_test_namespace().await;
5529
5530 let schema = create_test_schema();
5532
5533 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
5535 let arrow_schema = Arc::new(arrow_schema);
5536
5537 let id_array = Int32Array::from(vec![1, 2, 3]);
5539 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
5540 let batch = arrow::record_batch::RecordBatch::try_new(
5541 arrow_schema.clone(),
5542 vec![Arc::new(id_array), Arc::new(name_array)],
5543 )
5544 .unwrap();
5545
5546 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
5554 let mut request = CreateTableRequest::new();
5556 request.id = Some(vec!["test_table_with_data".to_string()]);
5557
5558 let response = namespace
5559 .create_table(request, Bytes::from(buffer))
5560 .await
5561 .unwrap();
5562
5563 assert_eq!(response.version, Some(1));
5564 assert!(
5565 response
5566 .location
5567 .unwrap()
5568 .contains("test_table_with_data.lance")
5569 );
5570
5571 let mut exists_request = TableExistsRequest::new();
5573 exists_request.id = Some(vec!["test_table_with_data".to_string()]);
5574 namespace.table_exists(exists_request).await.unwrap();
5575 }
5576
5577 #[tokio::test]
5578 async fn test_child_namespace_create_and_list() {
5579 let (namespace, _temp_dir) = create_test_namespace().await;
5580
5581 for i in 1..=3 {
5583 let mut create_req = CreateNamespaceRequest::new();
5584 create_req.id = Some(vec![format!("ns{}", i)]);
5585 let result = namespace.create_namespace(create_req).await;
5586 assert!(result.is_ok(), "Failed to create child namespace ns{}", i);
5587 }
5588
5589 let list_req = ListNamespacesRequest {
5591 id: Some(vec![]),
5592 ..Default::default()
5593 };
5594 let result = namespace.list_namespaces(list_req).await;
5595 assert!(result.is_ok());
5596 let namespaces = result.unwrap().namespaces;
5597 assert_eq!(namespaces.len(), 3);
5598 assert!(namespaces.contains(&"ns1".to_string()));
5599 assert!(namespaces.contains(&"ns2".to_string()));
5600 assert!(namespaces.contains(&"ns3".to_string()));
5601 }
5602
5603 #[tokio::test]
5604 async fn test_nested_namespace_hierarchy() {
5605 let (namespace, _temp_dir) = create_test_namespace().await;
5606
5607 let mut create_req = CreateNamespaceRequest::new();
5609 create_req.id = Some(vec!["parent".to_string()]);
5610 namespace.create_namespace(create_req).await.unwrap();
5611
5612 let mut create_req = CreateNamespaceRequest::new();
5614 create_req.id = Some(vec!["parent".to_string(), "child1".to_string()]);
5615 namespace.create_namespace(create_req).await.unwrap();
5616
5617 let mut create_req = CreateNamespaceRequest::new();
5618 create_req.id = Some(vec!["parent".to_string(), "child2".to_string()]);
5619 namespace.create_namespace(create_req).await.unwrap();
5620
5621 let list_req = ListNamespacesRequest {
5623 id: Some(vec!["parent".to_string()]),
5624 ..Default::default()
5625 };
5626 let result = namespace.list_namespaces(list_req).await;
5627 assert!(result.is_ok());
5628 let children = result.unwrap().namespaces;
5629 assert_eq!(children.len(), 2);
5630 assert!(children.contains(&"child1".to_string()));
5631 assert!(children.contains(&"child2".to_string()));
5632
5633 let list_req = ListNamespacesRequest {
5635 id: Some(vec![]),
5636 ..Default::default()
5637 };
5638 let result = namespace.list_namespaces(list_req).await;
5639 assert!(result.is_ok());
5640 let root_namespaces = result.unwrap().namespaces;
5641 assert_eq!(root_namespaces.len(), 1);
5642 assert_eq!(root_namespaces[0], "parent");
5643 }
5644
5645 #[tokio::test]
5646 async fn test_table_in_child_namespace() {
5647 let (namespace, _temp_dir) = create_test_namespace().await;
5648
5649 let mut create_ns_req = CreateNamespaceRequest::new();
5651 create_ns_req.id = Some(vec!["test_ns".to_string()]);
5652 namespace.create_namespace(create_ns_req).await.unwrap();
5653
5654 let schema = create_test_schema();
5656 let ipc_data = create_test_ipc_data(&schema);
5657 let mut create_table_req = CreateTableRequest::new();
5658 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5659 let result = namespace
5660 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
5661 .await;
5662 assert!(result.is_ok(), "Failed to create table in child namespace");
5663
5664 let list_req = ListTablesRequest {
5666 id: Some(vec!["test_ns".to_string()]),
5667 ..Default::default()
5668 };
5669 let result = namespace.list_tables(list_req).await;
5670 assert!(result.is_ok());
5671 let tables = result.unwrap().tables;
5672 assert_eq!(tables.len(), 1);
5673 assert_eq!(tables[0], "table1");
5674
5675 let mut exists_req = TableExistsRequest::new();
5677 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5678 let result = namespace.table_exists(exists_req).await;
5679 assert!(result.is_ok());
5680
5681 let mut describe_req = DescribeTableRequest::new();
5683 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5684 let result = namespace.describe_table(describe_req).await;
5685 assert!(result.is_ok());
5686 let response = result.unwrap();
5687 assert!(response.location.is_some());
5688 }
5689
5690 #[tokio::test]
5691 async fn test_multiple_tables_in_child_namespace() {
5692 let (namespace, _temp_dir) = create_test_namespace().await;
5693
5694 let mut create_ns_req = CreateNamespaceRequest::new();
5696 create_ns_req.id = Some(vec!["test_ns".to_string()]);
5697 namespace.create_namespace(create_ns_req).await.unwrap();
5698
5699 let schema = create_test_schema();
5701 let ipc_data = create_test_ipc_data(&schema);
5702 for i in 1..=3 {
5703 let mut create_table_req = CreateTableRequest::new();
5704 create_table_req.id = Some(vec!["test_ns".to_string(), format!("table{}", i)]);
5705 namespace
5706 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
5707 .await
5708 .unwrap();
5709 }
5710
5711 let list_req = ListTablesRequest {
5713 id: Some(vec!["test_ns".to_string()]),
5714 ..Default::default()
5715 };
5716 let result = namespace.list_tables(list_req).await;
5717 assert!(result.is_ok());
5718 let tables = result.unwrap().tables;
5719 assert_eq!(tables.len(), 3);
5720 assert!(tables.contains(&"table1".to_string()));
5721 assert!(tables.contains(&"table2".to_string()));
5722 assert!(tables.contains(&"table3".to_string()));
5723 }
5724
5725 #[tokio::test]
5726 async fn test_drop_table_in_child_namespace() {
5727 let (namespace, _temp_dir) = create_test_namespace().await;
5728
5729 let mut create_ns_req = CreateNamespaceRequest::new();
5731 create_ns_req.id = Some(vec!["test_ns".to_string()]);
5732 namespace.create_namespace(create_ns_req).await.unwrap();
5733
5734 let schema = create_test_schema();
5736 let ipc_data = create_test_ipc_data(&schema);
5737 let mut create_table_req = CreateTableRequest::new();
5738 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5739 namespace
5740 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
5741 .await
5742 .unwrap();
5743
5744 let mut drop_req = DropTableRequest::new();
5746 drop_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5747 let result = namespace.drop_table(drop_req).await;
5748 assert!(result.is_ok(), "Failed to drop table in child namespace");
5749
5750 let mut exists_req = TableExistsRequest::new();
5752 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5753 let result = namespace.table_exists(exists_req).await;
5754 assert!(result.is_err());
5755 }
5756
5757 #[tokio::test]
5758 async fn test_deeply_nested_namespace() {
5759 let (namespace, _temp_dir) = create_test_namespace().await;
5760
5761 let mut create_req = CreateNamespaceRequest::new();
5763 create_req.id = Some(vec!["level1".to_string()]);
5764 namespace.create_namespace(create_req).await.unwrap();
5765
5766 let mut create_req = CreateNamespaceRequest::new();
5767 create_req.id = Some(vec!["level1".to_string(), "level2".to_string()]);
5768 namespace.create_namespace(create_req).await.unwrap();
5769
5770 let mut create_req = CreateNamespaceRequest::new();
5771 create_req.id = Some(vec![
5772 "level1".to_string(),
5773 "level2".to_string(),
5774 "level3".to_string(),
5775 ]);
5776 namespace.create_namespace(create_req).await.unwrap();
5777
5778 let schema = create_test_schema();
5780 let ipc_data = create_test_ipc_data(&schema);
5781 let mut create_table_req = CreateTableRequest::new();
5782 create_table_req.id = Some(vec![
5783 "level1".to_string(),
5784 "level2".to_string(),
5785 "level3".to_string(),
5786 "table1".to_string(),
5787 ]);
5788 let result = namespace
5789 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
5790 .await;
5791 assert!(
5792 result.is_ok(),
5793 "Failed to create table in deeply nested namespace"
5794 );
5795
5796 let mut exists_req = TableExistsRequest::new();
5798 exists_req.id = Some(vec![
5799 "level1".to_string(),
5800 "level2".to_string(),
5801 "level3".to_string(),
5802 "table1".to_string(),
5803 ]);
5804 let result = namespace.table_exists(exists_req).await;
5805 assert!(result.is_ok());
5806 }
5807
5808 #[tokio::test]
5809 async fn test_namespace_with_properties() {
5810 let (namespace, _temp_dir) = create_test_namespace().await;
5811
5812 let mut properties = HashMap::new();
5814 properties.insert("owner".to_string(), "test_user".to_string());
5815 properties.insert("description".to_string(), "Test namespace".to_string());
5816
5817 let mut create_req = CreateNamespaceRequest::new();
5818 create_req.id = Some(vec!["test_ns".to_string()]);
5819 create_req.properties = Some(properties.clone());
5820 namespace.create_namespace(create_req).await.unwrap();
5821
5822 let describe_req = DescribeNamespaceRequest {
5824 id: Some(vec!["test_ns".to_string()]),
5825 ..Default::default()
5826 };
5827 let result = namespace.describe_namespace(describe_req).await;
5828 assert!(result.is_ok());
5829 let response = result.unwrap();
5830 assert!(response.properties.is_some());
5831 let props = response.properties.unwrap();
5832 assert_eq!(props.get("owner"), Some(&"test_user".to_string()));
5833 assert_eq!(
5834 props.get("description"),
5835 Some(&"Test namespace".to_string())
5836 );
5837 }
5838
5839 #[tokio::test]
5840 async fn test_cannot_drop_namespace_with_tables() {
5841 let (namespace, _temp_dir) = create_test_namespace().await;
5842
5843 let mut create_ns_req = CreateNamespaceRequest::new();
5845 create_ns_req.id = Some(vec!["test_ns".to_string()]);
5846 namespace.create_namespace(create_ns_req).await.unwrap();
5847
5848 let schema = create_test_schema();
5850 let ipc_data = create_test_ipc_data(&schema);
5851 let mut create_table_req = CreateTableRequest::new();
5852 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
5853 namespace
5854 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
5855 .await
5856 .unwrap();
5857
5858 let mut drop_req = DropNamespaceRequest::new();
5860 drop_req.id = Some(vec!["test_ns".to_string()]);
5861 let result = namespace.drop_namespace(drop_req).await;
5862 assert!(
5863 result.is_err(),
5864 "Should not be able to drop namespace with tables"
5865 );
5866 }
5867
5868 #[tokio::test]
5869 async fn test_isolation_between_namespaces() {
5870 let (namespace, _temp_dir) = create_test_namespace().await;
5871
5872 let mut create_req = CreateNamespaceRequest::new();
5874 create_req.id = Some(vec!["ns1".to_string()]);
5875 namespace.create_namespace(create_req).await.unwrap();
5876
5877 let mut create_req = CreateNamespaceRequest::new();
5878 create_req.id = Some(vec!["ns2".to_string()]);
5879 namespace.create_namespace(create_req).await.unwrap();
5880
5881 let schema = create_test_schema();
5883 let ipc_data = create_test_ipc_data(&schema);
5884
5885 let mut create_table_req = CreateTableRequest::new();
5886 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
5887 namespace
5888 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
5889 .await
5890 .unwrap();
5891
5892 let mut create_table_req = CreateTableRequest::new();
5893 create_table_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
5894 namespace
5895 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
5896 .await
5897 .unwrap();
5898
5899 let list_req = ListTablesRequest {
5901 id: Some(vec!["ns1".to_string()]),
5902 page_token: None,
5903 limit: None,
5904 ..Default::default()
5905 };
5906 let result = namespace.list_tables(list_req).await.unwrap();
5907 assert_eq!(result.tables.len(), 1);
5908 assert_eq!(result.tables[0], "table1");
5909
5910 let list_req = ListTablesRequest {
5911 id: Some(vec!["ns2".to_string()]),
5912 page_token: None,
5913 limit: None,
5914 ..Default::default()
5915 };
5916 let result = namespace.list_tables(list_req).await.unwrap();
5917 assert_eq!(result.tables.len(), 1);
5918 assert_eq!(result.tables[0], "table1");
5919
5920 let mut drop_req = DropTableRequest::new();
5922 drop_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
5923 namespace.drop_table(drop_req).await.unwrap();
5924
5925 let mut exists_req = TableExistsRequest::new();
5927 exists_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
5928 assert!(namespace.table_exists(exists_req).await.is_err());
5929
5930 let mut exists_req = TableExistsRequest::new();
5931 exists_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
5932 assert!(namespace.table_exists(exists_req).await.is_ok());
5933 }
5934
5935 #[tokio::test]
5936 async fn test_migrate_directory_tables() {
5937 let temp_dir = TempStdDir::default();
5938 let temp_path = temp_dir.to_str().unwrap();
5939
5940 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
5942 .manifest_enabled(false)
5943 .dir_listing_enabled(true)
5944 .build()
5945 .await
5946 .unwrap();
5947
5948 let schema = create_test_schema();
5950 let ipc_data = create_test_ipc_data(&schema);
5951
5952 for i in 1..=3 {
5953 let mut create_req = CreateTableRequest::new();
5954 create_req.id = Some(vec![format!("table{}", i)]);
5955 dir_only_ns
5956 .create_table(create_req, bytes::Bytes::from(ipc_data.clone()))
5957 .await
5958 .unwrap();
5959 }
5960
5961 drop(dir_only_ns);
5962
5963 let dual_mode_ns = DirectoryNamespaceBuilder::new(temp_path)
5965 .manifest_enabled(true)
5966 .dir_listing_enabled(true)
5967 .build()
5968 .await
5969 .unwrap();
5970
5971 let mut list_req = ListTablesRequest::new();
5973 list_req.id = Some(vec![]);
5974 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
5975 assert_eq!(tables.len(), 3);
5976
5977 let migrated_count = dual_mode_ns.migrate().await.unwrap();
5979 assert_eq!(migrated_count, 3, "Should migrate all 3 tables");
5980
5981 let mut list_req = ListTablesRequest::new();
5983 list_req.id = Some(vec![]);
5984 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
5985 assert_eq!(tables.len(), 3);
5986
5987 let migrated_count = dual_mode_ns.migrate().await.unwrap();
5989 assert_eq!(
5990 migrated_count, 0,
5991 "Should not migrate already-migrated tables"
5992 );
5993
5994 drop(dual_mode_ns);
5995
5996 let manifest_only_ns = DirectoryNamespaceBuilder::new(temp_path)
5998 .manifest_enabled(true)
5999 .dir_listing_enabled(false)
6000 .build()
6001 .await
6002 .unwrap();
6003
6004 let mut list_req = ListTablesRequest::new();
6006 list_req.id = Some(vec![]);
6007 let tables = manifest_only_ns.list_tables(list_req).await.unwrap().tables;
6008 assert_eq!(tables.len(), 3);
6009 assert!(tables.contains(&"table1".to_string()));
6010 assert!(tables.contains(&"table2".to_string()));
6011 assert!(tables.contains(&"table3".to_string()));
6012 }
6013
6014 #[tokio::test]
6015 async fn test_migrate_without_manifest() {
6016 let temp_dir = TempStdDir::default();
6017 let temp_path = temp_dir.to_str().unwrap();
6018
6019 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6021 .manifest_enabled(false)
6022 .dir_listing_enabled(true)
6023 .build()
6024 .await
6025 .unwrap();
6026
6027 let migrated_count = namespace.migrate().await.unwrap();
6029 assert_eq!(migrated_count, 0);
6030 }
6031
6032 #[tokio::test]
6033 async fn test_register_table() {
6034 use lance_namespace::models::{RegisterTableRequest, TableExistsRequest};
6035
6036 let temp_dir = TempStdDir::default();
6037 let temp_path = temp_dir.to_str().unwrap();
6038
6039 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6040 .dir_listing_to_manifest_migration_enabled(true)
6041 .build()
6042 .await
6043 .unwrap();
6044
6045 let schema = create_test_schema();
6047 let ipc_data = create_test_ipc_data(&schema);
6048
6049 let table_uri = format!("{}/external_table.lance", temp_path);
6050 let cursor = Cursor::new(ipc_data);
6051 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
6052 let batches: Vec<_> = stream_reader
6053 .collect::<std::result::Result<Vec<_>, _>>()
6054 .unwrap();
6055 let schema = batches[0].schema();
6056 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
6057 let reader = RecordBatchIterator::new(batch_results, schema);
6058 Dataset::write(Box::new(reader), &table_uri, None)
6059 .await
6060 .unwrap();
6061
6062 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
6064 register_req.id = Some(vec!["registered_table".to_string()]);
6065
6066 let response = namespace.register_table(register_req).await.unwrap();
6067 assert_eq!(response.location, Some("external_table.lance".to_string()));
6068
6069 let mut exists_req = TableExistsRequest::new();
6071 exists_req.id = Some(vec!["registered_table".to_string()]);
6072 assert!(namespace.table_exists(exists_req).await.is_ok());
6073
6074 let mut list_req = ListTablesRequest::new();
6076 list_req.id = Some(vec![]);
6077 let tables = namespace.list_tables(list_req).await.unwrap();
6078 assert!(tables.tables.contains(&"registered_table".to_string()));
6079 }
6080
6081 #[tokio::test]
6082 async fn test_register_table_duplicate_fails() {
6083 use lance_namespace::models::RegisterTableRequest;
6084
6085 let temp_dir = TempStdDir::default();
6086 let temp_path = temp_dir.to_str().unwrap();
6087
6088 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6089 .build()
6090 .await
6091 .unwrap();
6092
6093 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
6095 register_req.id = Some(vec!["test_table".to_string()]);
6096
6097 namespace
6098 .register_table(register_req.clone())
6099 .await
6100 .unwrap();
6101
6102 let result = namespace.register_table(register_req).await;
6104 assert!(result.is_err());
6105 assert!(result.unwrap_err().to_string().contains("already exists"));
6106 }
6107
6108 #[tokio::test]
6109 async fn test_deregister_table() {
6110 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
6111
6112 let temp_dir = TempStdDir::default();
6113 let temp_path = temp_dir.to_str().unwrap();
6114
6115 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6118 .manifest_enabled(true)
6119 .dir_listing_enabled(false)
6120 .build()
6121 .await
6122 .unwrap();
6123
6124 let schema = create_test_schema();
6126 let ipc_data = create_test_ipc_data(&schema);
6127
6128 let mut create_req = CreateTableRequest::new();
6129 create_req.id = Some(vec!["test_table".to_string()]);
6130 namespace
6131 .create_table(create_req, bytes::Bytes::from(ipc_data))
6132 .await
6133 .unwrap();
6134
6135 let mut exists_req = TableExistsRequest::new();
6137 exists_req.id = Some(vec!["test_table".to_string()]);
6138 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
6139
6140 let mut deregister_req = DeregisterTableRequest::new();
6142 deregister_req.id = Some(vec!["test_table".to_string()]);
6143 let response = namespace.deregister_table(deregister_req).await.unwrap();
6144
6145 assert!(
6147 response.location.is_some(),
6148 "Deregister should return location"
6149 );
6150 let location = response.location.as_ref().unwrap();
6151 let expected_url = lance_io::object_store::uri_to_url(temp_path)
6154 .expect("Failed to convert temp path to URL");
6155 let expected_prefix = expected_url.to_string();
6156 assert!(
6157 location.starts_with(&expected_prefix),
6158 "Location should start with '{}', got: {}",
6159 expected_prefix,
6160 location
6161 );
6162 assert!(
6163 location.contains("test_table"),
6164 "Location should contain table name: {}",
6165 location
6166 );
6167 assert_eq!(response.id, Some(vec!["test_table".to_string()]));
6168
6169 assert!(namespace.table_exists(exists_req).await.is_err());
6171
6172 let dataset = Dataset::open(location).await;
6174 assert!(
6175 dataset.is_ok(),
6176 "Physical table data should still exist at {}",
6177 location
6178 );
6179 }
6180
6181 #[tokio::test]
6182 async fn test_deregister_table_in_child_namespace() {
6183 use lance_namespace::models::{
6184 CreateNamespaceRequest, DeregisterTableRequest, TableExistsRequest,
6185 };
6186
6187 let temp_dir = TempStdDir::default();
6188 let temp_path = temp_dir.to_str().unwrap();
6189
6190 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6191 .build()
6192 .await
6193 .unwrap();
6194
6195 let mut create_ns_req = CreateNamespaceRequest::new();
6197 create_ns_req.id = Some(vec!["test_ns".to_string()]);
6198 namespace.create_namespace(create_ns_req).await.unwrap();
6199
6200 let schema = create_test_schema();
6202 let ipc_data = create_test_ipc_data(&schema);
6203
6204 let mut create_req = CreateTableRequest::new();
6205 create_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
6206 namespace
6207 .create_table(create_req, bytes::Bytes::from(ipc_data))
6208 .await
6209 .unwrap();
6210
6211 let mut deregister_req = DeregisterTableRequest::new();
6213 deregister_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
6214 let response = namespace.deregister_table(deregister_req).await.unwrap();
6215
6216 assert!(
6218 response.location.is_some(),
6219 "Deregister should return location"
6220 );
6221 let location = response.location.as_ref().unwrap();
6222 let expected_url = lance_io::object_store::uri_to_url(temp_path)
6225 .expect("Failed to convert temp path to URL");
6226 let expected_prefix = expected_url.to_string();
6227 assert!(
6228 location.starts_with(&expected_prefix),
6229 "Location should start with '{}', got: {}",
6230 expected_prefix,
6231 location
6232 );
6233 assert!(
6234 location.contains("test_ns") && location.contains("test_table"),
6235 "Location should contain namespace and table name: {}",
6236 location
6237 );
6238 assert_eq!(
6239 response.id,
6240 Some(vec!["test_ns".to_string(), "test_table".to_string()])
6241 );
6242
6243 let mut exists_req = TableExistsRequest::new();
6245 exists_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
6246 assert!(namespace.table_exists(exists_req).await.is_err());
6247 }
6248
6249 #[tokio::test]
6250 async fn test_register_without_manifest_fails() {
6251 use lance_namespace::models::RegisterTableRequest;
6252
6253 let temp_dir = TempStdDir::default();
6254 let temp_path = temp_dir.to_str().unwrap();
6255
6256 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6258 .manifest_enabled(false)
6259 .build()
6260 .await
6261 .unwrap();
6262
6263 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
6265 register_req.id = Some(vec!["test_table".to_string()]);
6266 let result = namespace.register_table(register_req).await;
6267 assert!(result.is_err());
6268 assert!(
6269 result
6270 .unwrap_err()
6271 .to_string()
6272 .contains("manifest mode is enabled")
6273 );
6274
6275 }
6278
6279 #[tokio::test]
6280 async fn test_register_table_rejects_absolute_uri() {
6281 use lance_namespace::models::RegisterTableRequest;
6282
6283 let temp_dir = TempStdDir::default();
6284 let temp_path = temp_dir.to_str().unwrap();
6285
6286 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6287 .build()
6288 .await
6289 .unwrap();
6290
6291 let mut register_req = RegisterTableRequest::new("s3://bucket/table.lance".to_string());
6293 register_req.id = Some(vec!["test_table".to_string()]);
6294 let result = namespace.register_table(register_req).await;
6295 assert!(result.is_err());
6296 let err_msg = result.unwrap_err().to_string();
6297 assert!(err_msg.contains("Absolute URIs are not allowed"));
6298 }
6299
6300 #[tokio::test]
6301 async fn test_register_table_rejects_absolute_path() {
6302 use lance_namespace::models::RegisterTableRequest;
6303
6304 let temp_dir = TempStdDir::default();
6305 let temp_path = temp_dir.to_str().unwrap();
6306
6307 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6308 .build()
6309 .await
6310 .unwrap();
6311
6312 let mut register_req = RegisterTableRequest::new("/tmp/table.lance".to_string());
6314 register_req.id = Some(vec!["test_table".to_string()]);
6315 let result = namespace.register_table(register_req).await;
6316 assert!(result.is_err());
6317 let err_msg = result.unwrap_err().to_string();
6318 assert!(err_msg.contains("Absolute paths are not allowed"));
6319 }
6320
6321 #[tokio::test]
6322 async fn test_register_table_rejects_path_traversal() {
6323 use lance_namespace::models::RegisterTableRequest;
6324
6325 let temp_dir = TempStdDir::default();
6326 let temp_path = temp_dir.to_str().unwrap();
6327
6328 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6329 .build()
6330 .await
6331 .unwrap();
6332
6333 let mut register_req = RegisterTableRequest::new("../outside/table.lance".to_string());
6335 register_req.id = Some(vec!["test_table".to_string()]);
6336 let result = namespace.register_table(register_req).await;
6337 assert!(result.is_err());
6338 let err_msg = result.unwrap_err().to_string();
6339 assert!(err_msg.contains("Path traversal is not allowed"));
6340 }
6341
6342 #[tokio::test]
6343 async fn test_namespace_write() {
6344 use arrow::array::Int32Array;
6345 use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
6346 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
6347 use lance::dataset::{Dataset, WriteMode, WriteParams};
6348 use lance_namespace::LanceNamespace;
6349
6350 let (namespace, _temp_dir) = create_test_namespace().await;
6351 let namespace = Arc::new(namespace) as Arc<dyn LanceNamespace>;
6352
6353 let table_id = vec!["test_ns".to_string(), "test_table".to_string()];
6355 let schema = Arc::new(ArrowSchema::new(vec![
6356 ArrowField::new("a", DataType::Int32, false),
6357 ArrowField::new("b", DataType::Int32, false),
6358 ]));
6359
6360 let data1 = RecordBatch::try_new(
6362 schema.clone(),
6363 vec![
6364 Arc::new(Int32Array::from(vec![1, 2, 3])),
6365 Arc::new(Int32Array::from(vec![10, 20, 30])),
6366 ],
6367 )
6368 .unwrap();
6369
6370 let reader1 = RecordBatchIterator::new(vec![data1].into_iter().map(Ok), schema.clone());
6371 let dataset =
6372 Dataset::write_into_namespace(reader1, namespace.clone(), table_id.clone(), None)
6373 .await
6374 .unwrap();
6375
6376 assert_eq!(dataset.count_rows(None).await.unwrap(), 3);
6377 assert_eq!(dataset.version().version, 1);
6378
6379 let data2 = RecordBatch::try_new(
6381 schema.clone(),
6382 vec![
6383 Arc::new(Int32Array::from(vec![4, 5])),
6384 Arc::new(Int32Array::from(vec![40, 50])),
6385 ],
6386 )
6387 .unwrap();
6388
6389 let params_append = WriteParams {
6390 mode: WriteMode::Append,
6391 ..Default::default()
6392 };
6393
6394 let reader2 = RecordBatchIterator::new(vec![data2].into_iter().map(Ok), schema.clone());
6395 let dataset = Dataset::write_into_namespace(
6396 reader2,
6397 namespace.clone(),
6398 table_id.clone(),
6399 Some(params_append),
6400 )
6401 .await
6402 .unwrap();
6403
6404 assert_eq!(dataset.count_rows(None).await.unwrap(), 5);
6405 assert_eq!(dataset.version().version, 2);
6406
6407 let data3 = RecordBatch::try_new(
6409 schema.clone(),
6410 vec![
6411 Arc::new(Int32Array::from(vec![100, 200])),
6412 Arc::new(Int32Array::from(vec![1000, 2000])),
6413 ],
6414 )
6415 .unwrap();
6416
6417 let params_overwrite = WriteParams {
6418 mode: WriteMode::Overwrite,
6419 ..Default::default()
6420 };
6421
6422 let reader3 = RecordBatchIterator::new(vec![data3].into_iter().map(Ok), schema.clone());
6423 let dataset = Dataset::write_into_namespace(
6424 reader3,
6425 namespace.clone(),
6426 table_id.clone(),
6427 Some(params_overwrite),
6428 )
6429 .await
6430 .unwrap();
6431
6432 assert_eq!(dataset.count_rows(None).await.unwrap(), 2);
6433 assert_eq!(dataset.version().version, 3);
6434
6435 let result = dataset.scan().try_into_batch().await.unwrap();
6437 let a_col = result
6438 .column_by_name("a")
6439 .unwrap()
6440 .as_any()
6441 .downcast_ref::<Int32Array>()
6442 .unwrap();
6443 assert_eq!(a_col.values(), &[100, 200]);
6444 }
6445
6446 #[tokio::test]
6451 async fn test_declare_table_v1_mode() {
6452 use lance_namespace::models::{
6453 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
6454 };
6455
6456 let temp_dir = TempStdDir::default();
6457 let temp_path = temp_dir.to_str().unwrap();
6458
6459 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6461 .manifest_enabled(false)
6462 .build()
6463 .await
6464 .unwrap();
6465
6466 let mut declare_req = DeclareTableRequest::new();
6468 declare_req.id = Some(vec!["test_table".to_string()]);
6469 let response = namespace.declare_table(declare_req).await.unwrap();
6470
6471 assert!(response.location.is_some());
6473 let location = response.location.as_ref().unwrap();
6474 assert!(location.ends_with("test_table.lance"));
6475
6476 let mut exists_req = TableExistsRequest::new();
6478 exists_req.id = Some(vec!["test_table".to_string()]);
6479 assert!(namespace.table_exists(exists_req).await.is_ok());
6480
6481 let mut describe_req = DescribeTableRequest::new();
6483 describe_req.id = Some(vec!["test_table".to_string()]);
6484 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6485 assert!(describe_response.location.is_some());
6486 assert!(describe_response.version.is_none()); assert!(describe_response.schema.is_none()); assert_eq!(describe_response.is_only_declared, None);
6489
6490 let mut describe_req = DescribeTableRequest::new();
6491 describe_req.id = Some(vec!["test_table".to_string()]);
6492 describe_req.check_declared = Some(true);
6493 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6494 assert_eq!(describe_response.is_only_declared, Some(true));
6495
6496 let mut list_req = ListTablesRequest::new();
6497 list_req.id = Some(vec![]);
6498 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
6499 assert_eq!(list_response.tables, vec!["test_table".to_string()]);
6500
6501 list_req.include_declared = Some(false);
6502 let list_response = namespace.list_tables(list_req).await.unwrap();
6503 assert!(list_response.tables.is_empty());
6504 }
6505
6506 #[tokio::test]
6507 async fn test_insert_into_declared_table_promotes_it_from_declared_state() {
6508 use lance_namespace::models::{
6509 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest,
6510 };
6511
6512 let temp_dir = TempStdDir::default();
6513 let temp_path = temp_dir.to_str().unwrap();
6514
6515 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6516 .manifest_enabled(false)
6517 .build()
6518 .await
6519 .unwrap();
6520
6521 let mut declare_req = DeclareTableRequest::new();
6522 declare_req.id = Some(vec!["test_table".to_string()]);
6523 namespace.declare_table(declare_req).await.unwrap();
6524
6525 let schema = create_test_schema();
6526 let ipc_data = create_test_ipc_data(&schema);
6527 let mut insert_req = InsertIntoTableRequest::new();
6528 insert_req.id = Some(vec!["test_table".to_string()]);
6529 namespace
6530 .insert_into_table(insert_req, bytes::Bytes::from(ipc_data))
6531 .await
6532 .unwrap();
6533
6534 let mut describe_req = DescribeTableRequest::new();
6535 describe_req.id = Some(vec!["test_table".to_string()]);
6536 describe_req.load_detailed_metadata = Some(true);
6537 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6538
6539 assert_eq!(describe_response.is_only_declared, Some(false));
6540 assert_eq!(describe_response.version, Some(1));
6541 assert!(describe_response.schema.is_some());
6542
6543 let mut list_req = ListTablesRequest::new();
6544 list_req.id = Some(vec![]);
6545 list_req.include_declared = Some(false);
6546 assert_eq!(
6547 namespace.list_tables(list_req).await.unwrap().tables,
6548 vec!["test_table".to_string()]
6549 );
6550 }
6551
6552 #[tokio::test]
6553 async fn test_create_table_after_declare_table_v1_mode_creates_table() {
6554 use lance_namespace::models::{
6555 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
6556 };
6557
6558 let temp_dir = TempStdDir::default();
6559 let temp_path = temp_dir.to_str().unwrap();
6560
6561 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6562 .manifest_enabled(false)
6563 .build()
6564 .await
6565 .unwrap();
6566
6567 let mut declare_req = DeclareTableRequest::new();
6568 declare_req.id = Some(vec!["test_table".to_string()]);
6569 namespace.declare_table(declare_req).await.unwrap();
6570
6571 let mut create_req = CreateTableRequest::new();
6572 create_req.id = Some(vec!["test_table".to_string()]);
6573 let response = namespace
6574 .create_table(
6575 create_req,
6576 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6577 )
6578 .await
6579 .unwrap();
6580
6581 assert_eq!(response.version, Some(1));
6582
6583 let mut describe_req = DescribeTableRequest::new();
6584 describe_req.id = Some(vec!["test_table".to_string()]);
6585 describe_req.load_detailed_metadata = Some(true);
6586 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6587 assert_eq!(describe_response.is_only_declared, Some(false));
6588 assert_eq!(describe_response.version, Some(1));
6589
6590 let mut list_req = ListTablesRequest::new();
6591 list_req.id = Some(vec![]);
6592 list_req.include_declared = Some(false);
6593 assert_eq!(
6594 namespace.list_tables(list_req).await.unwrap().tables,
6595 vec!["test_table".to_string()]
6596 );
6597 }
6598
6599 #[tokio::test]
6600 async fn test_insert_into_declared_table_with_manifest_promotes_it() {
6601 use lance_namespace::models::{
6602 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest, ListTablesRequest,
6603 };
6604
6605 let temp_dir = TempStdDir::default();
6606 let temp_path = temp_dir.to_str().unwrap();
6607
6608 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6609 .manifest_enabled(true)
6610 .dir_listing_enabled(false)
6611 .build()
6612 .await
6613 .unwrap();
6614
6615 let mut declare_req = DeclareTableRequest::new();
6616 declare_req.id = Some(vec!["test_table".to_string()]);
6617 namespace.declare_table(declare_req).await.unwrap();
6618
6619 let mut insert_req = InsertIntoTableRequest::new();
6620 insert_req.id = Some(vec!["test_table".to_string()]);
6621 namespace
6622 .insert_into_table(
6623 insert_req,
6624 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6625 )
6626 .await
6627 .unwrap();
6628
6629 let mut describe_req = DescribeTableRequest::new();
6630 describe_req.id = Some(vec!["test_table".to_string()]);
6631 describe_req.load_detailed_metadata = Some(true);
6632 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6633 assert_eq!(describe_response.is_only_declared, Some(false));
6634 assert_eq!(describe_response.version, Some(1));
6635
6636 let mut list_req = ListTablesRequest::new();
6637 list_req.id = Some(vec![]);
6638 list_req.include_declared = Some(false);
6639 assert_eq!(
6640 namespace.list_tables(list_req).await.unwrap().tables,
6641 vec!["test_table".to_string()]
6642 );
6643 }
6644
6645 #[tokio::test]
6646 async fn test_create_table_after_declare_table_with_manifest_creates_table() {
6647 use lance_namespace::models::{
6648 CreateTableRequest, DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
6649 };
6650
6651 let temp_dir = TempStdDir::default();
6652 let temp_path = temp_dir.to_str().unwrap();
6653
6654 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6655 .manifest_enabled(true)
6656 .dir_listing_enabled(false)
6657 .build()
6658 .await
6659 .unwrap();
6660
6661 let mut declare_req = DeclareTableRequest::new();
6662 declare_req.id = Some(vec!["test_table".to_string()]);
6663 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
6664 namespace.declare_table(declare_req).await.unwrap();
6665
6666 let mut create_req = CreateTableRequest::new();
6667 create_req.id = Some(vec!["test_table".to_string()]);
6668 create_req.mode = Some("Overwrite".to_string());
6669 let response = namespace
6670 .create_table(
6671 create_req,
6672 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6673 )
6674 .await
6675 .unwrap();
6676
6677 assert_eq!(response.version, Some(1));
6678 assert_eq!(
6679 response
6680 .properties
6681 .as_ref()
6682 .and_then(|properties| properties.get("owner")),
6683 Some(&"alice".to_string())
6684 );
6685
6686 let mut describe_req = DescribeTableRequest::new();
6687 describe_req.id = Some(vec!["test_table".to_string()]);
6688 describe_req.load_detailed_metadata = Some(true);
6689 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6690 assert_eq!(describe_response.is_only_declared, Some(false));
6691 assert_eq!(describe_response.version, Some(1));
6692 assert_eq!(
6693 describe_response
6694 .properties
6695 .as_ref()
6696 .and_then(|properties| properties.get("owner")),
6697 Some(&"alice".to_string())
6698 );
6699
6700 let mut list_req = ListTablesRequest::new();
6701 list_req.id = Some(vec![]);
6702 list_req.include_declared = Some(false);
6703 assert_eq!(
6704 namespace.list_tables(list_req).await.unwrap().tables,
6705 vec!["test_table".to_string()]
6706 );
6707 }
6708
6709 #[tokio::test]
6710 async fn test_create_table_after_declare_table_with_manifest_rejects_new_properties() {
6711 use lance_namespace::models::{CreateTableRequest, DeclareTableRequest};
6712
6713 let temp_dir = TempStdDir::default();
6714 let temp_path = temp_dir.to_str().unwrap();
6715
6716 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6717 .manifest_enabled(true)
6718 .dir_listing_enabled(false)
6719 .build()
6720 .await
6721 .unwrap();
6722
6723 let mut declare_req = DeclareTableRequest::new();
6724 declare_req.id = Some(vec!["test_table".to_string()]);
6725 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
6726 namespace.declare_table(declare_req).await.unwrap();
6727
6728 let mut create_req = CreateTableRequest::new();
6729 create_req.id = Some(vec!["test_table".to_string()]);
6730 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
6731
6732 let result = namespace
6733 .create_table(
6734 create_req,
6735 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6736 )
6737 .await;
6738
6739 assert!(result.is_err());
6740 assert!(
6741 result
6742 .unwrap_err()
6743 .to_string()
6744 .contains("cannot set properties for already declared table")
6745 );
6746 }
6747
6748 #[tokio::test]
6749 async fn test_create_table_with_manifest_exist_ok_keeps_existing_table() {
6750 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
6751
6752 let temp_dir = TempStdDir::default();
6753 let temp_path = temp_dir.to_str().unwrap();
6754
6755 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6756 .manifest_enabled(true)
6757 .dir_listing_enabled(false)
6758 .build()
6759 .await
6760 .unwrap();
6761
6762 let mut create_req = CreateTableRequest::new();
6763 create_req.id = Some(vec!["test_table".to_string()]);
6764 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
6765 namespace
6766 .create_table(
6767 create_req,
6768 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6769 )
6770 .await
6771 .unwrap();
6772
6773 let mut create_req = CreateTableRequest::new();
6774 create_req.id = Some(vec!["test_table".to_string()]);
6775 create_req.mode = Some("ExistOk".to_string());
6776 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
6777 let response = namespace
6778 .create_table(
6779 create_req,
6780 bytes::Bytes::from(create_single_row_test_ipc_data()),
6781 )
6782 .await
6783 .unwrap();
6784
6785 assert_eq!(
6786 response
6787 .properties
6788 .as_ref()
6789 .and_then(|properties| properties.get("owner")),
6790 Some(&"alice".to_string())
6791 );
6792 assert_eq!(
6793 open_dataset(&namespace, "test_table")
6794 .await
6795 .count_rows(None)
6796 .await
6797 .unwrap(),
6798 2
6799 );
6800
6801 let mut describe_req = DescribeTableRequest::new();
6802 describe_req.id = Some(vec!["test_table".to_string()]);
6803 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6804 assert_eq!(
6805 describe_response
6806 .properties
6807 .as_ref()
6808 .and_then(|properties| properties.get("owner")),
6809 Some(&"alice".to_string())
6810 );
6811 }
6812
6813 #[tokio::test]
6814 async fn test_create_table_with_manifest_overwrite_replaces_existing_table() {
6815 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
6816
6817 let temp_dir = TempStdDir::default();
6818 let temp_path = temp_dir.to_str().unwrap();
6819
6820 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6821 .manifest_enabled(true)
6822 .dir_listing_enabled(false)
6823 .build()
6824 .await
6825 .unwrap();
6826
6827 let mut create_req = CreateTableRequest::new();
6828 create_req.id = Some(vec!["test_table".to_string()]);
6829 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
6830 namespace
6831 .create_table(
6832 create_req,
6833 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6834 )
6835 .await
6836 .unwrap();
6837
6838 let mut create_req = CreateTableRequest::new();
6839 create_req.id = Some(vec!["test_table".to_string()]);
6840 create_req.mode = Some("overwrite".to_string());
6841 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
6842 let response = namespace
6843 .create_table(
6844 create_req,
6845 bytes::Bytes::from(create_single_row_test_ipc_data()),
6846 )
6847 .await
6848 .unwrap();
6849
6850 assert_eq!(response.version, Some(2));
6851 assert_eq!(
6852 response
6853 .properties
6854 .as_ref()
6855 .and_then(|properties| properties.get("owner")),
6856 Some(&"bob".to_string())
6857 );
6858 assert_eq!(
6859 open_dataset(&namespace, "test_table")
6860 .await
6861 .count_rows(None)
6862 .await
6863 .unwrap(),
6864 1
6865 );
6866
6867 let mut describe_req = DescribeTableRequest::new();
6868 describe_req.id = Some(vec!["test_table".to_string()]);
6869 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6870 assert_eq!(
6871 describe_response
6872 .properties
6873 .as_ref()
6874 .and_then(|properties| properties.get("owner")),
6875 Some(&"bob".to_string())
6876 );
6877 }
6878
6879 #[tokio::test]
6880 async fn test_create_table_with_manifest_invalid_mode_rejected() {
6881 use lance_namespace::models::CreateTableRequest;
6882
6883 let temp_dir = TempStdDir::default();
6884 let temp_path = temp_dir.to_str().unwrap();
6885
6886 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6887 .manifest_enabled(true)
6888 .dir_listing_enabled(false)
6889 .build()
6890 .await
6891 .unwrap();
6892
6893 let mut create_req = CreateTableRequest::new();
6894 create_req.id = Some(vec!["test_table".to_string()]);
6895 create_req.mode = Some("append".to_string());
6896 let result = namespace
6897 .create_table(
6898 create_req,
6899 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6900 )
6901 .await;
6902
6903 assert!(result.is_err());
6904 assert!(
6905 result
6906 .unwrap_err()
6907 .to_string()
6908 .contains("Unsupported create_table mode")
6909 );
6910 }
6911
6912 #[tokio::test]
6913 async fn test_merge_insert_into_declared_table_v1_mode_creates_table() {
6914 use lance_namespace::models::{
6915 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
6916 MergeInsertIntoTableRequest,
6917 };
6918
6919 let temp_dir = TempStdDir::default();
6920 let temp_path = temp_dir.to_str().unwrap();
6921
6922 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6923 .manifest_enabled(false)
6924 .build()
6925 .await
6926 .unwrap();
6927
6928 let mut declare_req = DeclareTableRequest::new();
6929 declare_req.id = Some(vec!["test_table".to_string()]);
6930 namespace.declare_table(declare_req).await.unwrap();
6931
6932 let mut merge_req = MergeInsertIntoTableRequest::new();
6933 merge_req.id = Some(vec!["test_table".to_string()]);
6934 merge_req.on = Some("id".to_string());
6935 let response = namespace
6936 .merge_insert_into_table(
6937 merge_req,
6938 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6939 )
6940 .await
6941 .unwrap();
6942
6943 assert_eq!(response.num_inserted_rows, Some(2));
6944 assert_eq!(response.num_updated_rows, Some(0));
6945
6946 let mut describe_req = DescribeTableRequest::new();
6947 describe_req.id = Some(vec!["test_table".to_string()]);
6948 describe_req.load_detailed_metadata = Some(true);
6949 let describe_response = namespace.describe_table(describe_req).await.unwrap();
6950 assert_eq!(describe_response.is_only_declared, Some(false));
6951 assert_eq!(describe_response.version, Some(1));
6952
6953 let mut list_req = ListTablesRequest::new();
6954 list_req.id = Some(vec![]);
6955 list_req.include_declared = Some(false);
6956 assert_eq!(
6957 namespace.list_tables(list_req).await.unwrap().tables,
6958 vec!["test_table".to_string()]
6959 );
6960 }
6961
6962 #[tokio::test]
6963 async fn test_merge_insert_into_declared_table_with_manifest_creates_table() {
6964 use lance_namespace::models::{
6965 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
6966 MergeInsertIntoTableRequest,
6967 };
6968
6969 let temp_dir = TempStdDir::default();
6970 let temp_path = temp_dir.to_str().unwrap();
6971
6972 let namespace = DirectoryNamespaceBuilder::new(temp_path)
6973 .manifest_enabled(true)
6974 .dir_listing_enabled(false)
6975 .build()
6976 .await
6977 .unwrap();
6978
6979 let mut declare_req = DeclareTableRequest::new();
6980 declare_req.id = Some(vec!["test_table".to_string()]);
6981 namespace.declare_table(declare_req).await.unwrap();
6982
6983 let mut merge_req = MergeInsertIntoTableRequest::new();
6984 merge_req.id = Some(vec!["test_table".to_string()]);
6985 merge_req.on = Some("id".to_string());
6986 let response = namespace
6987 .merge_insert_into_table(
6988 merge_req,
6989 bytes::Bytes::from(create_non_empty_test_ipc_data()),
6990 )
6991 .await
6992 .unwrap();
6993
6994 assert_eq!(response.num_inserted_rows, Some(2));
6995 assert_eq!(response.num_updated_rows, Some(0));
6996
6997 let mut describe_req = DescribeTableRequest::new();
6998 describe_req.id = Some(vec!["test_table".to_string()]);
6999 describe_req.load_detailed_metadata = Some(true);
7000 let describe_response = namespace.describe_table(describe_req).await.unwrap();
7001 assert_eq!(describe_response.is_only_declared, Some(false));
7002 assert_eq!(describe_response.version, Some(1));
7003
7004 let mut list_req = ListTablesRequest::new();
7005 list_req.id = Some(vec![]);
7006 list_req.include_declared = Some(false);
7007 assert_eq!(
7008 namespace.list_tables(list_req).await.unwrap().tables,
7009 vec!["test_table".to_string()]
7010 );
7011 }
7012
7013 #[tokio::test]
7014 async fn test_declare_table_with_manifest() {
7015 use lance_namespace::models::{
7016 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
7017 };
7018
7019 let temp_dir = TempStdDir::default();
7020 let temp_path = temp_dir.to_str().unwrap();
7021
7022 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7024 .manifest_enabled(true)
7025 .dir_listing_enabled(false)
7026 .build()
7027 .await
7028 .unwrap();
7029
7030 let mut declare_req = DeclareTableRequest::new();
7032 declare_req.id = Some(vec!["test_table".to_string()]);
7033 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
7034 let response = namespace.declare_table(declare_req).await.unwrap();
7035
7036 assert!(response.location.is_some());
7038 assert_eq!(
7039 response
7040 .properties
7041 .as_ref()
7042 .and_then(|properties| properties.get("owner")),
7043 Some(&"alice".to_string())
7044 );
7045
7046 let mut exists_req = TableExistsRequest::new();
7048 exists_req.id = Some(vec!["test_table".to_string()]);
7049 assert!(namespace.table_exists(exists_req).await.is_ok());
7050
7051 let mut describe_req = DescribeTableRequest::new();
7052 describe_req.id = Some(vec!["test_table".to_string()]);
7053 let describe_response = namespace.describe_table(describe_req).await.unwrap();
7054 assert_eq!(describe_response.is_only_declared, None);
7055
7056 let mut describe_req = DescribeTableRequest::new();
7057 describe_req.id = Some(vec!["test_table".to_string()]);
7058 describe_req.check_declared = Some(true);
7059 let describe_response = namespace.describe_table(describe_req).await.unwrap();
7060 assert_eq!(describe_response.is_only_declared, Some(true));
7061 assert_eq!(
7062 describe_response
7063 .properties
7064 .as_ref()
7065 .and_then(|properties| properties.get("owner")),
7066 Some(&"alice".to_string())
7067 );
7068
7069 let mut list_req = ListTablesRequest::new();
7070 list_req.id = Some(vec![]);
7071 assert_eq!(
7072 namespace
7073 .list_tables(list_req.clone())
7074 .await
7075 .unwrap()
7076 .tables,
7077 vec!["test_table".to_string()]
7078 );
7079 list_req.include_declared = Some(false);
7080 assert!(
7081 namespace
7082 .list_tables(list_req)
7083 .await
7084 .unwrap()
7085 .tables
7086 .is_empty()
7087 );
7088 }
7089
7090 #[tokio::test]
7091 async fn test_declare_table_when_table_exists() {
7092 use lance_namespace::models::DeclareTableRequest;
7093
7094 let temp_dir = TempStdDir::default();
7095 let temp_path = temp_dir.to_str().unwrap();
7096
7097 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7098 .manifest_enabled(false)
7099 .build()
7100 .await
7101 .unwrap();
7102
7103 let schema = create_test_schema();
7105 let ipc_data = create_test_ipc_data(&schema);
7106 let mut create_req = CreateTableRequest::new();
7107 create_req.id = Some(vec!["test_table".to_string()]);
7108 namespace
7109 .create_table(create_req, bytes::Bytes::from(ipc_data))
7110 .await
7111 .unwrap();
7112
7113 let mut declare_req = DeclareTableRequest::new();
7115 declare_req.id = Some(vec!["test_table".to_string()]);
7116 let result = namespace.declare_table(declare_req).await;
7117 assert!(result.is_err());
7118 }
7119
7120 #[tokio::test]
7125 async fn test_deregister_table_v1_mode() {
7126 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
7127
7128 let temp_dir = TempStdDir::default();
7129 let temp_path = temp_dir.to_str().unwrap();
7130
7131 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7133 .manifest_enabled(false)
7134 .dir_listing_enabled(true)
7135 .build()
7136 .await
7137 .unwrap();
7138
7139 let schema = create_test_schema();
7141 let ipc_data = create_test_ipc_data(&schema);
7142 let mut create_req = CreateTableRequest::new();
7143 create_req.id = Some(vec!["test_table".to_string()]);
7144 namespace
7145 .create_table(create_req, bytes::Bytes::from(ipc_data))
7146 .await
7147 .unwrap();
7148
7149 let mut exists_req = TableExistsRequest::new();
7151 exists_req.id = Some(vec!["test_table".to_string()]);
7152 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
7153
7154 let mut deregister_req = DeregisterTableRequest::new();
7156 deregister_req.id = Some(vec!["test_table".to_string()]);
7157 let response = namespace.deregister_table(deregister_req).await.unwrap();
7158
7159 assert!(response.location.is_some());
7161 let location = response.location.as_ref().unwrap();
7162 assert!(location.contains("test_table"));
7163
7164 let result = namespace.table_exists(exists_req).await;
7166 assert!(result.is_err());
7167 assert!(result.unwrap_err().to_string().contains("deregistered"));
7168
7169 let dataset = Dataset::open(location).await;
7171 assert!(dataset.is_ok(), "Physical table data should still exist");
7172 }
7173
7174 #[tokio::test]
7175 async fn test_deregister_table_v1_already_deregistered() {
7176 use lance_namespace::models::DeregisterTableRequest;
7177
7178 let temp_dir = TempStdDir::default();
7179 let temp_path = temp_dir.to_str().unwrap();
7180
7181 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7182 .manifest_enabled(false)
7183 .dir_listing_enabled(true)
7184 .build()
7185 .await
7186 .unwrap();
7187
7188 let schema = create_test_schema();
7190 let ipc_data = create_test_ipc_data(&schema);
7191 let mut create_req = CreateTableRequest::new();
7192 create_req.id = Some(vec!["test_table".to_string()]);
7193 namespace
7194 .create_table(create_req, bytes::Bytes::from(ipc_data))
7195 .await
7196 .unwrap();
7197
7198 let mut deregister_req = DeregisterTableRequest::new();
7200 deregister_req.id = Some(vec!["test_table".to_string()]);
7201 namespace
7202 .deregister_table(deregister_req.clone())
7203 .await
7204 .unwrap();
7205
7206 let result = namespace.deregister_table(deregister_req).await;
7208 assert!(result.is_err());
7209 assert!(
7210 result
7211 .unwrap_err()
7212 .to_string()
7213 .contains("already deregistered")
7214 );
7215 }
7216
7217 #[tokio::test]
7222 async fn test_list_tables_skips_deregistered_v1() {
7223 use lance_namespace::models::DeregisterTableRequest;
7224
7225 let temp_dir = TempStdDir::default();
7226 let temp_path = temp_dir.to_str().unwrap();
7227
7228 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7229 .manifest_enabled(false)
7230 .dir_listing_enabled(true)
7231 .build()
7232 .await
7233 .unwrap();
7234
7235 let schema = create_test_schema();
7237 let ipc_data = create_test_ipc_data(&schema);
7238
7239 let mut create_req1 = CreateTableRequest::new();
7240 create_req1.id = Some(vec!["table1".to_string()]);
7241 namespace
7242 .create_table(create_req1, bytes::Bytes::from(ipc_data.clone()))
7243 .await
7244 .unwrap();
7245
7246 let mut create_req2 = CreateTableRequest::new();
7247 create_req2.id = Some(vec!["table2".to_string()]);
7248 namespace
7249 .create_table(create_req2, bytes::Bytes::from(ipc_data))
7250 .await
7251 .unwrap();
7252
7253 let mut list_req = ListTablesRequest::new();
7255 list_req.id = Some(vec![]);
7256 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
7257 assert_eq!(list_response.tables.len(), 2);
7258
7259 let mut deregister_req = DeregisterTableRequest::new();
7261 deregister_req.id = Some(vec!["table1".to_string()]);
7262 namespace.deregister_table(deregister_req).await.unwrap();
7263
7264 let list_response = namespace.list_tables(list_req).await.unwrap();
7266 assert_eq!(list_response.tables.len(), 1);
7267 assert!(list_response.tables.contains(&"table2".to_string()));
7268 assert!(!list_response.tables.contains(&"table1".to_string()));
7269 }
7270
7271 #[tokio::test]
7276 async fn test_describe_table_fails_for_deregistered_v1() {
7277 use lance_namespace::models::{DeregisterTableRequest, DescribeTableRequest};
7278
7279 let temp_dir = TempStdDir::default();
7280 let temp_path = temp_dir.to_str().unwrap();
7281
7282 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7283 .manifest_enabled(false)
7284 .dir_listing_enabled(true)
7285 .build()
7286 .await
7287 .unwrap();
7288
7289 let schema = create_test_schema();
7291 let ipc_data = create_test_ipc_data(&schema);
7292 let mut create_req = CreateTableRequest::new();
7293 create_req.id = Some(vec!["test_table".to_string()]);
7294 namespace
7295 .create_table(create_req, bytes::Bytes::from(ipc_data))
7296 .await
7297 .unwrap();
7298
7299 let mut describe_req = DescribeTableRequest::new();
7301 describe_req.id = Some(vec!["test_table".to_string()]);
7302 assert!(namespace.describe_table(describe_req.clone()).await.is_ok());
7303
7304 let mut deregister_req = DeregisterTableRequest::new();
7306 deregister_req.id = Some(vec!["test_table".to_string()]);
7307 namespace.deregister_table(deregister_req).await.unwrap();
7308
7309 let result = namespace.describe_table(describe_req).await;
7311 assert!(result.is_err());
7312 let err = result.unwrap_err();
7313 assert!(matches!(err, Error::Namespace { .. }));
7314 let err_msg = err.to_string();
7315 assert!(err_msg.contains("deregistered"));
7316 assert!(err_msg.contains("table id 'test_table'"));
7317 }
7318
7319 #[tokio::test]
7320 async fn test_table_exists_fails_for_deregistered_v1() {
7321 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
7322
7323 let temp_dir = TempStdDir::default();
7324 let temp_path = temp_dir.to_str().unwrap();
7325
7326 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7327 .manifest_enabled(false)
7328 .dir_listing_enabled(true)
7329 .build()
7330 .await
7331 .unwrap();
7332
7333 let schema = create_test_schema();
7335 let ipc_data = create_test_ipc_data(&schema);
7336 let mut create_req = CreateTableRequest::new();
7337 create_req.id = Some(vec!["test_table".to_string()]);
7338 namespace
7339 .create_table(create_req, bytes::Bytes::from(ipc_data))
7340 .await
7341 .unwrap();
7342
7343 let mut exists_req = TableExistsRequest::new();
7345 exists_req.id = Some(vec!["test_table".to_string()]);
7346 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
7347
7348 let mut deregister_req = DeregisterTableRequest::new();
7350 deregister_req.id = Some(vec!["test_table".to_string()]);
7351 namespace.deregister_table(deregister_req).await.unwrap();
7352
7353 let result = namespace.table_exists(exists_req).await;
7355 assert!(result.is_err());
7356 let err = result.unwrap_err();
7357 assert!(matches!(err, Error::Namespace { .. }));
7358 let err_msg = err.to_string();
7359 assert!(err_msg.contains("deregistered"));
7360 assert!(err_msg.contains("table id 'test_table'"));
7361 }
7362
7363 #[tokio::test]
7364 async fn test_atomic_table_status_check() {
7365 let temp_dir = TempStdDir::default();
7369 let temp_path = temp_dir.to_str().unwrap();
7370
7371 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7372 .manifest_enabled(false)
7373 .dir_listing_enabled(true)
7374 .build()
7375 .await
7376 .unwrap();
7377
7378 let schema = create_test_schema();
7380 let ipc_data = create_test_ipc_data(&schema);
7381 let mut create_req = CreateTableRequest::new();
7382 create_req.id = Some(vec!["test_table".to_string()]);
7383 namespace
7384 .create_table(create_req, bytes::Bytes::from(ipc_data))
7385 .await
7386 .unwrap();
7387
7388 let status = namespace.check_table_status("test_table").await;
7390 assert!(status.exists);
7391 assert!(!status.is_deregistered);
7392 assert!(!status.has_reserved_file);
7393 }
7394
7395 #[tokio::test]
7396 async fn test_table_version_tracking_enabled_managed_versioning() {
7397 use lance_namespace::models::DescribeTableRequest;
7398
7399 let temp_dir = TempStdDir::default();
7400 let temp_path = temp_dir.to_str().unwrap();
7401
7402 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7404 .table_version_tracking_enabled(true)
7405 .build()
7406 .await
7407 .unwrap();
7408
7409 let schema = create_test_schema();
7411 let ipc_data = create_test_ipc_data(&schema);
7412 let mut create_req = CreateTableRequest::new();
7413 create_req.id = Some(vec!["test_table".to_string()]);
7414 namespace
7415 .create_table(create_req, bytes::Bytes::from(ipc_data))
7416 .await
7417 .unwrap();
7418
7419 let mut describe_req = DescribeTableRequest::new();
7421 describe_req.id = Some(vec!["test_table".to_string()]);
7422 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
7423
7424 assert_eq!(
7426 describe_resp.managed_versioning,
7427 Some(true),
7428 "managed_versioning should be true when table_version_tracking_enabled=true"
7429 );
7430 }
7431
7432 #[tokio::test]
7433 async fn test_table_version_tracking_disabled_no_managed_versioning() {
7434 use lance_namespace::models::DescribeTableRequest;
7435
7436 let temp_dir = TempStdDir::default();
7437 let temp_path = temp_dir.to_str().unwrap();
7438
7439 let namespace = DirectoryNamespaceBuilder::new(temp_path)
7441 .table_version_tracking_enabled(false)
7442 .build()
7443 .await
7444 .unwrap();
7445
7446 let schema = create_test_schema();
7448 let ipc_data = create_test_ipc_data(&schema);
7449 let mut create_req = CreateTableRequest::new();
7450 create_req.id = Some(vec!["test_table".to_string()]);
7451 namespace
7452 .create_table(create_req, bytes::Bytes::from(ipc_data))
7453 .await
7454 .unwrap();
7455
7456 let mut describe_req = DescribeTableRequest::new();
7458 describe_req.id = Some(vec!["test_table".to_string()]);
7459 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
7460
7461 assert!(
7463 describe_resp.managed_versioning.is_none(),
7464 "managed_versioning should be None when table_version_tracking_enabled=false, got: {:?}",
7465 describe_resp.managed_versioning
7466 );
7467 }
7468
7469 #[tokio::test]
7470 async fn test_list_table_versions() {
7471 use arrow::array::{Int32Array, RecordBatchIterator};
7472 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7473 use arrow::record_batch::RecordBatch;
7474 use lance::dataset::{Dataset, WriteMode, WriteParams};
7475 use lance_namespace::models::{CreateNamespaceRequest, ListTableVersionsRequest};
7476
7477 let temp_dir = TempStrDir::default();
7478 let temp_path: &str = &temp_dir;
7479
7480 let namespace: Arc<dyn LanceNamespace> = Arc::new(
7481 DirectoryNamespaceBuilder::new(temp_path)
7482 .table_version_tracking_enabled(true)
7483 .build()
7484 .await
7485 .unwrap(),
7486 );
7487
7488 let mut create_ns_req = CreateNamespaceRequest::new();
7490 create_ns_req.id = Some(vec!["workspace".to_string()]);
7491 namespace.create_namespace(create_ns_req).await.unwrap();
7492
7493 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
7495 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
7496 "id",
7497 DataType::Int32,
7498 false,
7499 )]));
7500 let batch = RecordBatch::try_new(
7501 arrow_schema.clone(),
7502 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
7503 )
7504 .unwrap();
7505 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
7506 let write_params = WriteParams {
7507 mode: WriteMode::Create,
7508 ..Default::default()
7509 };
7510 let mut dataset = Dataset::write_into_namespace(
7511 batches,
7512 namespace.clone(),
7513 table_id.clone(),
7514 Some(write_params),
7515 )
7516 .await
7517 .unwrap();
7518
7519 let batch2 = RecordBatch::try_new(
7521 arrow_schema.clone(),
7522 vec![Arc::new(Int32Array::from(vec![100, 200]))],
7523 )
7524 .unwrap();
7525 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
7526 dataset.append(batches, None).await.unwrap();
7527
7528 let batch3 = RecordBatch::try_new(
7530 arrow_schema.clone(),
7531 vec![Arc::new(Int32Array::from(vec![300, 400]))],
7532 )
7533 .unwrap();
7534 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
7535 dataset.append(batches, None).await.unwrap();
7536
7537 let mut list_req = ListTableVersionsRequest::new();
7539 list_req.id = Some(table_id.clone());
7540 let list_resp = namespace.list_table_versions(list_req).await.unwrap();
7541
7542 assert_eq!(
7543 list_resp.versions.len(),
7544 3,
7545 "Should have 3 versions, got: {:?}",
7546 list_resp.versions
7547 );
7548
7549 for expected_version in 1..=3 {
7551 let version = list_resp
7552 .versions
7553 .iter()
7554 .find(|v| v.version == expected_version)
7555 .unwrap_or_else(|| panic!("Expected version {}", expected_version));
7556
7557 assert!(
7558 !version.manifest_path.is_empty(),
7559 "manifest_path should be set for version {}",
7560 expected_version
7561 );
7562 assert!(
7563 version.manifest_path.contains(".manifest"),
7564 "manifest_path should contain .manifest for version {}",
7565 expected_version
7566 );
7567 assert!(
7568 version.manifest_size.is_some(),
7569 "manifest_size should be set for version {}",
7570 expected_version
7571 );
7572 assert!(
7573 version.manifest_size.unwrap() > 0,
7574 "manifest_size should be > 0 for version {}",
7575 expected_version
7576 );
7577 assert!(
7578 version.timestamp_millis.is_some(),
7579 "timestamp_millis should be set for version {}",
7580 expected_version
7581 );
7582 }
7583 }
7584
7585 #[tokio::test]
7586 async fn test_describe_table_version() {
7587 use arrow::array::{Int32Array, RecordBatchIterator};
7588 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7589 use arrow::record_batch::RecordBatch;
7590 use lance::dataset::{Dataset, WriteMode, WriteParams};
7591 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
7592
7593 let temp_dir = TempStrDir::default();
7594 let temp_path: &str = &temp_dir;
7595
7596 let namespace: Arc<dyn LanceNamespace> = Arc::new(
7597 DirectoryNamespaceBuilder::new(temp_path)
7598 .table_version_tracking_enabled(true)
7599 .build()
7600 .await
7601 .unwrap(),
7602 );
7603
7604 let mut create_ns_req = CreateNamespaceRequest::new();
7606 create_ns_req.id = Some(vec!["workspace".to_string()]);
7607 namespace.create_namespace(create_ns_req).await.unwrap();
7608
7609 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
7611 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
7612 "id",
7613 DataType::Int32,
7614 false,
7615 )]));
7616 let batch = RecordBatch::try_new(
7617 arrow_schema.clone(),
7618 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
7619 )
7620 .unwrap();
7621 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
7622 let write_params = WriteParams {
7623 mode: WriteMode::Create,
7624 ..Default::default()
7625 };
7626 let mut dataset = Dataset::write_into_namespace(
7627 batches,
7628 namespace.clone(),
7629 table_id.clone(),
7630 Some(write_params),
7631 )
7632 .await
7633 .unwrap();
7634
7635 let batch2 = RecordBatch::try_new(
7637 arrow_schema.clone(),
7638 vec![Arc::new(Int32Array::from(vec![100, 200]))],
7639 )
7640 .unwrap();
7641 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
7642 dataset.append(batches, None).await.unwrap();
7643
7644 let mut describe_req = DescribeTableVersionRequest::new();
7646 describe_req.id = Some(table_id.clone());
7647 describe_req.version = Some(1);
7648 let describe_resp = namespace
7649 .describe_table_version(describe_req)
7650 .await
7651 .unwrap();
7652
7653 let version = &describe_resp.version;
7654 assert_eq!(version.version, 1);
7655 assert!(version.timestamp_millis.is_some());
7656 assert!(
7657 !version.manifest_path.is_empty(),
7658 "manifest_path should be set"
7659 );
7660 assert!(
7661 version.manifest_path.contains(".manifest"),
7662 "manifest_path should contain .manifest"
7663 );
7664 assert!(
7665 version.manifest_size.is_some(),
7666 "manifest_size should be set"
7667 );
7668 assert!(
7669 version.manifest_size.unwrap() > 0,
7670 "manifest_size should be > 0"
7671 );
7672
7673 let mut describe_req = DescribeTableVersionRequest::new();
7675 describe_req.id = Some(table_id.clone());
7676 describe_req.version = Some(2);
7677 let describe_resp = namespace
7678 .describe_table_version(describe_req)
7679 .await
7680 .unwrap();
7681
7682 let version = &describe_resp.version;
7683 assert_eq!(version.version, 2);
7684 assert!(version.timestamp_millis.is_some());
7685 assert!(
7686 !version.manifest_path.is_empty(),
7687 "manifest_path should be set"
7688 );
7689 assert!(
7690 version.manifest_size.is_some(),
7691 "manifest_size should be set"
7692 );
7693 assert!(
7694 version.manifest_size.unwrap() > 0,
7695 "manifest_size should be > 0"
7696 );
7697 }
7698
7699 #[tokio::test]
7700 async fn test_describe_table_version_latest() {
7701 use arrow::array::{Int32Array, RecordBatchIterator};
7702 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7703 use arrow::record_batch::RecordBatch;
7704 use lance::dataset::{Dataset, WriteMode, WriteParams};
7705 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
7706
7707 let temp_dir = TempStrDir::default();
7708 let temp_path: &str = &temp_dir;
7709
7710 let namespace: Arc<dyn LanceNamespace> = Arc::new(
7711 DirectoryNamespaceBuilder::new(temp_path)
7712 .table_version_tracking_enabled(true)
7713 .build()
7714 .await
7715 .unwrap(),
7716 );
7717
7718 let mut create_ns_req = CreateNamespaceRequest::new();
7720 create_ns_req.id = Some(vec!["workspace".to_string()]);
7721 namespace.create_namespace(create_ns_req).await.unwrap();
7722
7723 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
7725 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
7726 "id",
7727 DataType::Int32,
7728 false,
7729 )]));
7730 let batch = RecordBatch::try_new(
7731 arrow_schema.clone(),
7732 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
7733 )
7734 .unwrap();
7735 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
7736 let write_params = WriteParams {
7737 mode: WriteMode::Create,
7738 ..Default::default()
7739 };
7740 let mut dataset = Dataset::write_into_namespace(
7741 batches,
7742 namespace.clone(),
7743 table_id.clone(),
7744 Some(write_params),
7745 )
7746 .await
7747 .unwrap();
7748
7749 let batch2 = RecordBatch::try_new(
7751 arrow_schema.clone(),
7752 vec![Arc::new(Int32Array::from(vec![100, 200]))],
7753 )
7754 .unwrap();
7755 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
7756 dataset.append(batches, None).await.unwrap();
7757
7758 let batch3 = RecordBatch::try_new(
7760 arrow_schema.clone(),
7761 vec![Arc::new(Int32Array::from(vec![300, 400]))],
7762 )
7763 .unwrap();
7764 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
7765 dataset.append(batches, None).await.unwrap();
7766
7767 let mut describe_req = DescribeTableVersionRequest::new();
7769 describe_req.id = Some(table_id.clone());
7770 describe_req.version = None;
7771 let describe_resp = namespace
7772 .describe_table_version(describe_req)
7773 .await
7774 .unwrap();
7775
7776 assert_eq!(describe_resp.version.version, 3);
7778 }
7779
7780 #[tokio::test]
7781 async fn test_create_table_version() {
7782 use futures::TryStreamExt;
7783 use lance::dataset::builder::DatasetBuilder;
7784 use lance_namespace::models::CreateTableVersionRequest;
7785
7786 let temp_dir = TempStrDir::default();
7787 let temp_path: &str = &temp_dir;
7788
7789 let namespace: Arc<dyn LanceNamespace> = Arc::new(
7790 DirectoryNamespaceBuilder::new(temp_path)
7791 .table_version_tracking_enabled(true)
7792 .build()
7793 .await
7794 .unwrap(),
7795 );
7796
7797 let schema = create_test_schema();
7799 let ipc_data = create_test_ipc_data(&schema);
7800 let mut create_req = CreateTableRequest::new();
7801 create_req.id = Some(vec!["test_table".to_string()]);
7802 namespace
7803 .create_table(create_req, bytes::Bytes::from(ipc_data))
7804 .await
7805 .unwrap();
7806
7807 let table_id = vec!["test_table".to_string()];
7809 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
7810 .await
7811 .unwrap()
7812 .load()
7813 .await
7814 .unwrap();
7815
7816 let versions_path = dataset.versions_dir();
7818 let manifest_metas: Vec<_> = dataset
7819 .object_store(None)
7820 .await
7821 .unwrap()
7822 .inner
7823 .list(Some(&versions_path))
7824 .try_collect()
7825 .await
7826 .unwrap();
7827
7828 let manifest_meta = manifest_metas
7829 .iter()
7830 .find(|m| {
7831 m.location
7832 .filename()
7833 .map(|f| f.ends_with(".manifest"))
7834 .unwrap_or(false)
7835 })
7836 .expect("No manifest file found");
7837
7838 let manifest_data = dataset
7840 .object_store(None)
7841 .await
7842 .unwrap()
7843 .inner
7844 .get(&manifest_meta.location)
7845 .await
7846 .unwrap()
7847 .bytes()
7848 .await
7849 .unwrap();
7850
7851 let staging_path = dataset.versions_dir().join("staging_manifest");
7853 dataset
7854 .object_store(None)
7855 .await
7856 .unwrap()
7857 .inner
7858 .put(&staging_path, manifest_data.into())
7859 .await
7860 .unwrap();
7861
7862 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
7865 create_version_req.id = Some(table_id.clone());
7866 create_version_req.naming_scheme = Some("V2".to_string());
7867
7868 let result = namespace.create_table_version(create_version_req).await;
7869 assert!(
7870 result.is_ok(),
7871 "create_table_version should succeed: {:?}",
7872 result
7873 );
7874
7875 let response = result.unwrap();
7877 let version_info = response
7878 .version
7879 .expect("response should contain version info");
7880 let version_2_path = Path::parse(&version_info.manifest_path).unwrap();
7881 let head_result = dataset
7882 .object_store(None)
7883 .await
7884 .unwrap()
7885 .inner
7886 .head(&version_2_path)
7887 .await;
7888 assert!(
7889 head_result.is_ok(),
7890 "Version 2 manifest should exist at {}",
7891 version_2_path
7892 );
7893
7894 let staging_head_result = dataset
7896 .object_store(None)
7897 .await
7898 .unwrap()
7899 .inner
7900 .head(&staging_path)
7901 .await;
7902 assert!(
7903 staging_head_result.is_err(),
7904 "Staging manifest should have been deleted after create_table_version"
7905 );
7906 }
7907
7908 #[tokio::test]
7909 async fn test_create_table_version_conflict() {
7910 use futures::TryStreamExt;
7913 use lance::dataset::builder::DatasetBuilder;
7914 use lance_namespace::models::CreateTableVersionRequest;
7915
7916 let temp_dir = TempStrDir::default();
7917 let temp_path: &str = &temp_dir;
7918
7919 let namespace: Arc<dyn LanceNamespace> = Arc::new(
7920 DirectoryNamespaceBuilder::new(temp_path)
7921 .table_version_tracking_enabled(true)
7922 .build()
7923 .await
7924 .unwrap(),
7925 );
7926
7927 let schema = create_test_schema();
7929 let ipc_data = create_test_ipc_data(&schema);
7930 let mut create_req = CreateTableRequest::new();
7931 create_req.id = Some(vec!["test_table".to_string()]);
7932 namespace
7933 .create_table(create_req, bytes::Bytes::from(ipc_data))
7934 .await
7935 .unwrap();
7936
7937 let table_id = vec!["test_table".to_string()];
7939 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
7940 .await
7941 .unwrap()
7942 .load()
7943 .await
7944 .unwrap();
7945
7946 let versions_path = dataset.versions_dir();
7948 let manifest_metas: Vec<_> = dataset
7949 .object_store(None)
7950 .await
7951 .unwrap()
7952 .inner
7953 .list(Some(&versions_path))
7954 .try_collect()
7955 .await
7956 .unwrap();
7957
7958 let manifest_meta = manifest_metas
7959 .iter()
7960 .find(|m| {
7961 m.location
7962 .filename()
7963 .map(|f| f.ends_with(".manifest"))
7964 .unwrap_or(false)
7965 })
7966 .expect("No manifest file found");
7967
7968 let manifest_data = dataset
7970 .object_store(None)
7971 .await
7972 .unwrap()
7973 .inner
7974 .get(&manifest_meta.location)
7975 .await
7976 .unwrap()
7977 .bytes()
7978 .await
7979 .unwrap();
7980
7981 let staging_path = dataset.versions_dir().join("staging_manifest");
7983 dataset
7984 .object_store(None)
7985 .await
7986 .unwrap()
7987 .inner
7988 .put(&staging_path, manifest_data.into())
7989 .await
7990 .unwrap();
7991
7992 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
7994 create_version_req.id = Some(table_id.clone());
7995 create_version_req.naming_scheme = Some("V2".to_string());
7996 let first_result = namespace.create_table_version(create_version_req).await;
7997 assert!(
7998 first_result.is_ok(),
7999 "First create_table_version for version 2 should succeed: {:?}",
8000 first_result
8001 );
8002
8003 let version_2_path = Path::parse(
8005 &first_result
8006 .unwrap()
8007 .version
8008 .expect("response should contain version info")
8009 .manifest_path,
8010 )
8011 .unwrap();
8012
8013 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
8015 create_version_req.id = Some(table_id.clone());
8016 create_version_req.naming_scheme = Some("V2".to_string());
8017
8018 let result = namespace.create_table_version(create_version_req).await;
8019 assert!(
8020 result.is_err(),
8021 "create_table_version should fail for existing version"
8022 );
8023
8024 let head_result = dataset
8026 .object_store(None)
8027 .await
8028 .unwrap()
8029 .inner
8030 .head(&version_2_path)
8031 .await;
8032 assert!(
8033 head_result.is_ok(),
8034 "Version 2 manifest should still exist at {}",
8035 version_2_path
8036 );
8037 }
8038
8039 #[tokio::test]
8040 async fn test_create_table_version_table_not_found() {
8041 use lance_namespace::models::CreateTableVersionRequest;
8042
8043 let temp_dir = TempStdDir::default();
8044 let temp_path = temp_dir.to_str().unwrap();
8045
8046 let namespace = DirectoryNamespaceBuilder::new(temp_path)
8047 .table_version_tracking_enabled(true)
8048 .build()
8049 .await
8050 .unwrap();
8051
8052 let mut create_version_req =
8054 CreateTableVersionRequest::new(1, "/some/staging/path".to_string());
8055 create_version_req.id = Some(vec!["non_existent_table".to_string()]);
8056
8057 let result = namespace.create_table_version(create_version_req).await;
8058 assert!(
8059 result.is_err(),
8060 "create_table_version should fail for non-existent table"
8061 );
8062 let err_msg = result.unwrap_err().to_string();
8063 assert!(
8064 err_msg.contains("Table not found"),
8065 "Error should mention table not found, got: {}",
8066 err_msg
8067 );
8068 }
8069
8070 mod e2e_table_version_tracking {
8072 use super::*;
8073 use std::sync::atomic::{AtomicUsize, Ordering};
8074
8075 struct TrackingNamespace {
8077 inner: DirectoryNamespace,
8078 create_table_version_count: AtomicUsize,
8079 describe_table_version_count: AtomicUsize,
8080 list_table_versions_count: AtomicUsize,
8081 }
8082
8083 impl TrackingNamespace {
8084 fn new(inner: DirectoryNamespace) -> Self {
8085 Self {
8086 inner,
8087 create_table_version_count: AtomicUsize::new(0),
8088 describe_table_version_count: AtomicUsize::new(0),
8089 list_table_versions_count: AtomicUsize::new(0),
8090 }
8091 }
8092
8093 fn create_table_version_calls(&self) -> usize {
8094 self.create_table_version_count.load(Ordering::SeqCst)
8095 }
8096
8097 fn describe_table_version_calls(&self) -> usize {
8098 self.describe_table_version_count.load(Ordering::SeqCst)
8099 }
8100
8101 fn list_table_versions_calls(&self) -> usize {
8102 self.list_table_versions_count.load(Ordering::SeqCst)
8103 }
8104 }
8105
8106 impl std::fmt::Debug for TrackingNamespace {
8107 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
8108 f.debug_struct("TrackingNamespace")
8109 .field(
8110 "create_table_version_calls",
8111 &self.create_table_version_calls(),
8112 )
8113 .finish()
8114 }
8115 }
8116
8117 #[async_trait]
8118 impl LanceNamespace for TrackingNamespace {
8119 async fn create_namespace(
8120 &self,
8121 request: CreateNamespaceRequest,
8122 ) -> Result<CreateNamespaceResponse> {
8123 self.inner.create_namespace(request).await
8124 }
8125
8126 async fn describe_namespace(
8127 &self,
8128 request: DescribeNamespaceRequest,
8129 ) -> Result<DescribeNamespaceResponse> {
8130 self.inner.describe_namespace(request).await
8131 }
8132
8133 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
8134 self.inner.namespace_exists(request).await
8135 }
8136
8137 async fn list_namespaces(
8138 &self,
8139 request: ListNamespacesRequest,
8140 ) -> Result<ListNamespacesResponse> {
8141 self.inner.list_namespaces(request).await
8142 }
8143
8144 async fn drop_namespace(
8145 &self,
8146 request: DropNamespaceRequest,
8147 ) -> Result<DropNamespaceResponse> {
8148 self.inner.drop_namespace(request).await
8149 }
8150
8151 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
8152 self.inner.list_tables(request).await
8153 }
8154
8155 async fn describe_table(
8156 &self,
8157 request: DescribeTableRequest,
8158 ) -> Result<DescribeTableResponse> {
8159 self.inner.describe_table(request).await
8160 }
8161
8162 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
8163 self.inner.table_exists(request).await
8164 }
8165
8166 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
8167 self.inner.drop_table(request).await
8168 }
8169
8170 async fn create_table(
8171 &self,
8172 request: CreateTableRequest,
8173 request_data: Bytes,
8174 ) -> Result<CreateTableResponse> {
8175 self.inner.create_table(request, request_data).await
8176 }
8177
8178 async fn declare_table(
8179 &self,
8180 request: DeclareTableRequest,
8181 ) -> Result<DeclareTableResponse> {
8182 self.inner.declare_table(request).await
8183 }
8184
8185 async fn list_table_versions(
8186 &self,
8187 request: ListTableVersionsRequest,
8188 ) -> Result<ListTableVersionsResponse> {
8189 self.list_table_versions_count
8190 .fetch_add(1, Ordering::SeqCst);
8191 self.inner.list_table_versions(request).await
8192 }
8193
8194 async fn create_table_version(
8195 &self,
8196 request: CreateTableVersionRequest,
8197 ) -> Result<CreateTableVersionResponse> {
8198 self.create_table_version_count
8199 .fetch_add(1, Ordering::SeqCst);
8200 self.inner.create_table_version(request).await
8201 }
8202
8203 async fn describe_table_version(
8204 &self,
8205 request: DescribeTableVersionRequest,
8206 ) -> Result<DescribeTableVersionResponse> {
8207 self.describe_table_version_count
8208 .fetch_add(1, Ordering::SeqCst);
8209 self.inner.describe_table_version(request).await
8210 }
8211
8212 async fn batch_delete_table_versions(
8213 &self,
8214 request: BatchDeleteTableVersionsRequest,
8215 ) -> Result<BatchDeleteTableVersionsResponse> {
8216 self.inner.batch_delete_table_versions(request).await
8217 }
8218
8219 fn namespace_id(&self) -> String {
8220 self.inner.namespace_id()
8221 }
8222 }
8223
8224 #[tokio::test]
8225 async fn test_describe_table_returns_managed_versioning() {
8226 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableRequest};
8227
8228 let temp_dir = TempStdDir::default();
8229 let temp_path = temp_dir.to_str().unwrap();
8230
8231 let ns = DirectoryNamespaceBuilder::new(temp_path)
8233 .table_version_tracking_enabled(true)
8234 .manifest_enabled(true)
8235 .build()
8236 .await
8237 .unwrap();
8238
8239 let mut create_ns_req = CreateNamespaceRequest::new();
8241 create_ns_req.id = Some(vec!["workspace".to_string()]);
8242 ns.create_namespace(create_ns_req).await.unwrap();
8243
8244 let schema = create_test_schema();
8246 let ipc_data = create_test_ipc_data(&schema);
8247 let mut create_req = CreateTableRequest::new();
8248 create_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
8249 ns.create_table(create_req, bytes::Bytes::from(ipc_data))
8250 .await
8251 .unwrap();
8252
8253 let mut describe_req = DescribeTableRequest::new();
8255 describe_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
8256 let describe_resp = ns.describe_table(describe_req).await.unwrap();
8257
8258 assert_eq!(
8260 describe_resp.managed_versioning,
8261 Some(true),
8262 "managed_versioning should be true when table_version_tracking_enabled=true"
8263 );
8264 }
8265
8266 #[tokio::test]
8267 async fn test_external_manifest_store_invokes_namespace_apis() {
8268 use arrow::array::{Int32Array, StringArray};
8269 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
8270 use arrow::record_batch::RecordBatch;
8271 use lance::Dataset;
8272 use lance::dataset::builder::DatasetBuilder;
8273 use lance::dataset::{WriteMode, WriteParams};
8274 use lance_namespace::models::CreateNamespaceRequest;
8275
8276 let temp_dir = TempStdDir::default();
8277 let temp_path = temp_dir.to_str().unwrap();
8278
8279 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
8281 .table_version_tracking_enabled(true)
8282 .manifest_enabled(true)
8283 .build()
8284 .await
8285 .unwrap();
8286
8287 let tracking_ns = Arc::new(TrackingNamespace::new(inner_ns));
8288 let ns: Arc<dyn LanceNamespace> = tracking_ns.clone();
8289
8290 let mut create_ns_req = CreateNamespaceRequest::new();
8292 create_ns_req.id = Some(vec!["workspace".to_string()]);
8293 ns.create_namespace(create_ns_req).await.unwrap();
8294
8295 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
8297
8298 let arrow_schema = Arc::new(ArrowSchema::new(vec![
8300 Field::new("id", DataType::Int32, false),
8301 Field::new("name", DataType::Utf8, true),
8302 ]));
8303 let batch = RecordBatch::try_new(
8304 arrow_schema.clone(),
8305 vec![
8306 Arc::new(Int32Array::from(vec![1, 2, 3])),
8307 Arc::new(StringArray::from(vec!["a", "b", "c"])),
8308 ],
8309 )
8310 .unwrap();
8311
8312 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
8314 let write_params = WriteParams {
8315 mode: WriteMode::Create,
8316 ..Default::default()
8317 };
8318 let mut dataset = Dataset::write_into_namespace(
8319 batches,
8320 ns.clone(),
8321 table_id.clone(),
8322 Some(write_params),
8323 )
8324 .await
8325 .unwrap();
8326 assert_eq!(dataset.version().version, 1);
8327
8328 assert_eq!(
8330 tracking_ns.create_table_version_calls(),
8331 1,
8332 "create_table_version should have been called once during initial write_into_namespace"
8333 );
8334
8335 let append_batch = RecordBatch::try_new(
8337 arrow_schema.clone(),
8338 vec![
8339 Arc::new(Int32Array::from(vec![4, 5, 6])),
8340 Arc::new(StringArray::from(vec!["d", "e", "f"])),
8341 ],
8342 )
8343 .unwrap();
8344 let append_batches = RecordBatchIterator::new(vec![Ok(append_batch)], arrow_schema);
8345 dataset.append(append_batches, None).await.unwrap();
8346
8347 assert_eq!(
8348 tracking_ns.create_table_version_calls(),
8349 2,
8350 "create_table_version should have been called twice (once for create, once for append)"
8351 );
8352
8353 let initial_list_calls = tracking_ns.list_table_versions_calls();
8355 let latest_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8356 .await
8357 .unwrap()
8358 .load()
8359 .await
8360 .unwrap();
8361 assert_eq!(latest_dataset.version().version, 2);
8362 assert_eq!(
8363 tracking_ns.list_table_versions_calls(),
8364 initial_list_calls + 1,
8365 "list_table_versions should have been called exactly once during checkout_latest"
8366 );
8367
8368 let initial_describe_calls = tracking_ns.describe_table_version_calls();
8370 let v1_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8371 .await
8372 .unwrap()
8373 .with_version(1)
8374 .load()
8375 .await
8376 .unwrap();
8377 assert_eq!(v1_dataset.version().version, 1);
8378 assert_eq!(
8379 tracking_ns.describe_table_version_calls(),
8380 initial_describe_calls + 1,
8381 "describe_table_version should have been called exactly once during checkout to version 1"
8382 );
8383 }
8384
8385 #[tokio::test]
8386 async fn test_dataset_commit_with_external_manifest_store() {
8387 use arrow::array::{Int32Array, StringArray};
8388 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
8389 use arrow::record_batch::RecordBatch;
8390 use futures::TryStreamExt;
8391 use lance::dataset::{Dataset, WriteMode, WriteParams};
8392 use lance_namespace::models::CreateNamespaceRequest;
8393 use lance_table::io::commit::ManifestNamingScheme;
8394
8395 let temp_dir = TempStdDir::default();
8396 let temp_path = temp_dir.to_str().unwrap();
8397
8398 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
8400 .table_version_tracking_enabled(true)
8401 .manifest_enabled(true)
8402 .build()
8403 .await
8404 .unwrap();
8405
8406 let tracking_ns: Arc<dyn LanceNamespace> = Arc::new(TrackingNamespace::new(inner_ns));
8407
8408 let mut create_ns_req = CreateNamespaceRequest::new();
8410 create_ns_req.id = Some(vec!["workspace".to_string()]);
8411 tracking_ns.create_namespace(create_ns_req).await.unwrap();
8412
8413 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
8415 let arrow_schema = Arc::new(ArrowSchema::new(vec![
8416 Field::new("id", DataType::Int32, false),
8417 Field::new("name", DataType::Utf8, true),
8418 ]));
8419 let batch = RecordBatch::try_new(
8420 arrow_schema.clone(),
8421 vec![
8422 Arc::new(Int32Array::from(vec![1, 2, 3])),
8423 Arc::new(StringArray::from(vec!["a", "b", "c"])),
8424 ],
8425 )
8426 .unwrap();
8427 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
8428 let write_params = WriteParams {
8429 mode: WriteMode::Create,
8430 ..Default::default()
8431 };
8432 let dataset = Dataset::write_into_namespace(
8433 batches,
8434 tracking_ns.clone(),
8435 table_id.clone(),
8436 Some(write_params),
8437 )
8438 .await
8439 .unwrap();
8440 assert_eq!(dataset.version().version, 1);
8441
8442 let batch2 = RecordBatch::try_new(
8444 arrow_schema.clone(),
8445 vec![
8446 Arc::new(Int32Array::from(vec![4, 5, 6])),
8447 Arc::new(StringArray::from(vec!["d", "e", "f"])),
8448 ],
8449 )
8450 .unwrap();
8451 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
8452 let write_params = WriteParams {
8453 mode: WriteMode::Append,
8454 ..Default::default()
8455 };
8456 Dataset::write_into_namespace(
8457 batches,
8458 tracking_ns.clone(),
8459 table_id.clone(),
8460 Some(write_params),
8461 )
8462 .await
8463 .unwrap();
8464
8465 let manifest_metas: Vec<_> = dataset
8468 .object_store(None)
8469 .await
8470 .unwrap()
8471 .inner
8472 .list(Some(&dataset.versions_dir()))
8473 .try_collect()
8474 .await
8475 .unwrap();
8476 let version_2_found = manifest_metas.iter().any(|m| {
8477 m.location
8478 .filename()
8479 .map(|f| {
8480 f.ends_with(".manifest")
8481 && ManifestNamingScheme::V2.parse_version(f) == Some(2)
8482 })
8483 .unwrap_or(false)
8484 });
8485 assert!(
8486 version_2_found,
8487 "Version 2 manifest should exist in versions directory"
8488 );
8489 }
8490
8491 async fn create_ns_with_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
8493 use arrow::array::{Int32Array, StringArray};
8494 use arrow::ipc::writer::StreamWriter;
8495
8496 let (namespace, temp_dir) = create_test_namespace().await;
8497
8498 let schema = create_test_schema();
8499 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8500 let arrow_schema = Arc::new(arrow_schema);
8501
8502 let id_array = Int32Array::from(vec![1, 2, 3]);
8503 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
8504 let batch = arrow::record_batch::RecordBatch::try_new(
8505 arrow_schema.clone(),
8506 vec![Arc::new(id_array), Arc::new(name_array)],
8507 )
8508 .unwrap();
8509
8510 let mut buffer = Vec::new();
8511 {
8512 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8513 writer.write(&batch).unwrap();
8514 writer.finish().unwrap();
8515 }
8516
8517 let mut request = CreateTableRequest::new();
8518 let table_id = vec!["test_ops_table".to_string()];
8519 request.id = Some(table_id.clone());
8520
8521 namespace
8522 .create_table(request, Bytes::from(buffer))
8523 .await
8524 .unwrap();
8525
8526 (namespace, temp_dir, table_id)
8527 }
8528
8529 #[tokio::test]
8530 async fn test_count_table_rows_basic() {
8531 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8532
8533 let request = CountTableRowsRequest {
8534 id: Some(table_id),
8535 version: None,
8536 predicate: None,
8537 ..Default::default()
8538 };
8539
8540 let count = namespace.count_table_rows(request).await.unwrap();
8541 assert_eq!(count, 3);
8542 }
8543
8544 #[tokio::test]
8545 async fn test_count_table_rows_with_predicate() {
8546 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8547
8548 let request = CountTableRowsRequest {
8549 id: Some(table_id),
8550 version: None,
8551 predicate: Some("id > 1".to_string()),
8552 ..Default::default()
8553 };
8554
8555 let count = namespace.count_table_rows(request).await.unwrap();
8556 assert_eq!(count, 2);
8557 }
8558
8559 #[tokio::test]
8560 async fn test_query_table_invalid_distance_type() {
8561 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
8562
8563 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
8564 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
8565 multi_vector: None,
8566 });
8567
8568 let request = QueryTableRequest {
8569 id: Some(table_id),
8570 k: 2,
8571 vector,
8572 vector_column: Some("vector".to_string()),
8573 distance_type: Some("invalid_metric".to_string()),
8574 filter: None,
8575 offset: None,
8576 version: None,
8577 ..Default::default()
8578 };
8579
8580 let result = namespace.query_table(request).await;
8581 assert!(result.is_err());
8582 let err_msg = result.unwrap_err().to_string();
8583 assert!(
8584 err_msg.contains("Unknown distance type"),
8585 "Expected error about unknown distance type, got: {}",
8586 err_msg
8587 );
8588 }
8589
8590 #[tokio::test]
8591 async fn test_insert_into_table_append() {
8592 use arrow::array::{Int32Array, StringArray};
8593 use arrow::ipc::writer::StreamWriter;
8594
8595 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8596
8597 let schema = create_test_schema();
8599 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8600 let arrow_schema = Arc::new(arrow_schema);
8601
8602 let id_array = Int32Array::from(vec![4, 5]);
8603 let name_array = StringArray::from(vec!["Dave", "Eve"]);
8604 let batch = arrow::record_batch::RecordBatch::try_new(
8605 arrow_schema.clone(),
8606 vec![Arc::new(id_array), Arc::new(name_array)],
8607 )
8608 .unwrap();
8609
8610 let mut buffer = Vec::new();
8611 {
8612 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8613 writer.write(&batch).unwrap();
8614 writer.finish().unwrap();
8615 }
8616
8617 let request = InsertIntoTableRequest {
8618 id: Some(table_id.clone()),
8619 mode: Some("append".to_string()),
8620 ..Default::default()
8621 };
8622
8623 let response = namespace
8624 .insert_into_table(request, Bytes::from(buffer))
8625 .await
8626 .unwrap();
8627 assert!(response.transaction_id.is_none());
8628
8629 let count_req = CountTableRowsRequest {
8631 id: Some(table_id),
8632 version: None,
8633 predicate: None,
8634 ..Default::default()
8635 };
8636 let count = namespace.count_table_rows(count_req).await.unwrap();
8637 assert_eq!(count, 5);
8638 }
8639
8640 #[tokio::test]
8641 async fn test_insert_into_table_overwrite() {
8642 use arrow::array::{Int32Array, StringArray};
8643 use arrow::ipc::writer::StreamWriter;
8644
8645 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8646
8647 let schema = create_test_schema();
8648 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8649 let arrow_schema = Arc::new(arrow_schema);
8650
8651 let id_array = Int32Array::from(vec![10, 20]);
8652 let name_array = StringArray::from(vec!["X", "Y"]);
8653 let batch = arrow::record_batch::RecordBatch::try_new(
8654 arrow_schema.clone(),
8655 vec![Arc::new(id_array), Arc::new(name_array)],
8656 )
8657 .unwrap();
8658
8659 let mut buffer = Vec::new();
8660 {
8661 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8662 writer.write(&batch).unwrap();
8663 writer.finish().unwrap();
8664 }
8665
8666 let request = InsertIntoTableRequest {
8667 id: Some(table_id.clone()),
8668 mode: Some("overwrite".to_string()),
8669 ..Default::default()
8670 };
8671
8672 namespace
8673 .insert_into_table(request, Bytes::from(buffer))
8674 .await
8675 .unwrap();
8676
8677 let count_req = CountTableRowsRequest {
8679 id: Some(table_id),
8680 version: None,
8681 predicate: None,
8682 ..Default::default()
8683 };
8684 let count = namespace.count_table_rows(count_req).await.unwrap();
8685 assert_eq!(count, 2);
8686 }
8687
8688 #[tokio::test]
8689 async fn test_insert_into_table_empty_data() {
8690 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8691
8692 let request = InsertIntoTableRequest {
8693 id: Some(table_id),
8694 mode: None,
8695 ..Default::default()
8696 };
8697
8698 let result = namespace.insert_into_table(request, Bytes::new()).await;
8699 assert!(result.is_err());
8700 assert!(
8701 result
8702 .unwrap_err()
8703 .to_string()
8704 .contains("Arrow IPC stream) is required")
8705 );
8706 }
8707
8708 #[tokio::test]
8709 async fn test_insert_into_table_with_storage_options() {
8710 use arrow::array::{Int32Array, StringArray};
8711 use arrow::ipc::writer::StreamWriter;
8712
8713 let temp_dir = TempStdDir::default();
8714
8715 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
8717 .storage_option("allow_http", "true")
8718 .build()
8719 .await
8720 .unwrap();
8721
8722 let schema = create_test_schema();
8724 let ipc_data = create_test_ipc_data(&schema);
8725 let mut create_req = CreateTableRequest::new();
8726 let table_id = vec!["so_table".to_string()];
8727 create_req.id = Some(table_id.clone());
8728 namespace
8729 .create_table(create_req, Bytes::from(ipc_data))
8730 .await
8731 .unwrap();
8732
8733 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8735 let arrow_schema = Arc::new(arrow_schema);
8736
8737 let id_array = Int32Array::from(vec![10, 20]);
8738 let name_array = StringArray::from(vec!["X", "Y"]);
8739 let batch = arrow::record_batch::RecordBatch::try_new(
8740 arrow_schema.clone(),
8741 vec![Arc::new(id_array), Arc::new(name_array)],
8742 )
8743 .unwrap();
8744
8745 let mut buffer = Vec::new();
8746 {
8747 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8748 writer.write(&batch).unwrap();
8749 writer.finish().unwrap();
8750 }
8751
8752 let request = InsertIntoTableRequest {
8753 id: Some(table_id.clone()),
8754 mode: Some("append".to_string()),
8755 ..Default::default()
8756 };
8757
8758 let response = namespace
8759 .insert_into_table(request, Bytes::from(buffer))
8760 .await
8761 .unwrap();
8762 assert!(response.transaction_id.is_none());
8763
8764 let count_req = CountTableRowsRequest {
8766 id: Some(table_id),
8767 version: None,
8768 predicate: None,
8769 ..Default::default()
8770 };
8771 let count = namespace.count_table_rows(count_req).await.unwrap();
8772 assert_eq!(count, 2);
8773 }
8774
8775 #[tokio::test]
8776 async fn test_query_table_basic() {
8777 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8778
8779 let request = QueryTableRequest {
8780 id: Some(table_id),
8781 k: 10,
8782 filter: None,
8783 offset: None,
8784 version: None,
8785 ..Default::default()
8786 };
8787
8788 let bytes = namespace.query_table(request).await.unwrap();
8789
8790 let cursor = Cursor::new(bytes.to_vec());
8792 let reader = FileReader::try_new(cursor, None).unwrap();
8793 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
8794 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
8795 assert_eq!(total_rows, 3);
8796 }
8797
8798 #[tokio::test]
8799 async fn test_query_table_with_filter() {
8800 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8801
8802 let request = QueryTableRequest {
8803 id: Some(table_id),
8804 k: 10,
8805 filter: Some("id <= 2".to_string()),
8806 offset: None,
8807 version: None,
8808 ..Default::default()
8809 };
8810
8811 let bytes = namespace.query_table(request).await.unwrap();
8812
8813 let cursor = Cursor::new(bytes.to_vec());
8814 let reader = FileReader::try_new(cursor, None).unwrap();
8815 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
8816 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
8817 assert_eq!(total_rows, 2);
8818 }
8819
8820 #[tokio::test]
8821 async fn test_query_table_with_limit_and_offset() {
8822 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8823
8824 let request = QueryTableRequest {
8825 id: Some(table_id),
8826 k: 2,
8827 filter: None,
8828 offset: Some(1),
8829 version: None,
8830 ..Default::default()
8831 };
8832
8833 let bytes = namespace.query_table(request).await.unwrap();
8834
8835 let cursor = Cursor::new(bytes.to_vec());
8836 let reader = FileReader::try_new(cursor, None).unwrap();
8837 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
8838 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
8839 assert_eq!(total_rows, 2);
8840 }
8841
8842 #[tokio::test]
8843 async fn test_query_table_no_limit() {
8844 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8845
8846 let request = QueryTableRequest {
8848 id: Some(table_id),
8849 k: 0,
8850 filter: None,
8851 offset: None,
8852 version: None,
8853 ..Default::default()
8854 };
8855
8856 let bytes = namespace.query_table(request).await.unwrap();
8857
8858 let cursor = Cursor::new(bytes.to_vec());
8859 let reader = FileReader::try_new(cursor, None).unwrap();
8860 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
8861 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
8862 assert_eq!(total_rows, 3);
8863 }
8864
8865 #[tokio::test]
8866 async fn test_query_table_with_columns() {
8867 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8868
8869 let columns = Box::new(lance_namespace::models::QueryTableRequestColumns {
8870 column_names: Some(vec!["id".to_string()]),
8871 column_aliases: None,
8872 });
8873
8874 let request = QueryTableRequest {
8875 id: Some(table_id),
8876 k: 10,
8877 filter: None,
8878 offset: None,
8879 version: None,
8880 columns: Some(columns),
8881 ..Default::default()
8882 };
8883
8884 let bytes = namespace.query_table(request).await.unwrap();
8885
8886 let cursor = Cursor::new(bytes.to_vec());
8887 let reader = FileReader::try_new(cursor, None).unwrap();
8888 let schema = reader.schema();
8889 assert_eq!(schema.fields().len(), 1);
8890 assert_eq!(schema.field(0).name(), "id");
8891 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
8892 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
8893 assert_eq!(total_rows, 3);
8894 }
8895
8896 #[tokio::test]
8897 async fn test_count_table_rows_with_version() {
8898 use arrow::array::{Int32Array, StringArray};
8899 use arrow::ipc::writer::StreamWriter;
8900
8901 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8902
8903 let schema = create_test_schema();
8905 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8906 let arrow_schema = Arc::new(arrow_schema);
8907
8908 let id_array = Int32Array::from(vec![4, 5]);
8909 let name_array = StringArray::from(vec!["Dave", "Eve"]);
8910 let batch = arrow::record_batch::RecordBatch::try_new(
8911 arrow_schema.clone(),
8912 vec![Arc::new(id_array), Arc::new(name_array)],
8913 )
8914 .unwrap();
8915
8916 let mut buffer = Vec::new();
8917 {
8918 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8919 writer.write(&batch).unwrap();
8920 writer.finish().unwrap();
8921 }
8922
8923 let request = InsertIntoTableRequest {
8924 id: Some(table_id.clone()),
8925 mode: None,
8926 ..Default::default()
8927 };
8928 namespace
8929 .insert_into_table(request, Bytes::from(buffer))
8930 .await
8931 .unwrap();
8932
8933 let count_req = CountTableRowsRequest {
8935 id: Some(table_id.clone()),
8936 version: Some(1),
8937 predicate: None,
8938 ..Default::default()
8939 };
8940 let count = namespace.count_table_rows(count_req).await.unwrap();
8941 assert_eq!(count, 3);
8942
8943 let count_req = CountTableRowsRequest {
8945 id: Some(table_id),
8946 version: None,
8947 predicate: None,
8948 ..Default::default()
8949 };
8950 let count = namespace.count_table_rows(count_req).await.unwrap();
8951 assert_eq!(count, 5);
8952 }
8953
8954 #[tokio::test]
8955 async fn test_query_table_with_version() {
8956 use arrow::array::{Int32Array, StringArray};
8957 use arrow::ipc::writer::StreamWriter;
8958
8959 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
8960
8961 let schema = create_test_schema();
8963 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
8964 let arrow_schema = Arc::new(arrow_schema);
8965
8966 let id_array = Int32Array::from(vec![4, 5]);
8967 let name_array = StringArray::from(vec!["Dave", "Eve"]);
8968 let batch = arrow::record_batch::RecordBatch::try_new(
8969 arrow_schema.clone(),
8970 vec![Arc::new(id_array), Arc::new(name_array)],
8971 )
8972 .unwrap();
8973
8974 let mut buffer = Vec::new();
8975 {
8976 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
8977 writer.write(&batch).unwrap();
8978 writer.finish().unwrap();
8979 }
8980
8981 let request = InsertIntoTableRequest {
8982 id: Some(table_id.clone()),
8983 mode: None,
8984 ..Default::default()
8985 };
8986 namespace
8987 .insert_into_table(request, Bytes::from(buffer))
8988 .await
8989 .unwrap();
8990
8991 let request = QueryTableRequest {
8993 id: Some(table_id.clone()),
8994 k: 100,
8995 filter: None,
8996 offset: None,
8997 version: Some(1),
8998 ..Default::default()
8999 };
9000
9001 let bytes = namespace.query_table(request).await.unwrap();
9002 let cursor = Cursor::new(bytes.to_vec());
9003 let reader = FileReader::try_new(cursor, None).unwrap();
9004 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9005 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9006 assert_eq!(total_rows, 3);
9007
9008 let request = QueryTableRequest {
9010 id: Some(table_id),
9011 k: 100,
9012 filter: None,
9013 offset: None,
9014 version: None,
9015 ..Default::default()
9016 };
9017
9018 let bytes = namespace.query_table(request).await.unwrap();
9019 let cursor = Cursor::new(bytes.to_vec());
9020 let reader = FileReader::try_new(cursor, None).unwrap();
9021 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9022 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9023 assert_eq!(total_rows, 5);
9024 }
9025
9026 async fn create_ns_with_vector_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
9029 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
9030 use arrow::ipc::writer::StreamWriter;
9031
9032 let (namespace, temp_dir) = create_test_namespace().await;
9033
9034 let arrow_schema = Arc::new(arrow::datatypes::Schema::new(vec![
9036 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
9037 arrow::datatypes::Field::new(
9038 "vector",
9039 arrow::datatypes::DataType::FixedSizeList(
9040 Arc::new(arrow::datatypes::Field::new(
9041 "item",
9042 arrow::datatypes::DataType::Float32,
9043 true,
9044 )),
9045 4,
9046 ),
9047 true,
9048 ),
9049 ]));
9050
9051 let id_array = Int32Array::from(vec![1, 2, 3]);
9052 let values = Float32Array::from(vec![
9053 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ]);
9057 let vector_array = FixedSizeListArray::try_new(
9058 Arc::new(arrow::datatypes::Field::new(
9059 "item",
9060 arrow::datatypes::DataType::Float32,
9061 true,
9062 )),
9063 4,
9064 Arc::new(values),
9065 None,
9066 )
9067 .unwrap();
9068
9069 let batch = arrow::record_batch::RecordBatch::try_new(
9070 arrow_schema.clone(),
9071 vec![Arc::new(id_array), Arc::new(vector_array)],
9072 )
9073 .unwrap();
9074
9075 let mut buffer = Vec::new();
9076 {
9077 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
9078 writer.write(&batch).unwrap();
9079 writer.finish().unwrap();
9080 }
9081
9082 let table_name = "vector_table";
9084 let table_uri = format!("{}/{}.lance", temp_dir.to_str().unwrap(), table_name);
9085 let reader = arrow::record_batch::RecordBatchIterator::new(
9086 vec![Ok(batch)],
9087 arrow_schema.clone(),
9088 );
9089 Dataset::write(reader, &table_uri, None).await.unwrap();
9090
9091 let table_id = vec![table_name.to_string()];
9092 (namespace, temp_dir, table_id)
9093 }
9094
9095 #[tokio::test]
9096 async fn test_query_table_vector_search() {
9097 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
9098
9099 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
9100 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
9101 multi_vector: None,
9102 });
9103
9104 let request = QueryTableRequest {
9105 id: Some(table_id),
9106 k: 2,
9107 vector,
9108 filter: None,
9109 offset: None,
9110 version: None,
9111 ..Default::default()
9112 };
9113
9114 let bytes = namespace.query_table(request).await.unwrap();
9115
9116 let cursor = Cursor::new(bytes.to_vec());
9117 let reader = FileReader::try_new(cursor, None).unwrap();
9118 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9119 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9120 assert_eq!(total_rows, 2);
9121 }
9122
9123 #[tokio::test]
9124 async fn test_query_table_vector_search_with_distance_type() {
9125 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
9126
9127 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
9128 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
9129 multi_vector: None,
9130 });
9131
9132 let request = QueryTableRequest {
9133 id: Some(table_id),
9134 k: 3,
9135 vector,
9136 filter: None,
9137 offset: None,
9138 version: None,
9139 distance_type: Some("cosine".to_string()),
9140 ..Default::default()
9141 };
9142
9143 let bytes = namespace.query_table(request).await.unwrap();
9144
9145 let cursor = Cursor::new(bytes.to_vec());
9146 let reader = FileReader::try_new(cursor, None).unwrap();
9147 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9148 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9149 assert_eq!(total_rows, 3);
9150 }
9151
9152 #[tokio::test]
9153 async fn test_query_table_vector_search_with_filter() {
9154 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
9155
9156 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
9157 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
9158 multi_vector: None,
9159 });
9160
9161 let request = QueryTableRequest {
9162 id: Some(table_id),
9163 k: 10,
9164 vector,
9165 filter: Some("id <= 2".to_string()),
9166 offset: None,
9167 version: None,
9168 ..Default::default()
9169 };
9170
9171 let bytes = namespace.query_table(request).await.unwrap();
9172
9173 let cursor = Cursor::new(bytes.to_vec());
9174 let reader = FileReader::try_new(cursor, None).unwrap();
9175 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9176 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9177 assert!(total_rows <= 2);
9178 }
9179
9180 #[tokio::test]
9181 async fn test_query_table_vector_search_with_nprobes_and_refine() {
9182 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
9183
9184 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
9185 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
9186 multi_vector: None,
9187 });
9188
9189 let request = QueryTableRequest {
9190 id: Some(table_id),
9191 k: 2,
9192 vector,
9193 filter: None,
9194 offset: None,
9195 version: None,
9196 nprobes: Some(1),
9197 refine_factor: Some(1),
9198 prefilter: Some(true),
9199 ..Default::default()
9200 };
9201
9202 let bytes = namespace.query_table(request).await.unwrap();
9203
9204 let cursor = Cursor::new(bytes.to_vec());
9205 let reader = FileReader::try_new(cursor, None).unwrap();
9206 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9207 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9208 assert_eq!(total_rows, 2);
9209 }
9210
9211 #[tokio::test]
9212 async fn test_namespace_id() {
9213 let (namespace, _temp_dir) = create_test_namespace().await;
9214 let id = namespace.namespace_id();
9215 assert!(id.contains("DirectoryNamespace"));
9216 assert!(id.contains("root"));
9217 }
9218
9219 #[tokio::test]
9220 async fn test_query_table_empty_table() {
9221 let (namespace, _temp_dir) = create_test_namespace().await;
9222
9223 let schema = create_test_schema();
9225 let ipc_data = create_test_ipc_data(&schema);
9226 let mut create_request = CreateTableRequest::new();
9227 create_request.id = Some(vec!["empty_table".to_string()]);
9228 namespace
9229 .create_table(create_request, bytes::Bytes::from(ipc_data))
9230 .await
9231 .unwrap();
9232
9233 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
9235 single_vector: None,
9236 multi_vector: None,
9237 });
9238 let request = QueryTableRequest {
9239 id: Some(vec!["empty_table".to_string()]),
9240 k: 10,
9241 vector,
9242 ..Default::default()
9243 };
9244 let bytes = namespace.query_table(request).await.unwrap();
9245
9246 let cursor = Cursor::new(bytes.to_vec());
9247 let reader = FileReader::try_new(cursor, None).unwrap();
9248 let batches: Vec<_> = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
9249 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9250 assert_eq!(total_rows, 0, "empty table should yield no rows");
9251 }
9252
9253 #[tokio::test]
9254 async fn test_query_table_with_plain_filter_no_vector() {
9255 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
9256
9257 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
9259 single_vector: None,
9260 multi_vector: None,
9261 });
9262 let request = QueryTableRequest {
9263 id: Some(table_id),
9264 k: 0,
9265 vector,
9266 filter: Some("id > 1".to_string()),
9267 ..Default::default()
9268 };
9269 let bytes = namespace.query_table(request).await.unwrap();
9270
9271 let cursor = Cursor::new(bytes.to_vec());
9272 let reader = FileReader::try_new(cursor, None).unwrap();
9273 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
9274 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
9275 assert!(total_rows > 0);
9276 assert!(total_rows < 3);
9277 }
9278 }
9279
9280 mod multi_table_transactions {
9282 use super::*;
9283 use futures::TryStreamExt;
9284 use lance::dataset::builder::DatasetBuilder;
9285 use lance_namespace::models::CreateTableVersionRequest;
9286
9287 async fn create_managed_namespace(temp_path: &str) -> Arc<DirectoryNamespace> {
9289 Arc::new(
9290 DirectoryNamespaceBuilder::new(temp_path)
9291 .table_version_tracking_enabled(true)
9292 .table_version_storage_enabled(true)
9293 .manifest_enabled(true)
9294 .build()
9295 .await
9296 .unwrap(),
9297 )
9298 }
9299
9300 async fn create_table_and_get_staging(
9302 namespace: Arc<dyn LanceNamespace>,
9303 table_name: &str,
9304 ) -> (Vec<String>, object_store::path::Path) {
9305 let schema = create_test_schema();
9306 let ipc_data = create_test_ipc_data(&schema);
9307 let mut create_req = CreateTableRequest::new();
9308 create_req.id = Some(vec![table_name.to_string()]);
9309 namespace
9310 .create_table(create_req, bytes::Bytes::from(ipc_data))
9311 .await
9312 .unwrap();
9313
9314 let table_id = vec![table_name.to_string()];
9315 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
9316 .await
9317 .unwrap()
9318 .load()
9319 .await
9320 .unwrap();
9321
9322 let versions_path = dataset.versions_dir();
9324 let manifest_metas: Vec<_> = dataset
9325 .object_store(None)
9326 .await
9327 .unwrap()
9328 .inner
9329 .list(Some(&versions_path))
9330 .try_collect()
9331 .await
9332 .unwrap();
9333
9334 let manifest_meta = manifest_metas
9335 .iter()
9336 .find(|m| {
9337 m.location
9338 .filename()
9339 .map(|f| f.ends_with(".manifest"))
9340 .unwrap_or(false)
9341 })
9342 .expect("No manifest file found");
9343
9344 let manifest_data = dataset
9345 .object_store(None)
9346 .await
9347 .unwrap()
9348 .inner
9349 .get(&manifest_meta.location)
9350 .await
9351 .unwrap()
9352 .bytes()
9353 .await
9354 .unwrap();
9355
9356 let staging_path = dataset
9357 .versions_dir()
9358 .join(format!("staging_{}", table_name));
9359 dataset
9360 .object_store(None)
9361 .await
9362 .unwrap()
9363 .inner
9364 .put(&staging_path, manifest_data.into())
9365 .await
9366 .unwrap();
9367
9368 (table_id, staging_path)
9369 }
9370
9371 #[tokio::test]
9372 async fn test_table_version_storage_enabled_requires_manifest() {
9373 let temp_dir = TempStdDir::default();
9375 let temp_path = temp_dir.to_str().unwrap();
9376
9377 let result = DirectoryNamespaceBuilder::new(temp_path)
9378 .table_version_storage_enabled(true)
9379 .manifest_enabled(false)
9380 .build()
9381 .await;
9382
9383 assert!(
9384 result.is_err(),
9385 "Should fail when table_version_storage_enabled=true but manifest_enabled=false"
9386 );
9387 }
9388
9389 #[tokio::test]
9390 async fn test_create_table_version_records_in_manifest() {
9391 let temp_dir = TempStrDir::default();
9394 let temp_path: &str = &temp_dir;
9395
9396 let namespace = create_managed_namespace(temp_path).await;
9397 let ns: Arc<dyn LanceNamespace> = namespace.clone();
9398
9399 let (table_id, staging_path) =
9400 create_table_and_get_staging(ns.clone(), "table_managed").await;
9401
9402 let mut create_req = CreateTableVersionRequest::new(2, staging_path.to_string());
9404 create_req.id = Some(table_id.clone());
9405 create_req.naming_scheme = Some("V2".to_string());
9406 let response = namespace.create_table_version(create_req).await.unwrap();
9407
9408 assert!(response.version.is_some());
9409 let version = response.version.unwrap();
9410 assert_eq!(version.version, 2);
9411
9412 let manifest_ns = namespace.manifest_ns.as_ref().unwrap();
9414 let table_id_str = manifest::ManifestNamespace::str_object_id(&table_id);
9415 let versions = manifest_ns
9416 .query_table_versions(&table_id_str, false, None)
9417 .await
9418 .unwrap();
9419
9420 assert!(
9421 !versions.is_empty(),
9422 "Version should be recorded in __manifest"
9423 );
9424 let (ver, _path) = &versions[0];
9425 assert_eq!(*ver, 2, "Recorded version should be 2");
9426 }
9427 }
9428
9429 #[tokio::test]
9430 async fn test_list_all_tables() {
9431 use lance_namespace::models::ListTablesRequest;
9432
9433 let (namespace, _temp_dir) = create_test_namespace().await;
9434 create_scalar_table(&namespace, "alpha").await;
9435 create_scalar_table(&namespace, "beta").await;
9436
9437 let request = ListTablesRequest {
9438 id: Some(vec![]),
9439 page_token: None,
9440 limit: None,
9441 ..Default::default()
9442 };
9443 let response = namespace.list_all_tables(request).await.unwrap();
9444 let mut tables = response.tables;
9445 tables.sort();
9446 assert_eq!(tables, vec!["alpha", "beta"]);
9447 }
9448
9449 #[tokio::test]
9450 async fn test_restore_table() {
9451 use lance_namespace::models::RestoreTableRequest;
9452
9453 let (namespace, _temp_dir) = create_test_namespace().await;
9454 create_scalar_table(&namespace, "users").await;
9455
9456 create_scalar_index(&namespace, "users", "users_id_idx").await;
9458
9459 let dataset = open_dataset(&namespace, "users").await;
9460 let current_version = dataset.version().version;
9461 assert!(current_version >= 2, "Should have at least 2 versions");
9462
9463 let mut restore_req = RestoreTableRequest::new(1);
9465 restore_req.id = Some(vec!["users".to_string()]);
9466 let response = namespace.restore_table(restore_req).await.unwrap();
9467
9468 assert!(
9470 response.transaction_id.is_some(),
9471 "restore_table should return a transaction_id"
9472 );
9473
9474 let dataset_after = open_dataset(&namespace, "users").await;
9476 assert!(
9477 dataset_after.version().version > current_version,
9478 "Restore should create a new version"
9479 );
9480 }
9481
9482 #[tokio::test]
9483 async fn test_update_table_schema_metadata() {
9484 use lance_namespace::models::UpdateTableSchemaMetadataRequest;
9485
9486 let (namespace, _temp_dir) = create_test_namespace().await;
9487 create_scalar_table(&namespace, "products").await;
9488
9489 let mut metadata = HashMap::new();
9490 metadata.insert("owner".to_string(), "team_a".to_string());
9491 metadata.insert("version".to_string(), "1.0".to_string());
9492
9493 let mut req = UpdateTableSchemaMetadataRequest::new();
9494 req.id = Some(vec!["products".to_string()]);
9495 req.metadata = Some(metadata.clone());
9496
9497 let response = namespace.update_table_schema_metadata(req).await.unwrap();
9498
9499 assert!(response.metadata.is_some());
9500 let returned = response.metadata.unwrap();
9501 assert_eq!(returned.get("owner"), Some(&"team_a".to_string()));
9502 assert_eq!(returned.get("version"), Some(&"1.0".to_string()));
9503 assert!(
9504 response.transaction_id.is_some(),
9505 "update_table_schema_metadata should return a transaction_id"
9506 );
9507 }
9508
9509 #[tokio::test]
9510 async fn test_get_table_stats() {
9511 use lance_namespace::models::GetTableStatsRequest;
9512
9513 let (namespace, _temp_dir) = create_test_namespace().await;
9514 create_scalar_table(&namespace, "items").await;
9515 create_scalar_index(&namespace, "items", "items_id_idx").await;
9516
9517 let mut req = GetTableStatsRequest::new();
9518 req.id = Some(vec!["items".to_string()]);
9519
9520 let response = namespace.get_table_stats(req).await.unwrap();
9521 assert_eq!(response.num_rows, 3);
9522 assert_eq!(response.num_indices, 1);
9523 }
9524
9525 #[tokio::test]
9526 async fn test_explain_table_query_plan() {
9527 use lance_namespace::models::QueryTableRequestVector;
9528 use lance_namespace::models::{ExplainTableQueryPlanRequest, QueryTableRequest};
9529
9530 let (namespace, _temp_dir) = create_test_namespace().await;
9531 create_scalar_table(&namespace, "catalog").await;
9532
9533 let mut query = QueryTableRequest::new(1, QueryTableRequestVector::new());
9534 query.filter = Some("id > 1".to_string());
9535 query.columns = Some(Box::new(QueryTableRequestColumns {
9536 column_names: Some(vec!["id".to_string(), "name".to_string()]),
9537 column_aliases: None,
9538 }));
9539 query.with_row_id = Some(true);
9540
9541 let mut req = ExplainTableQueryPlanRequest::new(query);
9542 req.id = Some(vec!["catalog".to_string()]);
9543
9544 let plan_str = namespace.explain_table_query_plan(req).await.unwrap();
9545 assert_plan_contains_all(
9546 &plan_str,
9547 &[
9548 "ProjectionExec: expr=[id@0 as id, name@2 as name",
9549 "Take: columns=\"id, _rowid, (name)\"",
9550 "LanceRead: uri=",
9551 "projection=[id]",
9552 "row_id=true, row_addr=false",
9553 "full_filter=id > Int32(1)",
9554 "refine_filter=id > Int32(1)",
9555 ],
9556 "Filtered explain plan should preserve late materialization and filter pushdown",
9557 );
9558 }
9559
9560 #[tokio::test]
9561 async fn test_analyze_table_query_plan() {
9562 use lance_namespace::models::AnalyzeTableQueryPlanRequest;
9563 use lance_namespace::models::QueryTableRequestVector;
9564
9565 let (namespace, _temp_dir) = create_test_namespace().await;
9566 create_scalar_table(&namespace, "catalog").await;
9567
9568 let mut req = AnalyzeTableQueryPlanRequest::new(1, QueryTableRequestVector::new());
9569 req.id = Some(vec!["catalog".to_string()]);
9570 req.filter = Some("id > 0".to_string());
9571 req.columns = Some(Box::new(QueryTableRequestColumns {
9572 column_names: Some(vec!["id".to_string(), "name".to_string()]),
9573 column_aliases: None,
9574 }));
9575 req.with_row_id = Some(true);
9576
9577 let analysis_str = namespace.analyze_table_query_plan(req).await.unwrap();
9578 assert_plan_contains_all(
9579 &analysis_str,
9580 &[
9581 "AnalyzeExec verbose=true",
9582 "ProjectionExec: elapsed=",
9583 "expr=[id@0 as id, name@2 as name",
9584 "Take: elapsed=",
9585 "columns=\"id, _rowid, (name)\"",
9586 "CoalesceBatchesExec: elapsed=",
9587 "LanceRead: elapsed=",
9588 "projection=[id]",
9589 "row_id=true, row_addr=false",
9590 "full_filter=id > Int32(0)",
9591 "refine_filter=id > Int32(0)",
9592 "metrics=[output_rows=",
9593 ],
9594 "Filtered analyze plan should preserve late materialization and filter pushdown",
9595 );
9596 }
9597
9598 #[tokio::test]
9599 async fn test_dir_listing_no_extra_calls_without_migration() {
9600 let temp_dir = TempStdDir::default();
9601 let temp_path = temp_dir.to_str().unwrap();
9602 let root_uri = file_object_store_uri(temp_path);
9603 let listing_count = Arc::new(AtomicUsize::new(0));
9604 let session = build_listing_counting_session(listing_count.clone());
9605
9606 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
9608 .session(session.clone())
9609 .manifest_enabled(false)
9610 .dir_listing_enabled(true)
9611 .build()
9612 .await
9613 .unwrap();
9614
9615 let schema = create_test_schema();
9616 let ipc_data = create_test_ipc_data(&schema);
9617 let mut create_req = CreateTableRequest::new();
9618 create_req.id = Some(vec!["test_table".to_string()]);
9619 dir_only_ns
9620 .create_table(create_req, Bytes::from(ipc_data))
9621 .await
9622 .unwrap();
9623
9624 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
9626 .session(session)
9627 .manifest_enabled(true)
9628 .dir_listing_enabled(true)
9629 .dir_listing_to_manifest_migration_enabled(false)
9630 .build()
9631 .await
9632 .unwrap();
9633
9634 listing_count.store(0, Ordering::SeqCst);
9636
9637 let mut exists_req = TableExistsRequest::new();
9639 exists_req.id = Some(vec!["test_table".to_string()]);
9640 hybrid_ns.table_exists(exists_req).await.unwrap();
9641
9642 let count = listing_count.load(Ordering::SeqCst);
9643 assert_eq!(
9644 count, 1,
9645 "Expected exactly 1 listing call for table_exists \
9646 without migration mode, but got {}",
9647 count
9648 );
9649
9650 listing_count.store(0, Ordering::SeqCst);
9652
9653 let mut describe_req = DescribeTableRequest::new();
9654 describe_req.id = Some(vec!["test_table".to_string()]);
9655 hybrid_ns.describe_table(describe_req).await.unwrap();
9656
9657 let count = listing_count.load(Ordering::SeqCst);
9658 assert_eq!(
9659 count, 1,
9660 "Expected exactly 1 listing call for describe_table \
9661 without migration mode, but got {}",
9662 count
9663 );
9664 }
9665
9666 #[tokio::test]
9667 async fn test_describe_declared_table_checks_versions_only_when_requested() {
9668 let temp_dir = TempStdDir::default();
9669 let temp_path = temp_dir.to_str().unwrap();
9670 let root_uri = file_object_store_uri(temp_path);
9671 let listing_count = Arc::new(AtomicUsize::new(0));
9672 let session = build_listing_counting_session(listing_count.clone());
9673
9674 let namespace = DirectoryNamespaceBuilder::new(root_uri)
9675 .session(session)
9676 .manifest_enabled(false)
9677 .dir_listing_enabled(true)
9678 .build()
9679 .await
9680 .unwrap();
9681
9682 let mut declare_req = DeclareTableRequest::new();
9683 declare_req.id = Some(vec!["test_table".to_string()]);
9684 namespace.declare_table(declare_req).await.unwrap();
9685
9686 listing_count.store(0, Ordering::SeqCst);
9687
9688 let mut describe_req = DescribeTableRequest::new();
9689 describe_req.id = Some(vec!["test_table".to_string()]);
9690 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9691
9692 assert_eq!(describe_response.is_only_declared, None);
9693 assert_eq!(
9694 listing_count.load(Ordering::SeqCst),
9695 1,
9696 "Default describe_table should only list the table directory"
9697 );
9698
9699 listing_count.store(0, Ordering::SeqCst);
9700
9701 let mut describe_req = DescribeTableRequest::new();
9702 describe_req.id = Some(vec!["test_table".to_string()]);
9703 describe_req.check_declared = Some(true);
9704 let describe_response = namespace.describe_table(describe_req).await.unwrap();
9705
9706 assert_eq!(describe_response.is_only_declared, Some(true));
9707 assert_eq!(
9708 listing_count.load(Ordering::SeqCst),
9709 2,
9710 "check_declared describe_table should list the table directory and _versions"
9711 );
9712 }
9713
9714 #[tokio::test]
9715 async fn test_dir_listing_extra_calls_with_migration() {
9716 let temp_dir = TempStdDir::default();
9717 let temp_path = temp_dir.to_str().unwrap();
9718 let root_uri = file_object_store_uri(temp_path);
9719 let listing_count = Arc::new(AtomicUsize::new(0));
9720 let session = build_listing_counting_session(listing_count.clone());
9721
9722 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
9724 .session(session.clone())
9725 .manifest_enabled(false)
9726 .dir_listing_enabled(true)
9727 .build()
9728 .await
9729 .unwrap();
9730
9731 let schema = create_test_schema();
9732 let ipc_data = create_test_ipc_data(&schema);
9733 let mut create_req = CreateTableRequest::new();
9734 create_req.id = Some(vec!["test_table".to_string()]);
9735 dir_only_ns
9736 .create_table(create_req, Bytes::from(ipc_data))
9737 .await
9738 .unwrap();
9739
9740 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
9741 .session(session)
9742 .manifest_enabled(true)
9743 .dir_listing_enabled(true)
9744 .dir_listing_to_manifest_migration_enabled(true)
9745 .build()
9746 .await
9747 .unwrap();
9748
9749 listing_count.store(0, Ordering::SeqCst);
9753
9754 let mut exists_req = TableExistsRequest::new();
9755 exists_req.id = Some(vec!["test_table".to_string()]);
9756 hybrid_ns.table_exists(exists_req).await.unwrap();
9757
9758 let count = listing_count.load(Ordering::SeqCst);
9759 assert_eq!(
9760 count, 1,
9761 "Expected exactly 1 listing call for table_exists with migration mode \
9762 (table directory fallback; manifest reload uses the version hint), but got {}",
9763 count
9764 );
9765
9766 listing_count.store(0, Ordering::SeqCst);
9768
9769 let mut describe_req = DescribeTableRequest::new();
9770 describe_req.id = Some(vec!["test_table".to_string()]);
9771 hybrid_ns.describe_table(describe_req).await.unwrap();
9772
9773 let count = listing_count.load(Ordering::SeqCst);
9774 assert_eq!(
9775 count, 1,
9776 "Expected exactly 1 listing call for describe_table with migration mode \
9777 (table directory fallback; manifest reload uses the version hint), but got {}",
9778 count
9779 );
9780 }
9781
9782 #[tokio::test]
9783 async fn test_migration_not_found_errors_include_table_id() {
9784 let temp_dir = TempStdDir::default();
9785 let temp_path = temp_dir.to_str().unwrap();
9786
9787 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9788 .manifest_enabled(true)
9789 .dir_listing_enabled(true)
9790 .dir_listing_to_manifest_migration_enabled(true)
9791 .build()
9792 .await
9793 .unwrap();
9794
9795 let mut exists_req = TableExistsRequest::new();
9796 exists_req.id = Some(vec!["missing_table".to_string()]);
9797 let err = namespace.table_exists(exists_req).await.unwrap_err();
9798 assert!(matches!(err, Error::Namespace { .. }));
9799 let err_msg = err.to_string();
9800 assert!(err_msg.contains("Table not found"));
9801 assert!(err_msg.contains("table id 'missing_table'"));
9802
9803 let mut describe_req = DescribeTableRequest::new();
9804 describe_req.id = Some(vec!["missing_table".to_string()]);
9805 let err = namespace.describe_table(describe_req).await.unwrap_err();
9806 assert!(matches!(err, Error::Namespace { .. }));
9807 let err_msg = err.to_string();
9808 assert!(err_msg.contains("Table not found"));
9809 assert!(err_msg.contains("table id 'missing_table'"));
9810 }
9811
9812 #[tokio::test]
9813 async fn test_manifest_not_found_errors_include_full_table_id() {
9814 use lance_namespace::models::CreateNamespaceRequest;
9815
9816 let temp_dir = TempStdDir::default();
9817 let temp_path = temp_dir.to_str().unwrap();
9818
9819 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9820 .manifest_enabled(true)
9821 .dir_listing_enabled(true)
9822 .build()
9823 .await
9824 .unwrap();
9825
9826 let mut create_ns_req = CreateNamespaceRequest::new();
9827 create_ns_req.id = Some(vec!["workspace".to_string()]);
9828 namespace.create_namespace(create_ns_req).await.unwrap();
9829
9830 let missing_table_id = vec!["workspace".to_string(), "missing_table".to_string()];
9831
9832 let mut exists_req = TableExistsRequest::new();
9833 exists_req.id = Some(missing_table_id.clone());
9834 let err = namespace.table_exists(exists_req).await.unwrap_err();
9835 assert!(matches!(err, Error::Namespace { .. }));
9836 let err_msg = err.to_string();
9837 assert!(err_msg.contains("Table not found"));
9838 assert!(err_msg.contains("table id 'workspace$missing_table'"));
9839
9840 let mut describe_req = DescribeTableRequest::new();
9841 describe_req.id = Some(missing_table_id);
9842 let err = namespace.describe_table(describe_req).await.unwrap_err();
9843 assert!(matches!(err, Error::Namespace { .. }));
9844 let err_msg = err.to_string();
9845 assert!(err_msg.contains("Table not found"));
9846 assert!(err_msg.contains("table id 'workspace$missing_table'"));
9847 }
9848
9849 async fn create_tagged_test_table(
9852 versions: u32,
9853 ) -> (Arc<DirectoryNamespace>, TempStdDir, Vec<String>) {
9854 use arrow::array::{Int32Array, RecordBatchIterator};
9855 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
9856 use arrow::record_batch::RecordBatch;
9857 use lance::dataset::{Dataset, WriteMode, WriteParams};
9858
9859 assert!(versions >= 1, "versions must be at least 1");
9860
9861 let temp_dir = TempStdDir::default();
9862 let temp_path = temp_dir.to_str().unwrap();
9863
9864 let namespace = Arc::new(
9865 DirectoryNamespaceBuilder::new(temp_path)
9866 .build()
9867 .await
9868 .unwrap(),
9869 );
9870 let table_id = vec!["tag_table".to_string()];
9871 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
9872 "id",
9873 DataType::Int32,
9874 false,
9875 )]));
9876 let initial_batch = RecordBatch::try_new(
9877 arrow_schema.clone(),
9878 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
9879 )
9880 .unwrap();
9881 let batches = RecordBatchIterator::new(vec![Ok(initial_batch)], arrow_schema.clone());
9882 let write_params = WriteParams {
9883 mode: WriteMode::Create,
9884 ..Default::default()
9885 };
9886
9887 let mut dataset = Dataset::write_into_namespace(
9888 batches,
9889 namespace.clone() as Arc<dyn LanceNamespace>,
9890 table_id.clone(),
9891 Some(write_params),
9892 )
9893 .await
9894 .unwrap();
9895
9896 for i in 1..versions {
9897 let value_start = (i as i32) * 10;
9898 let batch = RecordBatch::try_new(
9899 arrow_schema.clone(),
9900 vec![Arc::new(Int32Array::from(vec![
9901 value_start,
9902 value_start + 1,
9903 ]))],
9904 )
9905 .unwrap();
9906 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
9907 dataset.append(batches, None).await.unwrap();
9908 }
9909
9910 (namespace, temp_dir, table_id)
9911 }
9912
9913 #[tokio::test]
9914 async fn test_create_and_list_tags() {
9915 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
9916
9917 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
9918 req.id = Some(table_id.clone());
9919 namespace.create_table_tag(req).await.unwrap();
9920
9921 let mut req = CreateTableTagRequest::new("v2".to_string(), 2);
9922 req.id = Some(table_id.clone());
9923 namespace.create_table_tag(req).await.unwrap();
9924
9925 let mut list_req = ListTableTagsRequest::new();
9926 list_req.id = Some(table_id);
9927 let resp = namespace.list_table_tags(list_req).await.unwrap();
9928
9929 assert_eq!(resp.tags.len(), 2, "expected 2 tags, got: {:?}", resp.tags);
9930 assert_eq!(resp.tags.get("v1").unwrap().version, 1);
9931 assert_eq!(resp.tags.get("v2").unwrap().version, 2);
9932 assert!(resp.page_token.is_none());
9933 }
9934
9935 #[tokio::test]
9936 async fn test_create_existing_tag_conflict() {
9937 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
9938
9939 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
9940 req.id = Some(table_id.clone());
9941 namespace.create_table_tag(req).await.unwrap();
9942
9943 let mut req = CreateTableTagRequest::new("v1".to_string(), 2);
9944 req.id = Some(table_id);
9945 let err = namespace.create_table_tag(req).await.unwrap_err();
9946 assert!(
9947 err.to_string().to_lowercase().contains("already exists"),
9948 "expected already-exists error, got: {}",
9949 err
9950 );
9951 }
9952
9953 #[tokio::test]
9954 async fn test_get_tag_version() {
9955 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
9956
9957 let mut req = CreateTableTagRequest::new("release".to_string(), 2);
9958 req.id = Some(table_id.clone());
9959 namespace.create_table_tag(req).await.unwrap();
9960
9961 let mut get_req = GetTableTagVersionRequest::new("release".to_string());
9962 get_req.id = Some(table_id);
9963 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
9964 assert_eq!(resp.version, 2);
9965 }
9966
9967 #[tokio::test]
9968 async fn test_get_unknown_tag() {
9969 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
9970
9971 let mut get_req = GetTableTagVersionRequest::new("does-not-exist".to_string());
9972 get_req.id = Some(table_id);
9973 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
9974 assert!(
9975 err.to_string().to_lowercase().contains("not found"),
9976 "expected not-found error, got: {}",
9977 err
9978 );
9979 }
9980
9981 #[tokio::test]
9982 async fn test_update_tag_to_new_version() {
9983 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
9984
9985 let mut req = CreateTableTagRequest::new("rolling".to_string(), 1);
9986 req.id = Some(table_id.clone());
9987 namespace.create_table_tag(req).await.unwrap();
9988
9989 let mut update_req = UpdateTableTagRequest::new("rolling".to_string(), 3);
9990 update_req.id = Some(table_id.clone());
9991 namespace.update_table_tag(update_req).await.unwrap();
9992
9993 let mut get_req = GetTableTagVersionRequest::new("rolling".to_string());
9994 get_req.id = Some(table_id);
9995 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
9996 assert_eq!(resp.version, 3);
9997 }
9998
9999 #[tokio::test]
10000 async fn test_update_unknown_tag() {
10001 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
10002
10003 let mut update_req = UpdateTableTagRequest::new("ghost".to_string(), 1);
10004 update_req.id = Some(table_id);
10005 let err = namespace.update_table_tag(update_req).await.unwrap_err();
10006 assert!(
10007 err.to_string().to_lowercase().contains("not found"),
10008 "expected not-found error, got: {}",
10009 err
10010 );
10011 }
10012
10013 #[tokio::test]
10014 async fn test_delete_tag() {
10015 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
10016
10017 let mut req = CreateTableTagRequest::new("doomed".to_string(), 1);
10018 req.id = Some(table_id.clone());
10019 namespace.create_table_tag(req).await.unwrap();
10020
10021 let mut delete_req = DeleteTableTagRequest::new("doomed".to_string());
10022 delete_req.id = Some(table_id.clone());
10023 namespace.delete_table_tag(delete_req).await.unwrap();
10024
10025 let mut list_req = ListTableTagsRequest::new();
10026 list_req.id = Some(table_id.clone());
10027 let resp = namespace.list_table_tags(list_req).await.unwrap();
10028 assert!(resp.tags.is_empty(), "tag should be removed after delete");
10029
10030 let mut get_req = GetTableTagVersionRequest::new("doomed".to_string());
10032 get_req.id = Some(table_id);
10033 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
10034 assert!(err.to_string().to_lowercase().contains("not found"));
10035 }
10036
10037 #[tokio::test]
10038 async fn test_delete_unknown_tag() {
10039 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
10040
10041 let mut delete_req = DeleteTableTagRequest::new("nope".to_string());
10042 delete_req.id = Some(table_id);
10043 let err = namespace.delete_table_tag(delete_req).await.unwrap_err();
10044 assert!(
10045 err.to_string().to_lowercase().contains("not found"),
10046 "expected not-found error, got: {}",
10047 err
10048 );
10049 }
10050
10051 #[tokio::test]
10052 async fn test_create_tag_invalid_version() {
10053 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
10054
10055 let mut req = CreateTableTagRequest::new("v0".to_string(), 0);
10057 req.id = Some(table_id.clone());
10058 let err = namespace.create_table_tag(req).await.unwrap_err();
10059 assert!(
10060 err.to_string().to_lowercase().contains("positive"),
10061 "expected positive-version error, got: {}",
10062 err
10063 );
10064
10065 let mut req = CreateTableTagRequest::new(String::new(), 1);
10067 req.id = Some(table_id);
10068 let err = namespace.create_table_tag(req).await.unwrap_err();
10069 assert!(
10070 err.to_string().to_lowercase().contains("must not be empty"),
10071 "expected empty-tag-name error, got: {}",
10072 err
10073 );
10074 }
10075
10076 #[tokio::test]
10077 async fn test_create_tag_table_not_found() {
10078 let (namespace, _temp_dir) = create_test_namespace().await;
10079
10080 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
10081 req.id = Some(vec!["does_not_exist".to_string()]);
10082 let err = namespace.create_table_tag(req).await.unwrap_err();
10083 let msg = err.to_string();
10084 assert!(
10085 msg.contains("Table") && msg.to_lowercase().contains("not found"),
10086 "expected TableNotFound error, got: {}",
10087 err
10088 );
10089 }
10090}