1pub mod manifest;
10pub mod manifest_feature_flags;
11
12use arrow::array::Float32Array;
13use arrow::record_batch::RecordBatchIterator;
14use arrow_ipc::reader::StreamReader;
15use async_trait::async_trait;
16use bytes::Bytes;
17use futures::{StreamExt, TryStreamExt};
18use lance::dataset::builder::DatasetBuilder;
19use lance::dataset::refs::check_valid_branch;
20use lance::dataset::scanner::Scanner;
21use lance::dataset::statistics::DatasetStatisticsExt;
22use lance::dataset::transaction::{Operation, Transaction};
23use lance::dataset::{
24 Dataset, MergeInsertBuilder, UpdateBuilder, WhenMatched, WhenNotMatched,
25 WhenNotMatchedBySource, WriteMode, WriteParams,
26};
27use lance::index::{DatasetIndexExt, IndexParams, vector::VectorIndexParams};
28use lance::session::Session;
29use lance_index::scalar::{
30 BuiltinIndexType, FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams,
31};
32use lance_index::vector::{
33 bq::RQBuildParams, hnsw::builder::HnswBuildParams, ivf::IvfBuildParams, pq::PQBuildParams,
34 sq::builder::SQBuildParams,
35};
36use lance_index::{IndexType, is_system_index};
37use lance_io::object_store::throttle::is_throttle_error;
38use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
39use lance_linalg::distance::MetricType;
40use lance_table::io::commit::{ManifestNamingScheme, VERSIONS_DIR};
41use object_store::ObjectStoreExt;
42use object_store::path::Path;
43use object_store::{
44 Error as ObjectStoreError, ObjectMeta, ObjectStore as OSObjectStore, PutMode, PutOptions,
45};
46use std::collections::HashMap;
47use std::io::Cursor;
48use std::sync::{Arc, Mutex};
49use tokio::sync::OnceCell;
50
51use crate::context::DynamicContextProvider;
52use lance_namespace::models::{
53 AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest,
54 AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse,
55 AlterTransactionRequest, AlterTransactionResponse, AnalyzeTableQueryPlanRequest,
56 BatchDeleteTableVersionsRequest, BatchDeleteTableVersionsResponse,
57 BranchContents as ModelBranchContents, CountTableRowsRequest, CreateNamespaceRequest,
58 CreateNamespaceResponse, CreateTableBranchRequest, CreateTableBranchResponse,
59 CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest, CreateTableResponse,
60 CreateTableScalarIndexResponse, CreateTableTagRequest, CreateTableTagResponse,
61 CreateTableVersionRequest, CreateTableVersionResponse, DeclareTableRequest,
62 DeclareTableResponse, DeleteFromTableRequest, DeleteFromTableResponse,
63 DeleteTableBranchRequest, DeleteTableBranchResponse, DeleteTableTagRequest,
64 DeleteTableTagResponse, DescribeNamespaceRequest, DescribeNamespaceResponse,
65 DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, DescribeTableRequest,
66 DescribeTableResponse, DescribeTableVersionRequest, DescribeTableVersionResponse,
67 DescribeTransactionRequest, DescribeTransactionResponse, DropNamespaceRequest,
68 DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, DropTableRequest,
69 DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary,
70 GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest,
71 GetTableTagVersionResponse, Identity, IndexContent, InsertIntoTableRequest,
72 InsertIntoTableResponse, ListNamespacesRequest, ListNamespacesResponse,
73 ListTableBranchesRequest, ListTableBranchesResponse, ListTableIndicesRequest,
74 ListTableIndicesResponse, ListTableTagsRequest, ListTableTagsResponse,
75 ListTableVersionsRequest, ListTableVersionsResponse, ListTablesRequest, ListTablesResponse,
76 MergeInsertIntoTableRequest, MergeInsertIntoTableResponse, NamespaceExistsRequest,
77 QueryTableRequest, QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest,
78 RestoreTableResponse, TableExistsRequest, TableVersion, TagContents as ModelTagContents,
79 UpdateTableRequest, UpdateTableResponse, UpdateTableSchemaMetadataRequest,
80 UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, UpdateTableTagResponse,
81};
82
83use lance_core::{Error, Result, box_error};
84use lance_index::scalar::inverted::query::{
85 BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery,
86};
87use lance_namespace::LanceNamespace;
88use lance_namespace::error::NamespaceError;
89use lance_namespace::schema::arrow_schema_to_json;
90
91use crate::credentials::{
92 CredentialVendor, create_credential_vendor_for_location, has_credential_vendor_config,
93};
94
95#[derive(Debug, Default)]
100pub struct OpsMetrics {
101 counters: Mutex<HashMap<String, u64>>,
102}
103
104impl OpsMetrics {
105 pub fn increment(&self, operation: &str) {
107 if let Ok(mut counters) = self.counters.lock() {
108 *counters.entry(operation.to_string()).or_insert(0) += 1;
109 }
110 }
111
112 pub fn retrieve(&self) -> HashMap<String, u64> {
114 self.counters.lock().map(|c| c.clone()).unwrap_or_default()
115 }
116
117 pub fn reset(&self) {
119 if let Ok(mut counters) = self.counters.lock() {
120 counters.clear();
121 }
122 }
123}
124
125pub(crate) fn build_sql_expressions(
128 new_columns: &[lance_namespace::models::AddColumnsEntry],
129) -> Result<Vec<(String, String)>> {
130 new_columns
131 .iter()
132 .map(|col| {
133 let expression = col.expression.clone().and_then(|opt| opt).ok_or_else(|| {
136 Error::invalid_input(format!(
137 "Expression is required for new column '{}'",
138 col.name
139 ))
140 })?;
141 Ok((col.name.clone(), expression))
142 })
143 .collect()
144}
145
146pub(crate) fn build_column_alterations(
149 alterations: &[lance_namespace::models::AlterColumnsEntry],
150) -> Result<Vec<lance::dataset::ColumnAlteration>> {
151 alterations
152 .iter()
153 .map(|entry| {
154 let mut alteration = lance::dataset::ColumnAlteration::new(entry.path.clone());
155 if let Some(Some(rename)) = &entry.rename {
157 alteration = alteration.rename(rename.clone());
158 }
159 if let Some(Some(nullable)) = entry.nullable {
161 alteration = alteration.set_nullable(nullable);
162 }
163 if let Some(data_type) = &entry.data_type
165 && !data_type.is_null()
166 {
167 let type_str = data_type.as_str().ok_or_else(|| {
168 Error::invalid_input(format!(
169 "data_type for column '{}' must be a JSON string, got: {}",
170 entry.path, data_type
171 ))
172 })?;
173 let json_type =
174 lance_namespace::models::JsonArrowDataType::new(type_str.to_string());
175 let dt =
176 lance_namespace::schema::convert_json_arrow_type(&json_type).map_err(|e| {
177 Error::invalid_input(format!(
178 "Failed to parse data_type '{}' for column '{}': {}",
179 type_str, entry.path, e
180 ))
181 })?;
182 alteration = alteration.cast_to(dt);
183 }
184 Ok(alteration)
185 })
186 .collect()
187}
188
189pub(crate) struct TableStatus {
194 pub(crate) exists: bool,
196 pub(crate) is_deregistered: bool,
198 pub(crate) has_reserved_file: bool,
200}
201
202enum DirectoryIndexParams {
203 Scalar {
204 index_type: IndexType,
205 params: ScalarIndexParams,
206 },
207 Inverted(InvertedIndexParams),
208 Vector {
209 index_type: IndexType,
210 params: VectorIndexParams,
211 },
212}
213
214impl DirectoryIndexParams {
215 fn index_type(&self) -> IndexType {
216 match self {
217 Self::Scalar { index_type, .. } | Self::Vector { index_type, .. } => *index_type,
218 Self::Inverted(_) => IndexType::Inverted,
219 }
220 }
221
222 fn params(&self) -> &dyn IndexParams {
223 match self {
224 Self::Scalar { params, .. } => params,
225 Self::Inverted(params) => params,
226 Self::Vector { params, .. } => params,
227 }
228 }
229}
230
231#[derive(Clone)]
266pub struct DirectoryNamespaceBuilder {
267 root: String,
268 storage_options: Option<HashMap<String, String>>,
269 session: Option<Arc<Session>>,
270 manifest_enabled: bool,
271 dir_listing_enabled: bool,
272 inline_optimization_enabled: bool,
273 table_version_tracking_enabled: bool,
274 dir_listing_to_manifest_migration_enabled: bool,
279 credential_vendor_properties: HashMap<String, String>,
280 context_provider: Option<Arc<dyn DynamicContextProvider>>,
281 commit_retries: Option<u32>,
282 vend_input_storage_options: bool,
285 vend_input_storage_options_refresh_interval_millis: Option<u64>,
290 ops_metrics_enabled: bool,
292}
293
294impl std::fmt::Debug for DirectoryNamespaceBuilder {
295 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
296 f.debug_struct("DirectoryNamespaceBuilder")
297 .field("root", &self.root)
298 .field("storage_options", &self.storage_options)
299 .field("manifest_enabled", &self.manifest_enabled)
300 .field("dir_listing_enabled", &self.dir_listing_enabled)
301 .field(
302 "inline_optimization_enabled",
303 &self.inline_optimization_enabled,
304 )
305 .field(
306 "table_version_tracking_enabled",
307 &self.table_version_tracking_enabled,
308 )
309 .field(
310 "dir_listing_to_manifest_migration_enabled",
311 &self.dir_listing_to_manifest_migration_enabled,
312 )
313 .field(
314 "context_provider",
315 &self.context_provider.as_ref().map(|_| "Some(...)"),
316 )
317 .field(
318 "vend_input_storage_options",
319 &self.vend_input_storage_options,
320 )
321 .field(
322 "vend_input_storage_options_refresh_interval_millis",
323 &self.vend_input_storage_options_refresh_interval_millis,
324 )
325 .field("ops_metrics_enabled", &self.ops_metrics_enabled)
326 .finish()
327 }
328}
329
330impl DirectoryNamespaceBuilder {
331 pub fn new(root: impl Into<String>) -> Self {
337 Self {
338 root: root.into().trim_end_matches('/').to_string(),
339 storage_options: None,
340 session: None,
341 manifest_enabled: true,
342 dir_listing_enabled: true, inline_optimization_enabled: true,
344 table_version_tracking_enabled: false, dir_listing_to_manifest_migration_enabled: false, credential_vendor_properties: HashMap::new(),
347 context_provider: None,
348 commit_retries: None,
349 vend_input_storage_options: false,
350 vend_input_storage_options_refresh_interval_millis: None,
351 ops_metrics_enabled: false,
352 }
353 }
354
355 pub fn manifest_enabled(mut self, enabled: bool) -> Self {
360 self.manifest_enabled = enabled;
361 self
362 }
363
364 pub fn dir_listing_enabled(mut self, enabled: bool) -> Self {
369 self.dir_listing_enabled = enabled;
370 self
371 }
372
373 pub fn dir_listing_to_manifest_migration_enabled(mut self, enabled: bool) -> Self {
380 self.dir_listing_to_manifest_migration_enabled = enabled;
381 self
382 }
383
384 pub fn inline_optimization_enabled(mut self, enabled: bool) -> Self {
389 self.inline_optimization_enabled = enabled;
390 self
391 }
392
393 pub fn table_version_tracking_enabled(mut self, enabled: bool) -> Self {
401 self.table_version_tracking_enabled = enabled;
402 self
403 }
404
405 pub fn from_properties(
473 properties: HashMap<String, String>,
474 session: Option<Arc<Session>>,
475 ) -> Result<Self> {
476 let root = properties.get("root").cloned().ok_or_else(|| {
478 lance_core::Error::from(NamespaceError::InvalidInput {
479 message: "Missing required property 'root' for directory namespace".to_string(),
480 })
481 })?;
482
483 let storage_options: HashMap<String, String> = properties
485 .iter()
486 .filter_map(|(k, v)| {
487 k.strip_prefix("storage.")
488 .map(|key| (key.to_string(), v.clone()))
489 })
490 .collect();
491
492 let storage_options = if storage_options.is_empty() {
493 None
494 } else {
495 Some(storage_options)
496 };
497
498 let manifest_enabled = properties
500 .get("manifest_enabled")
501 .and_then(|v| v.parse::<bool>().ok())
502 .unwrap_or(true);
503
504 let dir_listing_enabled = properties
506 .get("dir_listing_enabled")
507 .and_then(|v| v.parse::<bool>().ok())
508 .unwrap_or(true);
509
510 let inline_optimization_enabled = properties
512 .get("inline_optimization_enabled")
513 .and_then(|v| v.parse::<bool>().ok())
514 .unwrap_or(true);
515
516 let table_version_tracking_enabled = properties
518 .get("table_version_tracking_enabled")
519 .and_then(|v| v.parse::<bool>().ok())
520 .unwrap_or(false);
521
522 let dir_listing_to_manifest_migration_enabled = properties
524 .get("dir_listing_to_manifest_migration_enabled")
525 .and_then(|v| v.parse::<bool>().ok())
526 .unwrap_or(false);
527
528 let credential_vendor_properties: HashMap<String, String> = properties
532 .iter()
533 .filter_map(|(k, v)| {
534 k.strip_prefix("credential_vendor.")
535 .map(|key| (key.to_string(), v.clone()))
536 })
537 .collect();
538
539 let commit_retries = properties
540 .get("commit_retries")
541 .and_then(|v| v.parse::<u32>().ok());
542
543 let vend_input_storage_options = properties
545 .get("vend_input_storage_options")
546 .and_then(|v| v.parse::<bool>().ok())
547 .unwrap_or(false);
548
549 let vend_input_storage_options_refresh_interval_millis = properties
551 .get("vend_input_storage_options_refresh_interval_millis")
552 .and_then(|v| v.parse::<u64>().ok());
553
554 let ops_metrics_enabled = properties
556 .get("ops_metrics_enabled")
557 .and_then(|v| v.parse::<bool>().ok())
558 .unwrap_or(false);
559
560 Ok(Self {
561 root: root.trim_end_matches('/').to_string(),
562 storage_options,
563 session,
564 manifest_enabled,
565 dir_listing_enabled,
566 inline_optimization_enabled,
567 table_version_tracking_enabled,
568 dir_listing_to_manifest_migration_enabled,
569 credential_vendor_properties,
570 context_provider: None,
571 commit_retries,
572 vend_input_storage_options,
573 vend_input_storage_options_refresh_interval_millis,
574 ops_metrics_enabled,
575 })
576 }
577
578 pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
585 self.storage_options
586 .get_or_insert_with(HashMap::new)
587 .insert(key.into(), value.into());
588 self
589 }
590
591 pub fn storage_options(mut self, options: HashMap<String, String>) -> Self {
597 self.storage_options
598 .get_or_insert_with(HashMap::new)
599 .extend(options);
600 self
601 }
602
603 pub fn session(mut self, session: Arc<Session>) -> Self {
613 self.session = Some(session);
614 self
615 }
616
617 pub fn commit_retries(mut self, retries: u32) -> Self {
620 self.commit_retries = Some(retries);
621 self
622 }
623
624 pub fn credential_vendor_property(
652 mut self,
653 key: impl Into<String>,
654 value: impl Into<String>,
655 ) -> Self {
656 self.credential_vendor_properties
657 .insert(key.into(), value.into());
658 self
659 }
660
661 pub fn credential_vendor_properties(mut self, properties: HashMap<String, String>) -> Self {
669 self.credential_vendor_properties.extend(properties);
670 self
671 }
672
673 pub fn context_provider(mut self, provider: Arc<dyn DynamicContextProvider>) -> Self {
683 self.context_provider = Some(provider);
684 self
685 }
686
687 pub fn vend_input_storage_options(mut self, enabled: bool) -> Self {
696 self.vend_input_storage_options = enabled;
697 self
698 }
699
700 pub fn vend_input_storage_options_refresh_interval_millis(
712 mut self,
713 interval_millis: u64,
714 ) -> Self {
715 self.vend_input_storage_options_refresh_interval_millis = Some(interval_millis);
716 self
717 }
718
719 pub fn ops_metrics_enabled(mut self, enabled: bool) -> Self {
727 self.ops_metrics_enabled = enabled;
728 self
729 }
730
731 pub async fn build(self) -> Result<DirectoryNamespace> {
744 let (object_store, base_path) =
745 Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?;
746
747 let manifest_ns = if self.manifest_enabled {
748 match manifest::ManifestNamespace::open_from_directory(
749 self.root.clone(),
750 self.storage_options.clone(),
751 self.session.clone(),
752 object_store.clone(),
753 base_path.clone(),
754 self.dir_listing_enabled,
755 self.inline_optimization_enabled,
756 self.commit_retries,
757 )
758 .await
759 {
760 Ok(ns) => Some(Arc::new(ns)),
761 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
762 return Err(e);
766 }
767 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => {
768 log::debug!("Manifest namespace does not exist yet: {}", e);
769 None
770 }
771 Err(e) => return Err(e),
772 }
773 } else {
774 None
775 };
776 let manifest_cell = OnceCell::new();
777 if let Some(manifest_ns) = manifest_ns {
778 let _ = manifest_cell.set(manifest_ns);
779 }
780
781 let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties)
783 {
784 create_credential_vendor_for_location(&self.root, &self.credential_vendor_properties)
785 .await?
786 .map(Arc::from)
787 } else {
788 None
789 };
790
791 let ops_metrics = if self.ops_metrics_enabled {
792 Some(Arc::new(OpsMetrics::default()))
793 } else {
794 None
795 };
796
797 Ok(DirectoryNamespace {
798 root: self.root,
799 storage_options: self.storage_options,
800 session: self.session,
801 object_store,
802 base_path,
803 manifest_ns: manifest_cell,
804 write_manifest_ns: OnceCell::new(),
805 manifest_enabled: self.manifest_enabled,
806 dir_listing_enabled: self.dir_listing_enabled,
807 inline_optimization_enabled: self.inline_optimization_enabled,
808 commit_retries: self.commit_retries,
809 dir_listing_to_manifest_migration_enabled: self
810 .dir_listing_to_manifest_migration_enabled,
811 table_version_tracking_enabled: self.table_version_tracking_enabled,
812 credential_vendor,
813 context_provider: self.context_provider,
814 vend_input_storage_options: self.vend_input_storage_options,
815 vend_input_storage_options_refresh_interval_millis: self
816 .vend_input_storage_options_refresh_interval_millis,
817 ops_metrics,
818 })
819 }
820
821 async fn initialize_object_store(
823 root: &str,
824 storage_options: &Option<HashMap<String, String>>,
825 session: &Option<Arc<Session>>,
826 ) -> Result<(Arc<ObjectStore>, Path)> {
827 let accessor = storage_options.clone().map(|opts| {
829 Arc::new(lance_io::object_store::StorageOptionsAccessor::with_static_options(opts))
830 });
831 let params = ObjectStoreParams {
832 storage_options_accessor: accessor,
833 ..Default::default()
834 };
835
836 let registry = if let Some(session) = session {
838 session.store_registry()
839 } else {
840 Arc::new(ObjectStoreRegistry::default())
841 };
842
843 let (object_store, base_path) = ObjectStore::from_uri_and_params(registry, root, ¶ms)
845 .await
846 .map_err(|e| {
847 lance_core::Error::from(NamespaceError::Internal {
848 message: format!("Failed to create object store: {:?}", e),
849 })
850 })?;
851
852 Ok((object_store, base_path))
853 }
854}
855
856pub struct DirectoryNamespace {
880 root: String,
881 storage_options: Option<HashMap<String, String>>,
882 session: Option<Arc<Session>>,
883 object_store: Arc<ObjectStore>,
884 base_path: Path,
885 manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
886 write_manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
887 manifest_enabled: bool,
888 dir_listing_enabled: bool,
889 inline_optimization_enabled: bool,
890 commit_retries: Option<u32>,
891 dir_listing_to_manifest_migration_enabled: bool,
895 table_version_tracking_enabled: bool,
898 credential_vendor: Option<Arc<dyn CredentialVendor>>,
901 #[allow(dead_code)]
904 context_provider: Option<Arc<dyn DynamicContextProvider>>,
905 vend_input_storage_options: bool,
907 vend_input_storage_options_refresh_interval_millis: Option<u64>,
910 ops_metrics: Option<Arc<OpsMetrics>>,
912}
913
914impl std::fmt::Debug for DirectoryNamespace {
915 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
916 write!(f, "{}", self.namespace_id())
917 }
918}
919
920impl std::fmt::Display for DirectoryNamespace {
921 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
922 write!(f, "{}", self.namespace_id())
923 }
924}
925
926struct ExistingTableVersionResolve<'a> {
928 staging_path: &'a Path,
929 final_path: &'a Path,
930 version: u64,
931 table_uri: &'a str,
932 final_meta: &'a ObjectMeta,
933 request_manifest_size: Option<i64>,
934}
935
936struct TableDeleteEntry {
939 table_id: Option<Vec<String>>,
940 ranges: Vec<(i64, i64)>,
941}
942
943#[derive(Debug, Clone, Default)]
953struct TransactionAlteration {
954 status: Option<String>,
956 properties: HashMap<String, String>,
960 removed_properties: std::collections::HashSet<String>,
964}
965
966impl TransactionAlteration {
967 const F_STATUS: &'static str = "status";
969 const F_PROPERTIES: &'static str = "properties";
970 const F_REMOVED_PROPERTIES: &'static str = "removed_properties";
971
972 fn to_json_bytes(&self) -> serde_json::Result<Vec<u8>> {
979 serde_json::to_vec(&serde_json::json!({
980 Self::F_STATUS: self.status,
981 Self::F_PROPERTIES: self.properties,
982 Self::F_REMOVED_PROPERTIES: self.removed_properties,
983 }))
984 }
985
986 fn from_json_slice(bytes: &[u8]) -> serde_json::Result<Self> {
991 let mut obj: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes)?;
992 Ok(Self {
993 status: serde_json::from_value(
994 obj.remove(Self::F_STATUS)
995 .unwrap_or(serde_json::Value::Null),
996 )?,
997 properties: serde_json::from_value(
998 obj.remove(Self::F_PROPERTIES)
999 .unwrap_or(serde_json::Value::Null),
1000 )
1001 .unwrap_or_default(),
1002 removed_properties: serde_json::from_value(
1003 obj.remove(Self::F_REMOVED_PROPERTIES)
1004 .unwrap_or(serde_json::Value::Null),
1005 )
1006 .unwrap_or_default(),
1007 })
1008 }
1009}
1010
1011impl DirectoryNamespace {
1012 fn manifest_ns_for_read(&self) -> Option<&Arc<manifest::ManifestNamespace>> {
1013 self.write_manifest_ns
1014 .get()
1015 .or_else(|| self.manifest_ns.get())
1016 }
1017
1018 async fn manifest_ns_for_write(&self) -> Result<Option<Arc<manifest::ManifestNamespace>>> {
1019 if !self.manifest_enabled {
1020 return Ok(None);
1021 }
1022
1023 let manifest_ns = self
1024 .write_manifest_ns
1025 .get_or_try_init(|| async {
1026 manifest::ManifestNamespace::from_directory(
1027 self.root.clone(),
1028 self.storage_options.clone(),
1029 self.session.clone(),
1030 self.object_store.clone(),
1031 self.base_path.clone(),
1032 self.dir_listing_enabled,
1033 self.inline_optimization_enabled,
1034 self.commit_retries,
1035 )
1036 .await
1037 .map(Arc::new)
1038 })
1039 .await?;
1040 Ok(Some(manifest_ns.clone()))
1041 }
1042
1043 fn child_namespace_requires_manifest_error(&self) -> Error {
1044 if self.manifest_enabled {
1045 NamespaceError::NamespaceNotFound {
1046 message: "Child namespace reads require an existing __manifest dataset".to_string(),
1047 }
1048 .into()
1049 } else {
1050 NamespaceError::Unsupported {
1051 message: "Child namespaces are only supported when manifest mode is enabled"
1052 .to_string(),
1053 }
1054 .into()
1055 }
1056 }
1057
1058 fn apply_pagination(
1071 names: &mut Vec<String>,
1072 page_token: Option<String>,
1073 limit: Option<i32>,
1074 ) -> Option<String> {
1075 names.sort();
1077
1078 if let Some(start_after) = page_token {
1080 if let Some(index) = names
1081 .iter()
1082 .position(|name| name.as_str() > start_after.as_str())
1083 {
1084 names.drain(0..index);
1085 } else {
1086 names.clear();
1087 }
1088 }
1089
1090 if let Some(limit) = limit
1092 && limit >= 0
1093 {
1094 let limit = limit as usize;
1095 if names.len() > limit {
1096 let next_page_token = if limit > 0 {
1097 Some(names[limit - 1].clone())
1098 } else {
1099 None
1100 };
1101 names.truncate(limit);
1102 return next_page_token;
1103 }
1104 }
1105
1106 None
1107 }
1108
1109 async fn list_directory_tables(&self) -> Result<Vec<String>> {
1111 let mut tables = Vec::new();
1112 let entries = self
1113 .object_store
1114 .read_dir(self.base_path.clone())
1115 .await
1116 .map_err(|e| {
1117 lance_core::Error::from(NamespaceError::Internal {
1118 message: format!("Failed to list directory: {:?}", e),
1119 })
1120 })?;
1121
1122 for entry in entries {
1123 let path = entry.trim_end_matches('/');
1124 if !path.ends_with(".lance") {
1125 continue;
1126 }
1127
1128 let table_name = &path[..path.len() - 6];
1129
1130 let status = self.check_table_status(table_name).await?;
1132 if status.is_deregistered {
1133 continue;
1134 }
1135
1136 tables.push(table_name.to_string());
1137 }
1138
1139 Ok(tables)
1140 }
1141
1142 fn validate_root_namespace_id(id: &Option<Vec<String>>) -> Result<()> {
1144 if let Some(id) = id
1145 && !id.is_empty()
1146 {
1147 return Err(NamespaceError::Unsupported {
1148 message: format!(
1149 "Directory namespace only supports root namespace operations, but got namespace ID: {:?}. Expected empty ID.",
1150 id
1151 ),
1152 }
1153 .into());
1154 }
1155 Ok(())
1156 }
1157
1158 fn table_name_from_id(id: &Option<Vec<String>>) -> Result<String> {
1160 let id = id.as_ref().ok_or_else(|| {
1161 lance_core::Error::from(NamespaceError::InvalidInput {
1162 message: "Directory namespace table ID cannot be empty".to_string(),
1163 })
1164 })?;
1165
1166 if id.len() != 1 {
1167 return Err(NamespaceError::Unsupported {
1168 message: format!(
1169 "Multi-level table IDs are only supported when manifest mode is enabled, but got: {:?}",
1170 id
1171 ),
1172 }
1173 .into());
1174 }
1175
1176 Ok(id[0].clone())
1177 }
1178
1179 fn format_table_id(table_id: &[String]) -> String {
1180 format!(
1181 "table id '{}'",
1182 manifest::ManifestNamespace::str_object_id(table_id)
1183 )
1184 }
1185
1186 fn format_table_id_from_request(id: &Option<Vec<String>>) -> String {
1187 id.as_ref()
1188 .map(|table_id| Self::format_table_id(table_id))
1189 .unwrap_or_else(|| "table id '<unknown>'".to_string())
1190 }
1191
1192 async fn resolve_table_location(&self, id: &Option<Vec<String>>) -> Result<String> {
1193 let mut describe_req = DescribeTableRequest::new();
1194 describe_req.id = id.clone();
1195 describe_req.load_detailed_metadata = Some(false);
1196
1197 let describe_resp = self.describe_table_impl(describe_req).await?;
1199
1200 describe_resp.location.ok_or_else(|| {
1201 lance_core::Error::from(NamespaceError::TableNotFound {
1202 message: format!("Table location not found for: {:?}", id),
1203 })
1204 })
1205 }
1206
1207 fn map_tag_error(err: lance_core::Error, tag: &str, table_uri: &str) -> lance_core::Error {
1210 match err {
1211 lance_core::Error::RefNotFound { .. } => NamespaceError::TableTagNotFound {
1212 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1213 }
1214 .into(),
1215 lance_core::Error::RefConflict { .. } => NamespaceError::TableTagAlreadyExists {
1216 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1217 }
1218 .into(),
1219 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1220 message: format!("invalid tag '{}': {}", tag, message),
1221 }
1222 .into(),
1223 lance_core::Error::VersionNotFound { message } => {
1224 NamespaceError::TableVersionNotFound {
1225 message: format!(
1226 "version referenced by tag '{}' not found for table at '{}': {}",
1227 tag, table_uri, message
1228 ),
1229 }
1230 .into()
1231 }
1232 other => NamespaceError::Internal {
1233 message: format!(
1234 "tag operation failed for tag '{}' on table at '{}': {}",
1235 tag, table_uri, other
1236 ),
1237 }
1238 .into(),
1239 }
1240 }
1241
1242 fn map_branch_error(
1248 err: lance_core::Error,
1249 branch: &str,
1250 table_uri: &str,
1251 ) -> lance_core::Error {
1252 match err {
1253 lance_core::Error::RefNotFound { .. } => NamespaceError::TableBranchNotFound {
1254 message: format!("branch '{}' for table at '{}'", branch, table_uri),
1255 }
1256 .into(),
1257 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1258 message: format!("invalid branch '{}': {}", branch, message),
1259 }
1260 .into(),
1261 lance_core::Error::VersionNotFound { message } => {
1262 NamespaceError::TableVersionNotFound {
1263 message: format!(
1264 "source version for branch '{}' not found for table at '{}': {}",
1265 branch, table_uri, message
1266 ),
1267 }
1268 .into()
1269 }
1270 other => NamespaceError::Internal {
1271 message: format!(
1272 "branch operation failed for branch '{}' on table at '{}': {}",
1273 branch, table_uri, other
1274 ),
1275 }
1276 .into(),
1277 }
1278 }
1279
1280 fn map_mutation_error(
1292 err: lance_core::Error,
1293 operation: &str,
1294 table_uri: &str,
1295 ) -> lance_core::Error {
1296 let detail = err.to_string();
1297 let ns_err = match &err {
1298 lance_core::Error::InvalidInput { .. }
1299 | lance_core::Error::Unprocessable { .. }
1300 | lance_core::Error::InvalidRef { .. } => NamespaceError::InvalidInput {
1301 message: format!(
1302 "Invalid input for {} on table at '{}': {}",
1303 operation, table_uri, detail
1304 ),
1305 },
1306 lance_core::Error::NotFound { .. } | lance_core::Error::DatasetNotFound { .. } => {
1307 NamespaceError::TableNotFound {
1308 message: format!(
1309 "Table at '{}' not found while running {}: {}",
1310 table_uri, operation, detail
1311 ),
1312 }
1313 }
1314 lance_core::Error::SchemaMismatch { .. } | lance_core::Error::Schema { .. } => {
1315 NamespaceError::TableSchemaValidationError {
1316 message: format!(
1317 "Schema validation failed for {} on table at '{}': {}",
1318 operation, table_uri, detail
1319 ),
1320 }
1321 }
1322 lance_core::Error::CommitConflict { .. } => NamespaceError::Throttling {
1326 message: format!(
1327 "Too many concurrent writes for {} on table at '{}', please retry later: {}",
1328 operation, table_uri, detail
1329 ),
1330 },
1331 lance_core::Error::TooMuchWriteContention { .. }
1335 | lance_core::Error::RetryableCommitConflict { .. }
1336 | lance_core::Error::IncompatibleTransaction { .. }
1337 | lance_core::Error::VersionConflict { .. } => NamespaceError::ConcurrentModification {
1338 message: format!(
1339 "Concurrent modification detected for {} on table at '{}': {}",
1340 operation, table_uri, detail
1341 ),
1342 },
1343 lance_core::Error::NotSupported { .. } => NamespaceError::Unsupported {
1344 message: format!(
1345 "{} is not supported on table at '{}': {}",
1346 operation, table_uri, detail
1347 ),
1348 },
1349 _ => NamespaceError::Internal {
1350 message: format!(
1351 "Failed to run {} on table at '{}': {}",
1352 operation, table_uri, detail
1353 ),
1354 },
1355 };
1356 ns_err.into()
1357 }
1358
1359 async fn table_has_actual_manifests(&self, table_name: &str) -> Result<bool> {
1360 manifest::ManifestNamespace::path_has_actual_manifests(
1361 &self.object_store,
1362 &self.table_path(table_name),
1363 )
1364 .await
1365 }
1366
1367 async fn filter_declared_tables(
1368 &self,
1369 tables: Vec<String>,
1370 include_declared: bool,
1371 ) -> Result<Vec<String>> {
1372 if include_declared {
1373 return Ok(tables);
1374 }
1375
1376 let mut stream = futures::stream::iter(tables.into_iter().map(|table_name| async move {
1377 if self.table_has_actual_manifests(&table_name).await? {
1381 Ok::<Option<String>, Error>(Some(table_name))
1382 } else {
1383 Ok::<Option<String>, Error>(None)
1384 }
1385 }))
1386 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1387
1388 let mut filtered = Vec::new();
1389 while let Some(result) = stream.next().await {
1390 if let Some(table_name) = result? {
1391 filtered.push(table_name);
1392 }
1393 }
1394 Ok(filtered)
1395 }
1396
1397 fn ipc_reader_from_request_data(
1398 request_data: &Bytes,
1399 operation: &str,
1400 ) -> Result<(
1401 Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1402 usize,
1403 )> {
1404 if request_data.is_empty() {
1405 return Err(NamespaceError::InvalidInput {
1406 message: format!(
1407 "Request data (Arrow IPC stream) is required for {}",
1408 operation
1409 ),
1410 }
1411 .into());
1412 }
1413
1414 let cursor = Cursor::new(request_data.as_ref());
1415 let stream_reader =
1416 StreamReader::try_new(cursor, None).map_err(|e| NamespaceError::InvalidInput {
1417 message: format!("Invalid Arrow IPC stream: {}", e),
1418 })?;
1419 let arrow_schema = stream_reader.schema();
1420
1421 let mut num_rows = 0usize;
1422 let mut batches = Vec::new();
1423 for batch_result in stream_reader {
1424 let batch = batch_result.map_err(|e| NamespaceError::Internal {
1425 message: format!("Failed to read batch from IPC stream: {}", e),
1426 })?;
1427 num_rows += batch.num_rows();
1428 batches.push(batch);
1429 }
1430
1431 let reader: Box<dyn arrow::record_batch::RecordBatchReader + Send> = if batches.is_empty() {
1432 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
1433 Box::new(RecordBatchIterator::new(vec![Ok(batch)], arrow_schema))
1434 } else {
1435 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
1436 Box::new(RecordBatchIterator::new(batch_results, arrow_schema))
1437 };
1438
1439 Ok((reader, num_rows))
1440 }
1441
1442 async fn table_uri_has_actual_manifests(&self, table_uri: &str) -> Result<bool> {
1443 let table_path = self.object_store_path_from_uri(table_uri)?;
1444 manifest::ManifestNamespace::path_has_actual_manifests(&self.object_store, &table_path)
1445 .await
1446 }
1447
1448 fn object_store_path_from_uri(&self, uri: &str) -> Result<Path> {
1449 let registry = self
1450 .session
1451 .as_ref()
1452 .map(|session| session.store_registry())
1453 .unwrap_or_else(|| Arc::new(ObjectStoreRegistry::default()));
1454 ObjectStore::extract_path_from_uri(registry, uri)
1455 }
1456
1457 fn normalized_branch(branch: Option<&str>) -> Result<Option<&str>> {
1462 match branch.filter(|b| !b.is_empty() && *b != "main") {
1463 Some(branch) => {
1464 check_valid_branch(branch).map_err(|e| {
1465 lance_core::Error::from(NamespaceError::InvalidInput {
1466 message: format!("invalid branch name '{}': {}", branch, e),
1467 })
1468 })?;
1469 Ok(Some(branch))
1470 }
1471 None => Ok(None),
1472 }
1473 }
1474
1475 async fn open_validated_branch(&self, table_uri: &str, branch: &str) -> Result<Dataset> {
1476 let dataset = self
1477 .configured_builder(table_uri)
1478 .with_branch(branch, None)
1479 .load()
1480 .await
1481 .map_err(|e| {
1482 let message = format!(
1483 "branch '{}' not found for table at '{}': {}",
1484 branch, table_uri, e
1485 );
1486 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1487 })?;
1488 dataset.branches().get(branch).await.map_err(|e| {
1489 Self::map_open_error(
1490 e,
1491 NamespaceError::TableNotFound {
1492 message: format!("branch '{}' not found for table at '{}'", branch, table_uri),
1493 },
1494 )
1495 })?;
1496 Ok(dataset)
1497 }
1498
1499 async fn resolve_branch_location(&self, table_uri: &str, branch: &str) -> Result<String> {
1500 Ok(self
1501 .open_validated_branch(table_uri, branch)
1502 .await?
1503 .branch_location()
1504 .uri)
1505 }
1506
1507 async fn resolve_branch_for_commit(
1517 &self,
1518 table_uri: &str,
1519 branch: &str,
1520 ) -> Result<(String, Path, Option<u64>)> {
1521 let main = self
1522 .configured_builder(table_uri)
1523 .load()
1524 .await
1525 .map_err(|e| {
1526 let message = format!("table at '{}' not found: {}", table_uri, e);
1527 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1528 })?;
1529 let branch_location = main.branch_location().find_branch(Some(branch))?;
1530 match main.branches().get(branch).await {
1531 Ok(contents) => Ok((
1532 branch_location.uri,
1533 branch_location.path,
1534 Some(contents.parent_version),
1535 )),
1536 Err(lance_core::Error::RefNotFound { .. }) => {
1537 if self
1538 .branch_has_committed_versions(&branch_location.path)
1539 .await?
1540 {
1541 return Err(NamespaceError::TableNotFound {
1542 message: format!(
1543 "branch '{}' not found for table at '{}'",
1544 branch, table_uri
1545 ),
1546 }
1547 .into());
1548 }
1549 Ok((branch_location.uri, branch_location.path, None))
1550 }
1551 Err(e) => Err(e),
1552 }
1553 }
1554
1555 async fn branch_has_committed_versions(&self, branch_path: &Path) -> Result<bool> {
1556 Ok(!self
1557 .list_versions_under(branch_path, false, Some(1))
1558 .await?
1559 .is_empty())
1560 }
1561
1562 fn validate_dir_only_properties(
1563 properties: Option<&HashMap<String, String>>,
1564 operation: &str,
1565 ) -> Result<()> {
1566 if properties.is_some_and(|properties| !properties.is_empty()) {
1570 return Err(NamespaceError::Unsupported {
1571 message: format!(
1572 "{} with non-empty table properties requires manifest_enabled=true",
1573 operation
1574 ),
1575 }
1576 .into());
1577 }
1578 Ok(())
1579 }
1580
1581 async fn write_reader_to_table(
1582 &self,
1583 table_uri: &str,
1584 reader: Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1585 mode: WriteMode,
1586 extra_storage_options: Option<HashMap<String, String>>,
1587 ) -> Result<Dataset> {
1588 let mut merged_storage_options = self.storage_options.clone().unwrap_or_default();
1591 if let Some(extra_storage_options) = extra_storage_options {
1592 merged_storage_options.extend(extra_storage_options);
1593 }
1594 let store_params = (!merged_storage_options.is_empty()).then(|| ObjectStoreParams {
1595 storage_options_accessor: Some(Arc::new(
1596 lance_io::object_store::StorageOptionsAccessor::with_static_options(
1597 merged_storage_options,
1598 ),
1599 )),
1600 ..Default::default()
1601 });
1602
1603 let write_params = WriteParams {
1604 mode,
1605 store_params,
1606 session: self.session.clone(),
1607 ..Default::default()
1608 };
1609
1610 let dataset = Dataset::write(reader, table_uri, Some(write_params))
1611 .await
1612 .map_err(|e| NamespaceError::Internal {
1613 message: format!("Failed to write table at '{}': {}", table_uri, e),
1614 })?;
1615
1616 Ok(dataset)
1617 }
1618
1619 fn manifest_version_from_filename(filename: &str) -> Option<u64> {
1623 ManifestNamingScheme::detect_scheme(filename)?.parse_version(filename)
1624 }
1625
1626 fn create_table_version_response(
1628 version: u64,
1629 final_path: &Path,
1630 final_meta: &ObjectMeta,
1631 ) -> CreateTableVersionResponse {
1632 CreateTableVersionResponse {
1633 transaction_id: None,
1634 version: Some(Box::new(TableVersion {
1635 version: version as i64,
1636 manifest_path: final_path.to_string(),
1637 manifest_size: Some(final_meta.size as i64),
1638 e_tag: final_meta.e_tag.clone(),
1639 timestamp_millis: None,
1640 metadata: None,
1641 })),
1642 }
1643 }
1644
1645 async fn staging_matches_final_manifest(
1653 &self,
1654 staging_path: &Path,
1655 final_path: &Path,
1656 final_meta: &ObjectMeta,
1657 request_manifest_size: Option<i64>,
1658 ) -> Result<bool> {
1659 if let Some(size) = request_manifest_size
1660 && size != final_meta.size as i64
1661 {
1662 return Ok(false);
1663 }
1664
1665 let staging_bytes = match self.object_store.inner.get(staging_path).await {
1666 Ok(r) => r.bytes().await.map_err(|e| {
1667 lance_core::Error::from(NamespaceError::Internal {
1668 message: format!(
1669 "Failed to read staging manifest at '{}': {}",
1670 staging_path, e
1671 ),
1672 })
1673 })?,
1674 Err(ObjectStoreError::NotFound { .. }) => return Ok(false),
1675 Err(e) => {
1676 return Err(lance_core::Error::from(NamespaceError::Internal {
1677 message: format!(
1678 "Failed to read staging manifest at '{}': {}",
1679 staging_path, e
1680 ),
1681 }));
1682 }
1683 };
1684
1685 let final_bytes = self
1686 .object_store
1687 .inner
1688 .get(final_path)
1689 .await
1690 .map_err(|e| {
1691 lance_core::Error::from(NamespaceError::Internal {
1692 message: format!(
1693 "Failed to read existing version manifest at '{}': {}",
1694 final_path, e
1695 ),
1696 })
1697 })?
1698 .bytes()
1699 .await
1700 .map_err(|e| {
1701 lance_core::Error::from(NamespaceError::Internal {
1702 message: format!(
1703 "Failed to read existing version manifest bytes at '{}': {}",
1704 final_path, e
1705 ),
1706 })
1707 })?;
1708
1709 Ok(staging_bytes.as_ref() == final_bytes.as_ref())
1710 }
1711
1712 async fn resolve_existing_table_version(
1714 &self,
1715 args: ExistingTableVersionResolve<'_>,
1716 ) -> Result<CreateTableVersionResponse> {
1717 if self
1718 .staging_matches_final_manifest(
1719 args.staging_path,
1720 args.final_path,
1721 args.final_meta,
1722 args.request_manifest_size,
1723 )
1724 .await?
1725 {
1726 if let Err(e) = self.object_store.inner.delete(args.staging_path).await {
1728 log::warn!(
1729 "Failed to delete staging manifest at '{}': {:?}",
1730 args.staging_path,
1731 e
1732 );
1733 }
1734 return Ok(Self::create_table_version_response(
1735 args.version,
1736 args.final_path,
1737 args.final_meta,
1738 ));
1739 }
1740
1741 Err(lance_core::Error::from(
1742 NamespaceError::ConcurrentModification {
1743 message: format!(
1744 "Version {} already exists for table at '{}' with different content",
1745 args.version, args.table_uri
1746 ),
1747 },
1748 ))
1749 }
1750
1751 async fn enforce_create_table_version_cas(
1760 &self,
1761 table_path: &Path,
1762 version: u64,
1763 table_uri: &str,
1764 is_branch: bool,
1765 branch_parent_version: Option<u64>,
1766 ) -> Result<()> {
1767 let latest = self.list_versions_under(table_path, true, Some(1)).await?;
1768 let expected = match latest.first() {
1769 Some(v) => (v.version as u64).checked_add(1).ok_or_else(|| {
1770 lance_core::Error::from(NamespaceError::ConcurrentModification {
1771 message: format!(
1772 "Version overflow computing next version for table at '{}': \
1773 latest version {} cannot advance",
1774 table_uri, v.version
1775 ),
1776 })
1777 })?,
1778 None => {
1779 if is_branch {
1780 match branch_parent_version {
1783 Some(parent_version) => parent_version,
1784 None => version,
1785 }
1786 } else {
1787 1
1788 }
1789 }
1790 };
1791 if version != expected {
1792 let latest_display = latest
1793 .first()
1794 .map(|v| v.version.to_string())
1795 .unwrap_or_else(|| "none".to_string());
1796 return Err(lance_core::Error::from(
1797 NamespaceError::ConcurrentModification {
1798 message: format!(
1799 "Version CAS failed for table at '{}': requested {}, expected {} (latest {})",
1800 table_uri, version, expected, latest_display
1801 ),
1802 },
1803 ));
1804 }
1805 Ok(())
1806 }
1807
1808 async fn materialize_version_manifest_create(
1810 &self,
1811 staging_path: &Path,
1812 final_path: &Path,
1813 staging_manifest_path: &str,
1814 ) -> std::result::Result<(), ObjectStoreError> {
1815 match self
1816 .object_store
1817 .inner
1818 .copy_if_not_exists(staging_path, final_path)
1819 .await
1820 {
1821 Ok(()) => Ok(()),
1822 Err(ObjectStoreError::NotImplemented { .. })
1823 | Err(ObjectStoreError::NotSupported { .. }) => {
1824 let manifest_data = self
1825 .object_store
1826 .inner
1827 .get(staging_path)
1828 .await?
1829 .bytes()
1830 .await
1831 .map_err(|e| ObjectStoreError::Generic {
1832 store: "DirectoryNamespace",
1833 source: Box::new(std::io::Error::other(format!(
1834 "Failed to read staging manifest bytes at '{}': {}",
1835 staging_manifest_path, e
1836 ))),
1837 })?;
1838 self.object_store
1839 .inner
1840 .put_opts(
1841 final_path,
1842 manifest_data.into(),
1843 PutOptions {
1844 mode: PutMode::Create,
1845 ..Default::default()
1846 },
1847 )
1848 .await
1849 .map(|_| ())
1850 }
1851 Err(e) => Err(e),
1852 }
1853 }
1854
1855 async fn list_table_versions_from_storage(
1856 &self,
1857 table_uri: &str,
1858 descending: bool,
1859 limit: Option<i32>,
1860 ) -> Result<Vec<TableVersion>> {
1861 let table_path = self.object_store_path_from_uri(table_uri)?;
1862 self.list_versions_under(&table_path, descending, limit)
1863 .await
1864 }
1865
1866 async fn list_versions_under(
1870 &self,
1871 table_path: &Path,
1872 descending: bool,
1873 limit: Option<i32>,
1874 ) -> Result<Vec<TableVersion>> {
1875 let versions_dir = table_path.clone().join(VERSIONS_DIR);
1876 let manifest_metas: Vec<_> = self
1877 .object_store
1878 .read_dir_all(&versions_dir, None)
1879 .try_collect()
1880 .await
1881 .map_err(|e| {
1882 lance_core::Error::from(NamespaceError::Internal {
1883 message: format!(
1884 "Failed to list manifest files under '{}': {}",
1885 versions_dir, e
1886 ),
1887 })
1888 })?;
1889
1890 let is_v2_naming = manifest_metas
1891 .first()
1892 .is_some_and(|meta| meta.location.filename().is_some_and(|f| f.len() == 29));
1893
1894 let mut table_versions: Vec<TableVersion> = manifest_metas
1895 .into_iter()
1896 .filter_map(|meta| {
1897 let filename = meta.location.filename()?;
1898 let actual_version = Self::manifest_version_from_filename(filename)?;
1899
1900 Some(TableVersion {
1901 version: actual_version as i64,
1902 manifest_path: meta.location.to_string(),
1903 manifest_size: Some(meta.size as i64),
1904 e_tag: meta.e_tag,
1905 timestamp_millis: Some(meta.last_modified.timestamp_millis()),
1906 metadata: None,
1907 })
1908 })
1909 .collect();
1910
1911 let list_is_ordered = self.object_store.list_is_lexically_ordered;
1912
1913 let needs_sort = if list_is_ordered {
1914 if is_v2_naming {
1915 !descending
1916 } else {
1917 descending
1918 }
1919 } else {
1920 true
1921 };
1922
1923 if needs_sort {
1924 if descending {
1925 table_versions.sort_by_key(|v| std::cmp::Reverse(v.version));
1926 } else {
1927 table_versions.sort_by_key(|v| v.version);
1928 }
1929 }
1930
1931 if let Some(limit) = limit {
1932 table_versions.truncate(limit as usize);
1933 }
1934
1935 Ok(table_versions)
1936 }
1937
1938 async fn describe_table_impl(
1942 &self,
1943 request: DescribeTableRequest,
1944 ) -> Result<DescribeTableResponse> {
1945 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
1946 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
1947 let skip_manifest_for_root = self.dir_listing_enabled
1948 && is_root_level
1949 && !self.dir_listing_to_manifest_migration_enabled;
1950 if let Some(manifest_ns) = self.manifest_ns_for_read()
1951 && !skip_manifest_for_root
1952 {
1953 match manifest_ns.describe_table(request.clone()).await {
1954 Ok(mut response) => {
1955 if let Some(ref table_uri) = response.table_uri {
1956 let vend = request.vend_credentials.unwrap_or(true);
1958 let identity = request.identity.as_deref();
1959 response.storage_options = self
1960 .get_storage_options_for_table(table_uri, vend, identity)
1961 .await?;
1962 }
1963 if self.table_version_tracking_enabled {
1965 response.managed_versioning = Some(true);
1966 }
1967 return Ok(response);
1968 }
1969 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
1970 return Err(e);
1973 }
1974 Err(e) if self.dir_listing_enabled && is_root_level => {
1975 if !Self::is_manifest_table_absent_error(&e) {
1979 return Err(Self::classify_storage_error(e));
1980 }
1981 }
1982 Err(e) => return Err(e),
1983 }
1984 }
1985 if is_child_table {
1986 return Err(self.child_namespace_requires_manifest_error());
1987 }
1988
1989 let table_name = Self::table_name_from_id(&request.id)?;
1990 let table_id = Self::format_table_id_from_request(&request.id);
1991 if !self.dir_listing_enabled {
1992 return Err(NamespaceError::TableNotFound { message: table_id }.into());
1993 }
1994
1995 let table_uri = self.table_full_uri(&table_name);
1996
1997 let status = self.check_table_status(&table_name).await?;
1999
2000 if !status.exists {
2001 return Err(NamespaceError::TableNotFound {
2002 message: table_id.clone(),
2003 }
2004 .into());
2005 }
2006
2007 if status.is_deregistered {
2008 return Err(NamespaceError::TableNotFound {
2009 message: format!("Table is deregistered: {}", table_id),
2010 }
2011 .into());
2012 }
2013
2014 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
2015 let should_check_declared =
2016 load_detailed_metadata || request.check_declared.unwrap_or(false);
2017 let vend_credentials = request.vend_credentials.unwrap_or(true);
2019 let identity = request.identity.as_deref();
2020 let is_only_declared = if should_check_declared {
2021 if status.has_reserved_file {
2022 Some(!self.table_has_actual_manifests(&table_name).await?)
2023 } else {
2024 Some(false)
2025 }
2026 } else {
2027 None
2028 };
2029
2030 if !load_detailed_metadata {
2031 let storage_options = self
2032 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2033 .await?;
2034 return Ok(DescribeTableResponse {
2035 table: Some(table_name),
2036 namespace: request.id.as_ref().map(|id| {
2037 if id.len() > 1 {
2038 id[..id.len() - 1].to_vec()
2039 } else {
2040 vec![]
2041 }
2042 }),
2043 location: Some(table_uri.clone()),
2044 table_uri: Some(table_uri),
2045 storage_options,
2046 is_only_declared,
2047 managed_versioning: if self.table_version_tracking_enabled {
2048 Some(true)
2049 } else {
2050 None
2051 },
2052 ..Default::default()
2053 });
2054 }
2055
2056 if is_only_declared == Some(true) {
2057 let storage_options = self
2058 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2059 .await?;
2060 return Ok(DescribeTableResponse {
2061 table: Some(table_name),
2062 namespace: request.id.as_ref().map(|id| {
2063 if id.len() > 1 {
2064 id[..id.len() - 1].to_vec()
2065 } else {
2066 vec![]
2067 }
2068 }),
2069 location: Some(table_uri.clone()),
2070 table_uri: Some(table_uri),
2071 storage_options,
2072 is_only_declared,
2073 managed_versioning: if self.table_version_tracking_enabled {
2074 Some(true)
2075 } else {
2076 None
2077 },
2078 ..Default::default()
2079 });
2080 }
2081
2082 let mut builder = DatasetBuilder::from_uri(&table_uri);
2085 if let Some(opts) = &self.storage_options {
2086 builder = builder.with_storage_options(opts.clone());
2087 }
2088 if let Some(sess) = &self.session {
2089 builder = builder.with_session(sess.clone());
2090 }
2091 match builder.load().await {
2092 Ok(mut dataset) => {
2093 if let Some(requested_version) = request.version {
2095 dataset = dataset
2096 .checkout_version(requested_version as u64)
2097 .await
2098 .map_err(|e| {
2099 let message = format!(
2100 "Version {} not found for table '{}': {}",
2101 requested_version, table_name, e
2102 );
2103 Self::map_open_error(
2104 e,
2105 NamespaceError::TableVersionNotFound { message },
2106 )
2107 })?;
2108 }
2109
2110 let version_info = dataset.version();
2111 let lance_schema = dataset.schema();
2112 let arrow_schema: arrow_schema::Schema = lance_schema.into();
2113 let json_schema = arrow_schema_to_json(&arrow_schema)?;
2114 let storage_options = self
2115 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2116 .await?;
2117
2118 let metadata: std::collections::HashMap<String, String> =
2120 version_info.metadata.into_iter().collect();
2121
2122 Ok(DescribeTableResponse {
2123 table: Some(table_name),
2124 namespace: request.id.as_ref().map(|id| {
2125 if id.len() > 1 {
2126 id[..id.len() - 1].to_vec()
2127 } else {
2128 vec![]
2129 }
2130 }),
2131 version: Some(version_info.version as i64),
2132 location: Some(table_uri.clone()),
2133 table_uri: Some(table_uri),
2134 schema: Some(Box::new(json_schema)),
2135 storage_options,
2136 metadata: Some(metadata),
2137 is_only_declared,
2138 managed_versioning: if self.table_version_tracking_enabled {
2139 Some(true)
2140 } else {
2141 None
2142 },
2143 ..Default::default()
2144 })
2145 }
2146 Err(err) => {
2147 if manifest::ManifestNamespace::is_not_found_load_error(&err)
2148 && is_only_declared == Some(true)
2149 {
2150 let storage_options = self
2151 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2152 .await?;
2153 Ok(DescribeTableResponse {
2154 table: Some(table_name),
2155 namespace: request.id.as_ref().map(|id| {
2156 if id.len() > 1 {
2157 id[..id.len() - 1].to_vec()
2158 } else {
2159 vec![]
2160 }
2161 }),
2162 location: Some(table_uri.clone()),
2163 table_uri: Some(table_uri),
2164 storage_options,
2165 is_only_declared,
2166 managed_versioning: if self.table_version_tracking_enabled {
2167 Some(true)
2168 } else {
2169 None
2170 },
2171 ..Default::default()
2172 })
2173 } else {
2174 Err(NamespaceError::Internal {
2175 message: format!(
2176 "Table directory exists but cannot load dataset {}: {:?}",
2177 table_name, err
2178 ),
2179 }
2180 .into())
2181 }
2182 }
2183 }
2184 }
2185
2186 fn configured_builder(&self, table_uri: &str) -> DatasetBuilder {
2189 let mut builder = DatasetBuilder::from_uri(table_uri);
2190 if let Some(opts) = &self.storage_options {
2191 builder = builder.with_storage_options(opts.clone());
2192 }
2193 if let Some(sess) = &self.session {
2194 builder = builder.with_session(sess.clone());
2195 }
2196 builder
2197 }
2198
2199 async fn load_dataset(
2200 &self,
2201 table_uri: &str,
2202 version: Option<i64>,
2203 operation: &str,
2204 ) -> Result<Dataset> {
2205 if let Some(version) = version
2206 && version < 0
2207 {
2208 return Err(NamespaceError::InvalidInput {
2209 message: format!(
2210 "Table version for {} must be non-negative, got {}",
2211 operation, version
2212 ),
2213 }
2214 .into());
2215 }
2216
2217 let builder = self.configured_builder(table_uri);
2218
2219 let dataset = builder.load().await.map_err(|e| {
2220 let message = format!(
2221 "Failed to open table at '{}' for {}: {}",
2222 table_uri, operation, e
2223 );
2224 Self::map_open_error(e, NamespaceError::TableNotFound { message })
2225 })?;
2226
2227 if let Some(version) = version {
2228 return dataset.checkout_version(version as u64).await.map_err(|e| {
2229 let message = format!(
2230 "Failed to checkout version {} for table at '{}' during {}: {}",
2231 version, table_uri, operation, e
2232 );
2233 Self::map_open_error(e, NamespaceError::TableVersionNotFound { message })
2234 });
2235 }
2236
2237 Ok(dataset)
2238 }
2239
2240 fn parse_index_type(index_type: &str) -> Result<IndexType> {
2241 match index_type.trim().to_ascii_uppercase().as_str() {
2242 "SCALAR" | "BTREE" => Ok(IndexType::BTree),
2243 "BITMAP" => Ok(IndexType::Bitmap),
2244 "LABEL_LIST" | "LABELLIST" => Ok(IndexType::LabelList),
2245 "INVERTED" | "FTS" => Ok(IndexType::Inverted),
2246 "NGRAM" => Ok(IndexType::NGram),
2247 "ZONEMAP" | "ZONE_MAP" => Ok(IndexType::ZoneMap),
2248 "BLOOMFILTER" | "BLOOM_FILTER" => Ok(IndexType::BloomFilter),
2249 "RTREE" | "R_TREE" => Ok(IndexType::RTree),
2250 "VECTOR" | "IVF_PQ" => Ok(IndexType::IvfPq),
2251 "IVF_FLAT" => Ok(IndexType::IvfFlat),
2252 "IVF_SQ" => Ok(IndexType::IvfSq),
2253 "IVF_RQ" => Ok(IndexType::IvfRq),
2254 "IVF_HNSW_FLAT" => Ok(IndexType::IvfHnswFlat),
2255 "IVF_HNSW_SQ" => Ok(IndexType::IvfHnswSq),
2256 "IVF_HNSW_PQ" => Ok(IndexType::IvfHnswPq),
2257 other => Err(NamespaceError::InvalidInput {
2258 message: format!("Unsupported index_type '{}'", other),
2259 }
2260 .into()),
2261 }
2262 }
2263
2264 fn parse_metric_type(distance_type: Option<&str>) -> Result<MetricType> {
2265 let distance_type = distance_type.unwrap_or("l2");
2266 MetricType::try_from(distance_type).map_err(|e| {
2267 lance_core::Error::from(NamespaceError::InvalidInput {
2268 message: format!(
2269 "Unsupported distance_type '{}' for vector index: {}",
2270 distance_type, e
2271 ),
2272 })
2273 })
2274 }
2275
2276 fn build_index_params(request: &CreateTableIndexRequest) -> Result<DirectoryIndexParams> {
2277 let index_type = Self::parse_index_type(&request.index_type)?;
2278 Ok(match index_type {
2279 IndexType::BTree => DirectoryIndexParams::Scalar {
2280 index_type,
2281 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
2282 },
2283 IndexType::Bitmap => DirectoryIndexParams::Scalar {
2284 index_type,
2285 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
2286 },
2287 IndexType::LabelList => DirectoryIndexParams::Scalar {
2288 index_type,
2289 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
2290 },
2291 IndexType::NGram => DirectoryIndexParams::Scalar {
2292 index_type,
2293 params: ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
2294 },
2295 IndexType::ZoneMap => DirectoryIndexParams::Scalar {
2296 index_type,
2297 params: ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
2298 },
2299 IndexType::BloomFilter => DirectoryIndexParams::Scalar {
2300 index_type,
2301 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter),
2302 },
2303 IndexType::RTree => DirectoryIndexParams::Scalar {
2304 index_type,
2305 params: ScalarIndexParams::for_builtin(BuiltinIndexType::RTree),
2306 },
2307 IndexType::Inverted => {
2308 let mut params = InvertedIndexParams::default();
2309 if let Some(with_position) = request.with_position {
2310 params = params.with_position(with_position);
2311 }
2312 if let Some(base_tokenizer) = &request.base_tokenizer {
2313 params = params.base_tokenizer(base_tokenizer.clone());
2314 }
2315 if let Some(language) = &request.language {
2316 params = params.language(language)?;
2317 }
2318 if let Some(max_token_length) = request.max_token_length {
2319 if max_token_length < 0 {
2320 return Err(NamespaceError::InvalidInput {
2321 message: format!(
2322 "FTS max_token_length must be non-negative, got {}",
2323 max_token_length
2324 ),
2325 }
2326 .into());
2327 }
2328 params = params.max_token_length(Some(max_token_length as usize));
2329 }
2330 if let Some(lower_case) = request.lower_case {
2331 params = params.lower_case(lower_case);
2332 }
2333 if let Some(stem) = request.stem {
2334 params = params.stem(stem);
2335 }
2336 if let Some(remove_stop_words) = request.remove_stop_words {
2337 params = params.remove_stop_words(remove_stop_words);
2338 }
2339 if let Some(ascii_folding) = request.ascii_folding {
2340 params = params.ascii_folding(ascii_folding);
2341 }
2342 DirectoryIndexParams::Inverted(params)
2343 }
2344 IndexType::IvfFlat => DirectoryIndexParams::Vector {
2345 index_type,
2346 params: VectorIndexParams::with_ivf_flat_params(
2347 Self::parse_metric_type(request.distance_type.as_deref())?,
2348 IvfBuildParams::default(),
2349 ),
2350 },
2351 IndexType::IvfPq => DirectoryIndexParams::Vector {
2352 index_type,
2353 params: VectorIndexParams::with_ivf_pq_params(
2354 Self::parse_metric_type(request.distance_type.as_deref())?,
2355 IvfBuildParams::default(),
2356 PQBuildParams::default(),
2357 ),
2358 },
2359 IndexType::IvfSq => DirectoryIndexParams::Vector {
2360 index_type,
2361 params: VectorIndexParams::with_ivf_sq_params(
2362 Self::parse_metric_type(request.distance_type.as_deref())?,
2363 IvfBuildParams::default(),
2364 SQBuildParams::default(),
2365 ),
2366 },
2367 IndexType::IvfRq => DirectoryIndexParams::Vector {
2368 index_type,
2369 params: VectorIndexParams::with_ivf_rq_params(
2370 Self::parse_metric_type(request.distance_type.as_deref())?,
2371 IvfBuildParams::default(),
2372 RQBuildParams::default(),
2373 ),
2374 },
2375 IndexType::IvfHnswFlat => DirectoryIndexParams::Vector {
2376 index_type,
2377 params: VectorIndexParams::ivf_hnsw(
2378 Self::parse_metric_type(request.distance_type.as_deref())?,
2379 IvfBuildParams::default(),
2380 HnswBuildParams::default(),
2381 ),
2382 },
2383 IndexType::IvfHnswSq => DirectoryIndexParams::Vector {
2384 index_type,
2385 params: VectorIndexParams::with_ivf_hnsw_sq_params(
2386 Self::parse_metric_type(request.distance_type.as_deref())?,
2387 IvfBuildParams::default(),
2388 HnswBuildParams::default(),
2389 SQBuildParams::default(),
2390 ),
2391 },
2392 IndexType::IvfHnswPq => DirectoryIndexParams::Vector {
2393 index_type,
2394 params: VectorIndexParams::with_ivf_hnsw_pq_params(
2395 Self::parse_metric_type(request.distance_type.as_deref())?,
2396 IvfBuildParams::default(),
2397 HnswBuildParams::default(),
2398 PQBuildParams::default(),
2399 ),
2400 },
2401 other => {
2402 return Err(NamespaceError::InvalidInput {
2403 message: format!("Unsupported index type for namespace API: {}", other),
2404 }
2405 .into());
2406 }
2407 })
2408 }
2409
2410 fn paginate_indices(
2411 indices: &mut Vec<IndexContent>,
2412 page_token: Option<String>,
2413 limit: Option<i32>,
2414 ) -> Option<String> {
2415 indices.sort_by(|a, b| a.index_name.cmp(&b.index_name));
2416
2417 if let Some(start_after) = page_token {
2418 if let Some(index) = indices
2419 .iter()
2420 .position(|index| index.index_name.as_str() > start_after.as_str())
2421 {
2422 indices.drain(0..index);
2423 } else {
2424 indices.clear();
2425 }
2426 }
2427
2428 let mut next_page_token = None;
2429 if let Some(limit) = limit
2430 && limit >= 0
2431 {
2432 let limit = limit as usize;
2433 if limit > 0 && indices.len() > limit {
2434 next_page_token = Some(indices[limit - 1].index_name.clone());
2435 }
2436 indices.truncate(limit);
2437 }
2438 if indices.is_empty() {
2439 None
2440 } else {
2441 next_page_token
2442 }
2443 }
2444
2445 fn transaction_operation_name(transaction: &Transaction) -> String {
2446 match &transaction.operation {
2447 Operation::CreateIndex {
2448 new_indices,
2449 removed_indices,
2450 } if new_indices.is_empty() && !removed_indices.is_empty() => "DropIndex".to_string(),
2451 _ => transaction.operation.to_string(),
2452 }
2453 }
2454
2455 fn transaction_response(
2456 version: u64,
2457 transaction: &Transaction,
2458 alteration: Option<TransactionAlteration>,
2459 ) -> DescribeTransactionResponse {
2460 let mut properties = transaction
2461 .transaction_properties
2462 .as_ref()
2463 .map(|properties| (**properties).clone())
2464 .unwrap_or_default();
2465
2466 let mut effective_status = "SUCCEEDED".to_string();
2469 if let Some(alteration) = alteration {
2470 for key in &alteration.removed_properties {
2471 properties.remove(key);
2472 }
2473 for (key, value) in alteration.properties {
2474 properties.insert(key, value);
2475 }
2476 if let Some(status) = alteration.status {
2477 effective_status = status;
2478 }
2479 }
2480
2481 properties.insert("uuid".to_string(), transaction.uuid.clone());
2482 properties.insert("version".to_string(), version.to_string());
2483 properties.insert(
2484 "read_version".to_string(),
2485 transaction.read_version.to_string(),
2486 );
2487 properties.insert(
2488 "operation".to_string(),
2489 Self::transaction_operation_name(transaction),
2490 );
2491 if let Some(tag) = &transaction.tag {
2492 properties.insert("tag".to_string(), tag.clone());
2493 }
2494
2495 DescribeTransactionResponse {
2496 status: effective_status,
2497 properties: Some(properties),
2498 }
2499 }
2500
2501 fn describe_table_index_stats_response(
2502 stats: &serde_json::Value,
2503 ) -> DescribeTableIndexStatsResponse {
2504 let get_i64 = |key: &str| {
2505 stats.get(key).and_then(|value| {
2506 value
2507 .as_i64()
2508 .or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
2509 })
2510 };
2511
2512 DescribeTableIndexStatsResponse {
2513 distance_type: stats
2514 .get("distance_type")
2515 .and_then(|value| value.as_str())
2516 .map(str::to_string),
2517 index_type: stats
2518 .get("index_type")
2519 .and_then(|value| value.as_str())
2520 .map(str::to_string),
2521 num_indexed_rows: get_i64("num_indexed_rows"),
2522 num_unindexed_rows: get_i64("num_unindexed_rows"),
2523 num_indices: get_i64("num_indices").and_then(|value| i32::try_from(value).ok()),
2524 }
2525 }
2526
2527 async fn find_transaction(&self, dataset: &Dataset, id: &str) -> Result<(u64, Transaction)> {
2532 if let Ok(version) = id.parse::<u64>() {
2533 let transaction = dataset
2534 .read_transaction_by_version(version)
2535 .await
2536 .map_err(|e| {
2537 lance_core::Error::from(NamespaceError::TransactionNotFound {
2538 message: format!(
2539 "Failed to read transaction for version {}: {}",
2540 version, e
2541 ),
2542 })
2543 })?
2544 .ok_or_else(|| {
2545 lance_core::Error::from(NamespaceError::TransactionNotFound {
2546 message: format!("version {}", version),
2547 })
2548 })?;
2549 return Ok((version, transaction));
2550 }
2551
2552 let versions = dataset.versions().await.map_err(|e| {
2553 lance_core::Error::from(NamespaceError::Internal {
2554 message: format!(
2555 "Failed to list table versions while resolving transaction '{}': {}",
2556 id, e
2557 ),
2558 })
2559 })?;
2560
2561 for version in versions.into_iter().rev() {
2562 if let Some(transaction) = dataset
2563 .read_transaction_by_version(version.version)
2564 .await
2565 .map_err(|e| {
2566 lance_core::Error::from(NamespaceError::Internal {
2567 message: format!(
2568 "Failed to read transaction for version {} while resolving '{}': {}",
2569 version.version, id, e
2570 ),
2571 })
2572 })?
2573 && transaction.uuid == id
2574 {
2575 return Ok((version.version, transaction));
2576 }
2577 }
2578
2579 Err(NamespaceError::TransactionNotFound {
2580 message: id.to_string(),
2581 }
2582 .into())
2583 }
2584
2585 const TRANSACTION_ALTERATIONS_DIR: &'static str = "_alter_transactions";
2589
2590 fn transaction_alteration_path(&self, table_uri: &str, txn_uuid: &str) -> Result<Path> {
2591 let table_path = self.object_store_path_from_uri(table_uri)?;
2592 Ok(table_path
2593 .join(Self::TRANSACTION_ALTERATIONS_DIR)
2594 .join(format!("{}.json", txn_uuid).as_str()))
2595 }
2596
2597 async fn load_transaction_alteration(
2598 &self,
2599 table_uri: &str,
2600 txn_uuid: &str,
2601 ) -> Result<Option<TransactionAlteration>> {
2602 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2603 match self.object_store.inner.get(&path).await {
2604 Ok(get_result) => {
2605 let bytes = get_result.bytes().await.map_err(|e| {
2606 lance_core::Error::from(NamespaceError::Internal {
2607 message: format!(
2608 "Failed to read alter_transaction sidecar for '{}': {}",
2609 txn_uuid, e
2610 ),
2611 })
2612 })?;
2613 let alteration = TransactionAlteration::from_json_slice(&bytes).map_err(|e| {
2614 lance_core::Error::from(NamespaceError::Internal {
2615 message: format!(
2616 "Failed to parse alter_transaction sidecar for '{}': {}",
2617 txn_uuid, e
2618 ),
2619 })
2620 })?;
2621 Ok(Some(alteration))
2622 }
2623 Err(ObjectStoreError::NotFound { .. }) => Ok(None),
2624 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2625 message: format!(
2626 "Failed to load alter_transaction sidecar for '{}': {}",
2627 txn_uuid, e
2628 ),
2629 })),
2630 }
2631 }
2632
2633 async fn save_transaction_alteration(
2634 &self,
2635 table_uri: &str,
2636 txn_uuid: &str,
2637 alteration: &TransactionAlteration,
2638 ) -> Result<()> {
2639 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2640 let bytes = alteration.to_json_bytes().map_err(|e| {
2641 lance_core::Error::from(NamespaceError::Internal {
2642 message: format!(
2643 "Failed to serialize alter_transaction sidecar for '{}': {}",
2644 txn_uuid, e
2645 ),
2646 })
2647 })?;
2648 self.object_store
2649 .inner
2650 .put(&path, bytes.into())
2651 .await
2652 .map_err(|e| {
2653 lance_core::Error::from(NamespaceError::Internal {
2654 message: format!(
2655 "Failed to persist alter_transaction sidecar for '{}': {}",
2656 txn_uuid, e
2657 ),
2658 })
2659 })?;
2660 Ok(())
2661 }
2662
2663 fn table_full_uri(&self, table_name: &str) -> String {
2664 format!("{}/{}.lance", self.root, table_name)
2665 }
2666
2667 fn table_path(&self, table_name: &str) -> Path {
2669 self.base_path
2670 .clone()
2671 .join(format!("{}.lance", table_name).as_str())
2672 }
2673
2674 fn table_reserved_file_path(&self, table_name: &str) -> Path {
2676 self.base_path
2677 .clone()
2678 .join(format!("{}.lance", table_name).as_str())
2679 .join(".lance-reserved")
2680 }
2681
2682 fn table_deregistered_file_path(&self, table_name: &str) -> Path {
2684 self.base_path
2685 .clone()
2686 .join(format!("{}.lance", table_name).as_str())
2687 .join(".lance-deregistered")
2688 }
2689
2690 pub(crate) async fn check_table_status(&self, table_name: &str) -> Result<TableStatus> {
2696 let table_path = self.table_path(table_name);
2697 match self.object_store.read_dir(table_path).await {
2698 Ok(entries) => {
2699 let exists = !entries.is_empty();
2700 let is_deregistered = entries.iter().any(|e| e.ends_with(".lance-deregistered"));
2701 let has_reserved_file = entries.iter().any(|e| e.ends_with(".lance-reserved"));
2702 Ok(TableStatus {
2703 exists,
2704 is_deregistered,
2705 has_reserved_file,
2706 })
2707 }
2708 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(TableStatus {
2711 exists: false,
2712 is_deregistered: false,
2713 has_reserved_file: false,
2714 }),
2715 Err(e) => Err(Self::classify_storage_error(e)),
2719 }
2720 }
2721
2722 fn classify_storage_error(err: Error) -> Error {
2727 if matches!(&err, Error::Namespace { .. }) {
2728 return err;
2729 }
2730 let detail = err.to_string();
2731 if let Error::IO { source, .. } = &err
2732 && let Some(os_err) = source.downcast_ref::<ObjectStoreError>()
2733 {
2734 if is_throttle_error(os_err) {
2735 return NamespaceError::Throttling {
2736 message: format!(
2737 "Storage request was throttled while resolving table: {detail}"
2738 ),
2739 }
2740 .into();
2741 }
2742 if Self::is_service_unavailable_error(os_err) {
2743 return NamespaceError::ServiceUnavailable {
2744 message: format!("Storage service unavailable while resolving table: {detail}"),
2745 }
2746 .into();
2747 }
2748 }
2749 NamespaceError::Internal {
2750 message: format!("Storage error while resolving table: {detail}"),
2751 }
2752 .into()
2753 }
2754
2755 fn is_service_unavailable_error(err: &ObjectStoreError) -> bool {
2759 if let ObjectStoreError::Generic { source, .. } = err {
2760 let message = source.to_string().to_ascii_lowercase();
2761 message.contains("503 service unavailable")
2762 || message.contains("502 bad gateway")
2763 || message.contains("504 gateway timeout")
2764 } else {
2765 false
2766 }
2767 }
2768
2769 fn is_manifest_table_absent_error(err: &Error) -> bool {
2773 if manifest::ManifestNamespace::is_not_found_load_error(err) {
2774 return true;
2775 }
2776 if let Error::Namespace { source, .. } = err
2777 && let Some(ns_err) = source.downcast_ref::<NamespaceError>()
2778 {
2779 return matches!(ns_err, NamespaceError::TableNotFound { .. });
2780 }
2781 false
2782 }
2783
2784 fn map_open_error(err: Error, not_found: NamespaceError) -> Error {
2788 if matches!(&err, Error::IO { .. })
2789 && !manifest::ManifestNamespace::is_not_found_load_error(&err)
2790 {
2791 return Self::classify_storage_error(err);
2792 }
2793 not_found.into()
2794 }
2795
2796 async fn get_storage_options_for_table(
2816 &self,
2817 table_uri: &str,
2818 vend_credentials: bool,
2819 identity: Option<&Identity>,
2820 ) -> Result<Option<HashMap<String, String>>> {
2821 if vend_credentials && let Some(ref vendor) = self.credential_vendor {
2822 let vended = vendor.vend_credentials(table_uri, identity).await?;
2823 return Ok(Some(vended.storage_options));
2824 }
2825 if self.vend_input_storage_options {
2828 let mut options = self.storage_options.clone().unwrap_or_default();
2829 if let Some(refresh_interval_millis) =
2831 self.vend_input_storage_options_refresh_interval_millis
2832 {
2833 let now_millis = std::time::SystemTime::now()
2834 .duration_since(std::time::UNIX_EPOCH)
2835 .unwrap()
2836 .as_millis() as u64;
2837 let expires_at_millis = now_millis + refresh_interval_millis;
2838 options.insert(
2839 "expires_at_millis".to_string(),
2840 expires_at_millis.to_string(),
2841 );
2842 }
2843 return Ok(Some(options));
2844 }
2845 Ok(None)
2848 }
2849
2850 pub async fn migrate(&self) -> Result<usize> {
2904 let Some(manifest_ns) = self.manifest_ns_for_write().await? else {
2906 return Ok(0); };
2908
2909 let manifest_locations = manifest_ns.list_manifest_table_locations().await?;
2911
2912 let dir_tables = self
2915 .filter_declared_tables(self.list_directory_tables().await?, false)
2916 .await?;
2917
2918 let mut migrated_count = 0;
2923 for table_name in dir_tables {
2924 let dir_name = format!("{}.lance", table_name);
2926 if !manifest_locations.contains(&dir_name) {
2927 manifest_ns.register_table(&table_name, dir_name).await?;
2928 migrated_count += 1;
2929 }
2930 }
2931
2932 Ok(migrated_count)
2933 }
2934
2935 async fn delete_physical_version_files(
2943 &self,
2944 table_entries: &[TableDeleteEntry],
2945 branch: Option<&str>,
2946 ) -> Result<i64> {
2947 let mut deleted_count = 0i64;
2948 for te in table_entries {
2949 let table_uri = self.resolve_table_location(&te.table_id).await?;
2950 let table_uri = match branch {
2951 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
2952 None => table_uri,
2953 };
2954 let table_path = self.object_store_path_from_uri(&table_uri)?;
2955 let versions_dir_path = table_path.clone().join(VERSIONS_DIR);
2956
2957 let manifest_metas: Vec<_> = self
2959 .object_store
2960 .read_dir_all(&versions_dir_path, None)
2961 .try_collect()
2962 .await
2963 .map_err(|e| {
2964 lance_core::Error::from(NamespaceError::Internal {
2965 message: format!(
2966 "Failed to list manifest files for table at '{}': {}",
2967 table_uri, e
2968 ),
2969 })
2970 })?;
2971 let location_by_version: HashMap<u64, Path> = manifest_metas
2972 .into_iter()
2973 .filter_map(|meta| {
2974 let version = Self::manifest_version_from_filename(meta.location.filename()?)?;
2975 Some((version, meta.location))
2976 })
2977 .collect();
2978
2979 for (&v, version_path) in &location_by_version {
2980 let vi = v as i64;
2981 if !te.ranges.iter().any(|&(s, e)| vi >= s && (e < 0 || vi < e)) {
2982 continue;
2983 }
2984 match self.object_store.inner.delete(version_path).await {
2985 Ok(_) => {
2986 deleted_count += 1;
2987 }
2988 Err(object_store::Error::NotFound { .. }) => {}
2989 Err(e) => {
2990 return Err(NamespaceError::Internal {
2991 message: format!(
2992 "Failed to delete version {} for table at '{}': {}",
2993 v, table_uri, e
2994 ),
2995 }
2996 .into());
2997 }
2998 }
2999 }
3000 }
3001 Ok(deleted_count)
3002 }
3003
3004 #[allow(clippy::too_many_arguments)]
3009 fn apply_query_params_to_scanner(
3010 scanner: &mut Scanner,
3011 filter: Option<&str>,
3012 columns: Option<&QueryTableRequestColumns>,
3013 vector_column: Option<&str>,
3014 vector: &QueryTableRequestVector,
3015 k: i32,
3016 offset: Option<i32>,
3017 prefilter: Option<bool>,
3018 bypass_vector_index: Option<bool>,
3019 nprobes: Option<i32>,
3020 ef: Option<i32>,
3021 refine_factor: Option<i32>,
3022 distance_type: Option<&str>,
3023 fast_search_flag: Option<bool>,
3024 with_row_id: Option<bool>,
3025 lower_bound: Option<f32>,
3026 upper_bound: Option<f32>,
3027 operation: &str,
3028 ) -> Result<()> {
3029 if let Some(pf) = prefilter {
3031 scanner.prefilter(pf);
3032 }
3033
3034 if let Some(filter) = filter {
3035 scanner.filter(filter).map_err(|e| {
3036 Error::invalid_input_source(
3037 format!("Invalid filter expression for {}: {}", operation, e).into(),
3038 )
3039 })?;
3040 }
3041
3042 if let Some(cols) = columns {
3043 if let Some(ref names) = cols.column_names {
3044 scanner.project(names.as_slice()).map_err(|e| {
3045 Error::invalid_input_source(
3046 format!("Invalid column projection for {}: {}", operation, e).into(),
3047 )
3048 })?;
3049 } else if let Some(ref aliases) = cols.column_aliases {
3050 let pairs: Vec<(&str, &str)> = aliases
3052 .iter()
3053 .map(|(alias, src)| (alias.as_str(), src.as_str()))
3054 .collect();
3055 scanner.project_with_transform(&pairs).map_err(|e| {
3056 Error::invalid_input_source(
3057 format!("Invalid column aliases for {}: {}", operation, e).into(),
3058 )
3059 })?;
3060 }
3061 }
3062
3063 let query_vec: Option<Vec<f32>> = vector
3065 .single_vector
3066 .as_ref()
3067 .filter(|v| !v.is_empty())
3068 .cloned()
3069 .or_else(|| {
3070 vector
3071 .multi_vector
3072 .as_ref()
3073 .and_then(|mv| mv.first())
3074 .filter(|v| !v.is_empty())
3075 .cloned()
3076 });
3077
3078 if let Some(q_vec) = query_vec {
3079 let col = vector_column.unwrap_or("vector");
3080 let q = Arc::new(Float32Array::from(q_vec));
3081 scanner
3082 .nearest(col, q.as_ref(), k.max(1) as usize)
3083 .map_err(|e| {
3084 Error::invalid_input_source(
3085 format!("Invalid vector query for {}: {}", operation, e).into(),
3086 )
3087 })?;
3088
3089 if let Some(n) = nprobes {
3091 scanner.nprobes(n.max(1) as usize);
3092 }
3093 if let Some(e) = ef {
3094 scanner.ef(e.max(1) as usize);
3095 }
3096 if let Some(rf) = refine_factor {
3097 scanner.refine(rf.max(0) as u32);
3098 }
3099 if let Some(true) = bypass_vector_index {
3101 scanner.use_index(false);
3102 }
3103 if let Some(true) = fast_search_flag {
3104 scanner.fast_search();
3105 }
3106 if lower_bound.is_some() || upper_bound.is_some() {
3107 scanner.distance_range(lower_bound, upper_bound);
3108 }
3109 if let Some(dt) = distance_type {
3110 let metric = Self::parse_metric_type(Some(dt))?;
3111 scanner.distance_metric(metric);
3112 }
3113 if let Some(off) = offset.filter(|&o| o > 0) {
3115 scanner.limit(None, Some(off as i64)).map_err(|e| {
3116 Error::invalid_input_source(
3117 format!("Invalid offset for {}: {}", operation, e).into(),
3118 )
3119 })?;
3120 }
3121 } else {
3122 let limit = if k > 0 { Some(k as i64) } else { None };
3124 scanner
3125 .limit(limit, offset.map(|o| o as i64))
3126 .map_err(|e| {
3127 Error::invalid_input_source(
3128 format!("Invalid limit/offset for {}: {}", operation, e).into(),
3129 )
3130 })?;
3131 }
3132
3133 if let Some(true) = with_row_id {
3134 scanner.with_row_id();
3135 }
3136
3137 Ok(())
3138 }
3139
3140 pub fn retrieve_ops_metrics(&self) -> HashMap<String, u64> {
3147 self.ops_metrics
3148 .as_ref()
3149 .map(|m| m.retrieve())
3150 .unwrap_or_default()
3151 }
3152
3153 pub fn reset_ops_metrics(&self) {
3157 if let Some(ref metrics) = self.ops_metrics {
3158 metrics.reset();
3159 }
3160 }
3161
3162 fn record_op(&self, operation: &str) {
3164 if let Some(ref metrics) = self.ops_metrics {
3165 metrics.increment(operation);
3166 }
3167 }
3168}
3169
3170#[async_trait]
3171impl LanceNamespace for DirectoryNamespace {
3172 async fn list_namespaces(
3173 &self,
3174 request: ListNamespacesRequest,
3175 ) -> Result<ListNamespacesResponse> {
3176 self.record_op("list_namespaces");
3177 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3178 return manifest_ns.list_namespaces(request).await;
3179 }
3180
3181 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3182 return Err(self.child_namespace_requires_manifest_error());
3183 }
3184 Self::validate_root_namespace_id(&request.id)?;
3185 Ok(ListNamespacesResponse::new(vec![]))
3186 }
3187
3188 async fn describe_namespace(
3189 &self,
3190 request: DescribeNamespaceRequest,
3191 ) -> Result<DescribeNamespaceResponse> {
3192 self.record_op("describe_namespace");
3193 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3194 return manifest_ns.describe_namespace(request).await;
3195 }
3196
3197 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3198 return Err(self.child_namespace_requires_manifest_error());
3199 }
3200 Self::validate_root_namespace_id(&request.id)?;
3201 #[allow(clippy::needless_update)]
3202 Ok(DescribeNamespaceResponse {
3203 properties: Some(HashMap::new()),
3204 ..Default::default()
3205 })
3206 }
3207
3208 async fn create_namespace(
3209 &self,
3210 request: CreateNamespaceRequest,
3211 ) -> Result<CreateNamespaceResponse> {
3212 self.record_op("create_namespace");
3213 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3214 return manifest_ns.create_namespace(request).await;
3215 }
3216
3217 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3218 return Err(NamespaceError::NamespaceAlreadyExists {
3219 message: "root namespace".to_string(),
3220 }
3221 .into());
3222 }
3223
3224 Err(NamespaceError::Unsupported {
3225 message: "Child namespaces are only supported when manifest mode is enabled"
3226 .to_string(),
3227 }
3228 .into())
3229 }
3230
3231 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
3232 self.record_op("drop_namespace");
3233 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3234 return manifest_ns.drop_namespace(request).await;
3235 }
3236
3237 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3238 return Err(NamespaceError::InvalidInput {
3239 message: "Root namespace cannot be dropped".to_string(),
3240 }
3241 .into());
3242 }
3243
3244 Err(NamespaceError::Unsupported {
3245 message: "Child namespaces are only supported when manifest mode is enabled"
3246 .to_string(),
3247 }
3248 .into())
3249 }
3250
3251 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
3252 self.record_op("namespace_exists");
3253 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3254 return manifest_ns.namespace_exists(request).await;
3255 }
3256
3257 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3258 return Ok(());
3259 }
3260
3261 Err(self.child_namespace_requires_manifest_error())
3262 }
3263
3264 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
3265 self.record_op("list_tables");
3266 let namespace_id = request.id.as_ref().ok_or_else(|| {
3268 lance_core::Error::from(NamespaceError::InvalidInput {
3269 message: "Namespace ID is required".to_string(),
3270 })
3271 })?;
3272
3273 if !namespace_id.is_empty() {
3275 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3276 return manifest_ns.list_tables(request).await;
3277 }
3278 return Err(self.child_namespace_requires_manifest_error());
3279 }
3280
3281 if let Some(manifest_ns) = self.manifest_ns_for_read()
3283 && !self.dir_listing_enabled
3284 {
3285 return manifest_ns.list_tables(request).await;
3286 }
3287 if !self.dir_listing_enabled {
3288 return Ok(ListTablesResponse::new(vec![]));
3289 }
3290
3291 let mut tables = if self.manifest_ns_for_read().is_some()
3294 && self.dir_listing_enabled
3295 && self.dir_listing_to_manifest_migration_enabled
3296 {
3297 let manifest_locations = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3299 manifest_ns.list_manifest_table_locations().await?
3300 } else {
3301 std::collections::HashSet::new()
3302 };
3303
3304 let mut manifest_request = request.clone();
3306 manifest_request.limit = None;
3307 manifest_request.page_token = None;
3308 let manifest_tables = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3309 let manifest_response = manifest_ns.list_tables(manifest_request).await?;
3310 manifest_response.tables
3311 } else {
3312 vec![]
3313 };
3314
3315 let mut all_tables: Vec<String> = manifest_tables;
3318 let dir_tables = self.list_directory_tables().await?;
3319 for table_name in dir_tables {
3320 let full_location = format!("{}/{}.lance", self.root, table_name);
3323 let relative_location = format!("{}.lance", table_name);
3324 if !manifest_locations.contains(&full_location)
3325 && !manifest_locations.contains(&relative_location)
3326 {
3327 all_tables.push(table_name);
3328 }
3329 }
3330
3331 all_tables
3332 } else {
3333 self.list_directory_tables().await?
3334 };
3335
3336 tables = self
3337 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
3338 .await?;
3339
3340 let next_page_token =
3342 Self::apply_pagination(&mut tables, request.page_token, request.limit);
3343 let mut response = ListTablesResponse::new(tables);
3344 response.page_token = next_page_token;
3345 Ok(response)
3346 }
3347
3348 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
3349 self.record_op("describe_table");
3350 self.describe_table_impl(request).await
3351 }
3352
3353 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
3354 self.record_op("table_exists");
3355 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
3356 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
3357 let skip_manifest_for_root = self.dir_listing_enabled
3358 && is_root_level
3359 && !self.dir_listing_to_manifest_migration_enabled;
3360 if let Some(manifest_ns) = self.manifest_ns_for_read()
3361 && !skip_manifest_for_root
3362 {
3363 match manifest_ns.table_exists(request.clone()).await {
3364 Ok(()) => return Ok(()),
3365 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
3366 return Err(e);
3369 }
3370 Err(e) if self.dir_listing_enabled && is_root_level => {
3371 if !Self::is_manifest_table_absent_error(&e) {
3375 return Err(Self::classify_storage_error(e));
3376 }
3377 }
3378 Err(e) => return Err(e),
3379 }
3380 }
3381 if is_child_table {
3382 return Err(self.child_namespace_requires_manifest_error());
3383 }
3384
3385 let table_name = Self::table_name_from_id(&request.id)?;
3386 let table_id = Self::format_table_id_from_request(&request.id);
3387 if !self.dir_listing_enabled {
3388 return Err(NamespaceError::TableNotFound { message: table_id }.into());
3389 }
3390
3391 let status = self.check_table_status(&table_name).await?;
3393
3394 if !status.exists {
3395 return Err(NamespaceError::TableNotFound {
3396 message: table_id.clone(),
3397 }
3398 .into());
3399 }
3400
3401 if status.is_deregistered {
3402 return Err(NamespaceError::TableNotFound {
3403 message: format!("Table is deregistered: {}", table_id),
3404 }
3405 .into());
3406 }
3407
3408 Ok(())
3409 }
3410
3411 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3412 self.record_op("drop_table");
3413 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3414 return manifest_ns.drop_table(request).await;
3415 }
3416
3417 let table_name = Self::table_name_from_id(&request.id)?;
3418 let table_uri = self.table_full_uri(&table_name);
3419 let table_path = self.table_path(&table_name);
3420
3421 self.object_store
3422 .remove_dir_all(table_path)
3423 .await
3424 .map_err(|e| {
3425 lance_core::Error::from(NamespaceError::Internal {
3426 message: format!("Failed to drop table {}: {:?}", table_name, e),
3427 })
3428 })?;
3429
3430 Ok(DropTableResponse {
3431 id: request.id,
3432 location: Some(table_uri),
3433 ..Default::default()
3434 })
3435 }
3436
3437 async fn create_table(
3438 &self,
3439 request: CreateTableRequest,
3440 request_data: Bytes,
3441 ) -> Result<CreateTableResponse> {
3442 self.record_op("create_table");
3443 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3444 return manifest_ns.create_table(request, request_data).await;
3445 }
3446
3447 Self::validate_dir_only_properties(request.properties.as_ref(), "create_table")?;
3448
3449 let table_name = Self::table_name_from_id(&request.id)?;
3450 let table_uri = self.table_full_uri(&table_name);
3451 let status = self.check_table_status(&table_name).await?;
3452 let (reader, _num_rows) =
3453 Self::ipc_reader_from_request_data(&request_data, "create_table")?;
3454
3455 if status.exists && self.table_has_actual_manifests(&table_name).await? {
3456 return Err(NamespaceError::TableAlreadyExists {
3457 message: table_name,
3458 }
3459 .into());
3460 }
3461
3462 let write_result = self
3463 .write_reader_to_table(
3464 &table_uri,
3465 reader,
3466 WriteMode::Create,
3467 request.storage_options.clone(),
3468 )
3469 .await;
3470 if let Err(err) = write_result {
3471 if self.table_uri_has_actual_manifests(&table_uri).await? {
3472 return Err(NamespaceError::TableAlreadyExists {
3473 message: table_name,
3474 }
3475 .into());
3476 }
3477 return Err(err);
3478 }
3479 Ok(CreateTableResponse {
3480 version: Some(1),
3481 location: Some(table_uri),
3482 storage_options: self.storage_options.clone(),
3483 properties: request.properties,
3484 ..Default::default()
3485 })
3486 }
3487
3488 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3489 self.record_op("declare_table");
3490 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3491 let mut response = manifest_ns.declare_table(request.clone()).await?;
3492 if let Some(ref location) = response.location {
3493 let vend = request.vend_credentials.unwrap_or(true);
3495 let identity = request.identity.as_deref();
3496 response.storage_options = self
3497 .get_storage_options_for_table(location, vend, identity)
3498 .await?;
3499 }
3500 if self.table_version_tracking_enabled {
3502 response.managed_versioning = Some(true);
3503 }
3504 return Ok(response);
3505 }
3506
3507 Self::validate_dir_only_properties(request.properties.as_ref(), "declare_table")?;
3508
3509 let table_name = Self::table_name_from_id(&request.id)?;
3510 let table_uri = self.table_full_uri(&table_name);
3511
3512 if let Some(location) = &request.location {
3514 let location = location.trim_end_matches('/');
3515 if location != table_uri {
3516 return Err(NamespaceError::InvalidInput {
3517 message: format!(
3518 "Cannot declare table {} at location {}, must be at location {}",
3519 table_name, location, table_uri
3520 ),
3521 }
3522 .into());
3523 }
3524 }
3525
3526 let status = self.check_table_status(&table_name).await?;
3530 if status.exists && !status.has_reserved_file {
3531 return Err(NamespaceError::TableAlreadyExists {
3533 message: table_name.to_string(),
3534 }
3535 .into());
3536 }
3537
3538 let reserved_file_path = self.table_reserved_file_path(&table_name);
3542
3543 put_marker_file_atomic(
3544 &self.object_store,
3545 &reserved_file_path,
3546 &format!("table {}", table_name),
3547 )
3548 .await
3549 .map_err(|e| match e {
3550 MarkerFileError::AlreadyExists { .. } => {
3551 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3552 message: table_name.to_string(),
3553 })
3554 }
3555 MarkerFileError::Other { message } => {
3556 lance_core::Error::from(NamespaceError::Internal { message })
3557 }
3558 })?;
3559
3560 let vend_credentials = request.vend_credentials.unwrap_or(true);
3562 let identity = request.identity.as_deref();
3563 let storage_options = self
3564 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
3565 .await?;
3566
3567 Ok(DeclareTableResponse {
3568 location: Some(table_uri),
3569 storage_options,
3570 properties: request.properties,
3571 managed_versioning: if self.table_version_tracking_enabled {
3572 Some(true)
3573 } else {
3574 None
3575 },
3576 ..Default::default()
3577 })
3578 }
3579
3580 async fn register_table(
3581 &self,
3582 request: lance_namespace::models::RegisterTableRequest,
3583 ) -> Result<lance_namespace::models::RegisterTableResponse> {
3584 self.record_op("register_table");
3585 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3587 return LanceNamespace::register_table(manifest_ns.as_ref(), request).await;
3588 }
3589
3590 Err(NamespaceError::Unsupported {
3592 message: "register_table is only supported when manifest mode is enabled".to_string(),
3593 }
3594 .into())
3595 }
3596
3597 async fn deregister_table(
3598 &self,
3599 request: lance_namespace::models::DeregisterTableRequest,
3600 ) -> Result<lance_namespace::models::DeregisterTableResponse> {
3601 self.record_op("deregister_table");
3602 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3604 return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await;
3605 }
3606
3607 let table_name = Self::table_name_from_id(&request.id)?;
3609 let table_uri = self.table_full_uri(&table_name);
3610
3611 let status = self.check_table_status(&table_name).await?;
3614
3615 if !status.exists {
3616 return Err(NamespaceError::TableNotFound {
3617 message: table_name.to_string(),
3618 }
3619 .into());
3620 }
3621
3622 if status.is_deregistered {
3623 return Err(NamespaceError::TableNotFound {
3624 message: format!("Table is already deregistered: {}", table_name),
3625 }
3626 .into());
3627 }
3628
3629 let deregistered_path = self.table_deregistered_file_path(&table_name);
3635 put_marker_file_atomic(
3636 &self.object_store,
3637 &deregistered_path,
3638 &format!("deregistration marker for table {}", table_name),
3639 )
3640 .await
3641 .map_err(|e| match e {
3642 MarkerFileError::AlreadyExists { .. } => {
3643 lance_core::Error::from(NamespaceError::InvalidTableState {
3644 message: format!("Table is already deregistered: {}", table_name),
3645 })
3646 }
3647 MarkerFileError::Other { message } => {
3648 lance_core::Error::from(NamespaceError::Internal { message })
3649 }
3650 })?;
3651
3652 Ok(lance_namespace::models::DeregisterTableResponse {
3653 id: request.id,
3654 location: Some(table_uri),
3655 ..Default::default()
3656 })
3657 }
3658
3659 async fn alter_table_add_columns(
3660 &self,
3661 request: AlterTableAddColumnsRequest,
3662 ) -> Result<AlterTableAddColumnsResponse> {
3663 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3664 return manifest_ns.alter_table_add_columns(request).await;
3665 }
3666
3667 let table_name = Self::table_name_from_id(&request.id)?;
3669 let table_uri = self.table_full_uri(&table_name);
3670
3671 let status = self.check_table_status(&table_name).await?;
3673 if !status.exists {
3674 return Err(NamespaceError::TableNotFound {
3675 message: table_name,
3676 }
3677 .into());
3678 }
3679 if status.is_deregistered {
3680 return Err(NamespaceError::TableNotFound {
3681 message: format!("Table is deregistered: {}", table_name),
3682 }
3683 .into());
3684 }
3685
3686 let mut dataset = self
3687 .configured_builder(&table_uri)
3688 .load()
3689 .await
3690 .map_err(|e| {
3691 Error::io_source(box_error(std::io::Error::other(format!(
3692 "Failed to open dataset: {}",
3693 e
3694 ))))
3695 })?;
3696
3697 let sql_expressions = build_sql_expressions(&request.new_columns)?;
3698
3699 dataset
3700 .add_columns(
3701 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3702 None,
3703 None,
3704 )
3705 .await
3706 .map_err(|e| {
3707 Error::io_source(box_error(std::io::Error::other(format!(
3708 "Failed to add columns: {}",
3709 e
3710 ))))
3711 })?;
3712
3713 let version = dataset.version().version as i64;
3714 Ok(AlterTableAddColumnsResponse::new(version))
3715 }
3716
3717 async fn alter_table_alter_columns(
3718 &self,
3719 request: AlterTableAlterColumnsRequest,
3720 ) -> Result<AlterTableAlterColumnsResponse> {
3721 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3722 return manifest_ns.alter_table_alter_columns(request).await;
3723 }
3724
3725 let table_name = Self::table_name_from_id(&request.id)?;
3726 let table_uri = self.table_full_uri(&table_name);
3727
3728 let status = self.check_table_status(&table_name).await?;
3730 if !status.exists {
3731 return Err(NamespaceError::TableNotFound {
3732 message: table_name,
3733 }
3734 .into());
3735 }
3736 if status.is_deregistered {
3737 return Err(NamespaceError::TableNotFound {
3738 message: format!("Table is deregistered: {}", table_name),
3739 }
3740 .into());
3741 }
3742
3743 let mut dataset = self
3744 .configured_builder(&table_uri)
3745 .load()
3746 .await
3747 .map_err(|e| {
3748 Error::io_source(box_error(std::io::Error::other(format!(
3749 "Failed to open dataset: {}",
3750 e
3751 ))))
3752 })?;
3753
3754 let alterations = build_column_alterations(&request.alterations)?;
3755
3756 dataset.alter_columns(&alterations).await.map_err(|e| {
3757 Error::io_source(box_error(std::io::Error::other(format!(
3758 "Failed to alter columns: {}",
3759 e
3760 ))))
3761 })?;
3762
3763 let version = dataset.version().version as i64;
3764 Ok(AlterTableAlterColumnsResponse::new(version))
3765 }
3766
3767 async fn alter_table_drop_columns(
3768 &self,
3769 request: AlterTableDropColumnsRequest,
3770 ) -> Result<AlterTableDropColumnsResponse> {
3771 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3772 return manifest_ns.alter_table_drop_columns(request).await;
3773 }
3774
3775 let table_name = Self::table_name_from_id(&request.id)?;
3776 let table_uri = self.table_full_uri(&table_name);
3777
3778 let status = self.check_table_status(&table_name).await?;
3780 if !status.exists {
3781 return Err(NamespaceError::TableNotFound {
3782 message: table_name,
3783 }
3784 .into());
3785 }
3786 if status.is_deregistered {
3787 return Err(NamespaceError::TableNotFound {
3788 message: format!("Table is deregistered: {}", table_name),
3789 }
3790 .into());
3791 }
3792
3793 let mut dataset = self
3794 .configured_builder(&table_uri)
3795 .load()
3796 .await
3797 .map_err(|e| {
3798 Error::io_source(box_error(std::io::Error::other(format!(
3799 "Failed to open dataset: {}",
3800 e
3801 ))))
3802 })?;
3803
3804 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
3805 dataset.drop_columns(&columns).await.map_err(|e| {
3806 Error::io_source(box_error(std::io::Error::other(format!(
3807 "Failed to drop columns: {}",
3808 e
3809 ))))
3810 })?;
3811
3812 let version = dataset.version().version as i64;
3813 Ok(AlterTableDropColumnsResponse::new(version))
3814 }
3815
3816 async fn list_table_versions(
3817 &self,
3818 request: ListTableVersionsRequest,
3819 ) -> Result<ListTableVersionsResponse> {
3820 self.record_op("list_table_versions");
3821 let branch = Self::normalized_branch(request.branch.as_deref())?;
3822 let table_uri = self.resolve_table_location(&request.id).await?;
3823 let table_uri = match branch {
3824 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3825 None => table_uri,
3826 };
3827 let want_descending = request.descending == Some(true);
3828 let table_versions = self
3829 .list_table_versions_from_storage(&table_uri, want_descending, request.limit)
3830 .await?;
3831
3832 Ok(ListTableVersionsResponse {
3833 versions: table_versions,
3834 page_token: None,
3835 })
3836 }
3837
3838 async fn create_table_version(
3839 &self,
3840 request: CreateTableVersionRequest,
3841 ) -> Result<CreateTableVersionResponse> {
3842 self.record_op("create_table_version");
3843 let branch = Self::normalized_branch(request.branch.as_deref())?;
3844 let table_uri = self.resolve_table_location(&request.id).await?;
3845 let (table_uri, table_path, branch_parent_version) = match branch {
3846 Some(b) => self.resolve_branch_for_commit(&table_uri, b).await?,
3847 None => {
3848 let table_path = self.object_store_path_from_uri(&table_uri)?;
3849 (table_uri, table_path, None)
3850 }
3851 };
3852
3853 let staging_manifest_path = &request.manifest_path;
3854 let version = request.version as u64;
3855
3856 let naming_scheme = match request.naming_scheme.as_deref() {
3858 Some("V1") => ManifestNamingScheme::V1,
3859 _ => ManifestNamingScheme::V2,
3860 };
3861
3862 let final_path = naming_scheme.manifest_path(&table_path, version);
3864
3865 let staging_path = Path::parse(staging_manifest_path).map_err(|e| {
3866 lance_core::Error::from(NamespaceError::InvalidInput {
3867 message: format!(
3868 "Invalid staging manifest path '{}': {}",
3869 staging_manifest_path, e
3870 ),
3871 })
3872 })?;
3873
3874 match self.object_store.inner.head(&final_path).await {
3876 Ok(existing_meta) => {
3877 return self
3878 .resolve_existing_table_version(ExistingTableVersionResolve {
3879 staging_path: &staging_path,
3880 final_path: &final_path,
3881 version,
3882 table_uri: &table_uri,
3883 final_meta: &existing_meta,
3884 request_manifest_size: request.manifest_size,
3885 })
3886 .await;
3887 }
3888 Err(ObjectStoreError::NotFound { .. }) => {}
3889 Err(e) => {
3890 return Err(lance_core::Error::from(NamespaceError::Internal {
3891 message: format!(
3892 "Failed to stat version {} for table at '{}': {}",
3893 version, table_uri, e
3894 ),
3895 }));
3896 }
3897 }
3898
3899 let is_branch = branch.is_some();
3902 self.enforce_create_table_version_cas(
3903 &table_path,
3904 version,
3905 &table_uri,
3906 is_branch,
3907 branch_parent_version,
3908 )
3909 .await?;
3910
3911 let copy_result = self
3913 .materialize_version_manifest_create(&staging_path, &final_path, staging_manifest_path)
3914 .await;
3915
3916 match copy_result {
3917 Ok(()) => {}
3918 Err(ObjectStoreError::AlreadyExists { .. })
3919 | Err(ObjectStoreError::Precondition { .. }) => {
3920 let existing_meta = self.object_store.inner.head(&final_path).await.map_err(|e| {
3922 lance_core::Error::from(NamespaceError::Internal {
3923 message: format!(
3924 "Version {} conflict for table at '{}' but failed to stat winner: {}",
3925 version, table_uri, e
3926 ),
3927 })
3928 })?;
3929 return self
3930 .resolve_existing_table_version(ExistingTableVersionResolve {
3931 staging_path: &staging_path,
3932 final_path: &final_path,
3933 version,
3934 table_uri: &table_uri,
3935 final_meta: &existing_meta,
3936 request_manifest_size: request.manifest_size,
3937 })
3938 .await;
3939 }
3940 Err(ObjectStoreError::NotFound { .. }) => {
3941 return Err(lance_core::Error::from(NamespaceError::InvalidInput {
3942 message: format!(
3943 "Staging manifest not found at '{}' for version {} of table at '{}'",
3944 staging_manifest_path, version, table_uri
3945 ),
3946 }));
3947 }
3948 Err(e) => {
3949 return Err(lance_core::Error::from(NamespaceError::Internal {
3950 message: format!(
3951 "Failed to create version {} for table at '{}': {}",
3952 version, table_uri, e
3953 ),
3954 }));
3955 }
3956 }
3957
3958 let final_meta = self
3959 .object_store
3960 .inner
3961 .head(&final_path)
3962 .await
3963 .map_err(|e| {
3964 lance_core::Error::from(NamespaceError::Internal {
3965 message: format!(
3966 "Failed to stat created version {} for table at '{}': {}",
3967 version, table_uri, e
3968 ),
3969 })
3970 })?;
3971
3972 if let Err(e) = self.object_store.inner.delete(&staging_path).await {
3974 log::warn!(
3975 "Failed to delete staging manifest at '{}': {:?}",
3976 staging_path,
3977 e
3978 );
3979 }
3980
3981 Ok(Self::create_table_version_response(
3982 version,
3983 &final_path,
3984 &final_meta,
3985 ))
3986 }
3987
3988 async fn describe_table_version(
3989 &self,
3990 request: DescribeTableVersionRequest,
3991 ) -> Result<DescribeTableVersionResponse> {
3992 self.record_op("describe_table_version");
3993 let branch = Self::normalized_branch(request.branch.as_deref())?;
3994 let table_uri = self.resolve_table_location(&request.id).await?;
3995 let table_uri = match branch {
3996 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3997 None => table_uri,
3998 };
3999 let versions = self
4000 .list_table_versions_from_storage(&table_uri, true, None)
4001 .await?;
4002 let table_version = if let Some(requested_version) = request.version {
4003 versions
4004 .into_iter()
4005 .find(|version| version.version == requested_version)
4006 .ok_or_else(|| {
4007 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4008 message: format!(
4009 "version {} for table {}",
4010 requested_version,
4011 Self::format_table_id_from_request(&request.id)
4012 ),
4013 })
4014 })?
4015 } else {
4016 versions.into_iter().next().ok_or_else(|| {
4017 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4018 message: format!(
4019 "latest version for table {}",
4020 Self::format_table_id_from_request(&request.id)
4021 ),
4022 })
4023 })?
4024 };
4025
4026 Ok(DescribeTableVersionResponse {
4027 version: Box::new(table_version),
4028 })
4029 }
4030
4031 async fn batch_delete_table_versions(
4032 &self,
4033 request: BatchDeleteTableVersionsRequest,
4034 ) -> Result<BatchDeleteTableVersionsResponse> {
4035 self.record_op("batch_delete_table_versions");
4036 let branch = Self::normalized_branch(request.branch.as_deref())?;
4037 let ranges: Vec<(i64, i64)> = request
4040 .ranges
4041 .iter()
4042 .map(|r| (r.start_version, r.end_version))
4043 .collect();
4044
4045 const MAX_VERSIONS_PER_REQUEST: i128 = 1_000_000;
4049 let requested: i128 = ranges
4050 .iter()
4051 .map(|(s, e)| {
4052 if *e < 0 {
4053 0
4054 } else {
4055 (*e as i128 - *s as i128).max(0)
4056 }
4057 })
4058 .sum();
4059 if requested > MAX_VERSIONS_PER_REQUEST {
4060 return Err(NamespaceError::InvalidInput {
4061 message: format!(
4062 "batch_delete requested {} versions; limit is {}",
4063 requested, MAX_VERSIONS_PER_REQUEST
4064 ),
4065 }
4066 .into());
4067 }
4068
4069 let table_entries = vec![TableDeleteEntry {
4070 table_id: request.id.clone(),
4071 ranges,
4072 }];
4073
4074 let total_deleted_count = self
4075 .delete_physical_version_files(&table_entries, branch)
4076 .await?;
4077
4078 Ok(BatchDeleteTableVersionsResponse {
4079 deleted_count: Some(total_deleted_count),
4080 transaction_id: None,
4081 })
4082 }
4083
4084 async fn create_table_index(
4085 &self,
4086 request: CreateTableIndexRequest,
4087 ) -> Result<CreateTableIndexResponse> {
4088 self.record_op("create_table_index");
4089 let table_uri = self.resolve_table_location(&request.id).await?;
4090 let mut dataset = self
4091 .load_dataset(&table_uri, None, "create_table_index")
4092 .await?;
4093 let index_request = Self::build_index_params(&request)?;
4094
4095 dataset
4096 .create_index(
4097 &[request.column.as_str()],
4098 index_request.index_type(),
4099 request.name.clone(),
4100 index_request.params(),
4101 false,
4102 )
4103 .await
4104 .map_err(|e| {
4105 let err_msg = format!("{}", e);
4106 let ns_err = if err_msg.contains("already exists") {
4107 NamespaceError::TableIndexAlreadyExists {
4108 message: format!(
4109 "Index '{}' already exists on table '{}': {:?}",
4110 request.name.as_deref().unwrap_or("<auto-generated>"),
4111 table_uri,
4112 e
4113 ),
4114 }
4115 } else if err_msg.contains("not found") || err_msg.contains("does not exist") {
4116 NamespaceError::TableColumnNotFound {
4117 message: format!(
4118 "Column '{}' not found for table '{}': {:?}",
4119 request.column, table_uri, e
4120 ),
4121 }
4122 } else {
4123 NamespaceError::Internal {
4124 message: format!(
4125 "Failed to create {} index '{}' on column '{}' for table '{}': {:?}",
4126 request.index_type,
4127 request.name.as_deref().unwrap_or("<auto-generated>"),
4128 request.column,
4129 table_uri,
4130 e
4131 ),
4132 }
4133 };
4134 lance_core::Error::from(ns_err)
4135 })?;
4136
4137 let transaction_id = dataset
4138 .read_transaction()
4139 .await
4140 .map_err(|e| {
4141 lance_core::Error::from(NamespaceError::Internal {
4142 message: format!(
4143 "Failed to read committed transaction after creating index on '{}': {}",
4144 table_uri, e
4145 ),
4146 })
4147 })?
4148 .map(|transaction| transaction.uuid);
4149
4150 Ok(CreateTableIndexResponse { transaction_id })
4151 }
4152
4153 async fn list_table_indices(
4154 &self,
4155 request: ListTableIndicesRequest,
4156 ) -> Result<ListTableIndicesResponse> {
4157 self.record_op("list_table_indices");
4158 let table_uri = self.resolve_table_location(&request.id).await?;
4159 let dataset = self
4160 .load_dataset(&table_uri, request.version, "list_table_indices")
4161 .await?;
4162 let total_rows = dataset.count_rows(None).await.map_err(|e| {
4163 lance_core::Error::from(NamespaceError::Internal {
4164 message: format!("Failed to count rows for table '{}': {:?}", table_uri, e),
4165 })
4166 })? as u64;
4167 let mut indices = dataset
4168 .describe_indices(None)
4169 .await
4170 .map_err(|e| {
4171 lance_core::Error::from(NamespaceError::Internal {
4172 message: format!("Failed to describe table indices for '{}': {:?}", table_uri, e),
4173 })
4174 })?
4175 .into_iter()
4176 .filter(|description| {
4177 description
4178 .metadata()
4179 .first()
4180 .map(|metadata| !is_system_index(metadata))
4181 .unwrap_or(false)
4182 })
4183 .map(|description| {
4184 let columns = description
4185 .field_ids()
4186 .iter()
4187 .map(|field_id| {
4188 dataset
4189 .schema()
4190 .field_path_minimal(i32::try_from(*field_id).map_err(|e| {
4191 lance_core::Error::from(NamespaceError::Internal {
4192 message: format!(
4193 "Field id {} does not fit in i32 for table '{}': {}",
4194 field_id, table_uri, e
4195 ),
4196 })
4197 })?)
4198 .map_err(|e| {
4199 lance_core::Error::from(NamespaceError::Internal {
4200 message: format!(
4201 "Failed to resolve field path for field_id {} in table '{}': {}",
4202 field_id, table_uri, e
4203 ),
4204 })
4205 })
4206 })
4207 .collect::<Result<Vec<_>>>()?;
4208
4209 let segments = description.segments();
4210 let created_at = segments
4211 .iter()
4212 .filter_map(|segment| segment.created_at)
4213 .min()
4214 .map(|ts| ts.to_rfc3339());
4215
4216 #[allow(clippy::needless_update)]
4219 let content = IndexContent {
4220 index_name: description.name().to_string(),
4221 index_uuid: description.metadata()[0].uuid.to_string(),
4222 columns,
4223 status: "SUCCEEDED".to_string(),
4224 index_type: Some(description.index_type().to_string()),
4225 type_url: Some(description.type_url().to_string()),
4226 num_indexed_rows: Some(description.rows_indexed() as i64),
4227 num_unindexed_rows: Some(
4228 total_rows.saturating_sub(description.rows_indexed()) as i64,
4229 ),
4230 size_bytes: description.total_size_bytes().map(|size| size as i64),
4231 num_segments: Some(segments.len() as i32),
4232 created_at,
4233 index_version: segments.first().map(|segment| segment.index_version),
4234 index_details: description.details().ok(),
4235 ..Default::default()
4236 };
4237 Ok(content)
4238 })
4239 .collect::<Result<Vec<_>>>()?;
4240
4241 let page_token = Self::paginate_indices(&mut indices, request.page_token, request.limit);
4242 Ok(ListTableIndicesResponse {
4243 indexes: indices,
4244 page_token,
4245 })
4246 }
4247
4248 async fn describe_table_index_stats(
4249 &self,
4250 request: DescribeTableIndexStatsRequest,
4251 ) -> Result<DescribeTableIndexStatsResponse> {
4252 self.record_op("describe_table_index_stats");
4253 let table_uri = self.resolve_table_location(&request.id).await?;
4254 let dataset = self
4255 .load_dataset(&table_uri, request.version, "describe_table_index_stats")
4256 .await?;
4257 let index_name = request.index_name.as_deref().ok_or_else(|| {
4258 lance_core::Error::from(NamespaceError::InvalidInput {
4259 message: "Index name is required for describe_table_index_stats".to_string(),
4260 })
4261 })?;
4262 let metadatas = dataset
4263 .load_indices_by_name(index_name)
4264 .await
4265 .map_err(|e| {
4266 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4267 message: format!(
4268 "Failed to load index '{}' metadata for table '{}': {}",
4269 index_name, table_uri, e
4270 ),
4271 })
4272 })?;
4273 if metadatas.first().is_some_and(is_system_index) {
4274 return Err(NamespaceError::Unsupported {
4275 message: format!("System index '{}' is not exposed by this API", index_name),
4276 }
4277 .into());
4278 }
4279
4280 let stats = <Dataset as DatasetIndexExt>::index_statistics(&dataset, index_name)
4281 .await
4282 .map_err(|e| {
4283 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4284 message: format!(
4285 "Failed to describe index statistics for '{}' on table '{}': {}",
4286 index_name, table_uri, e
4287 ),
4288 })
4289 })?;
4290 let stats: serde_json::Value = serde_json::from_str(&stats).map_err(|e| {
4291 lance_core::Error::from(NamespaceError::Internal {
4292 message: format!(
4293 "Failed to parse index statistics for '{}' on table '{}': {}",
4294 index_name, table_uri, e
4295 ),
4296 })
4297 })?;
4298
4299 Ok(Self::describe_table_index_stats_response(&stats))
4300 }
4301
4302 async fn describe_transaction(
4303 &self,
4304 request: DescribeTransactionRequest,
4305 ) -> Result<DescribeTransactionResponse> {
4306 self.record_op("describe_transaction");
4307 let mut request_id = request.id.ok_or_else(|| {
4308 lance_core::Error::from(NamespaceError::InvalidInput {
4309 message: "Transaction id must include table id and transaction identifier"
4310 .to_string(),
4311 })
4312 })?;
4313 if request_id.len() < 2 {
4314 return Err(NamespaceError::InvalidInput {
4315 message: format!(
4316 "Transaction request id must include table id and transaction identifier, got {:?}",
4317 request_id
4318 ),
4319 }
4320 .into());
4321 }
4322
4323 let id = request_id.pop().expect("request_id len checked above");
4324 let table_id = Some(request_id);
4325 let table_uri = self.resolve_table_location(&table_id).await?;
4326 let dataset = self
4327 .load_dataset(&table_uri, None, "describe_transaction")
4328 .await?;
4329 let (version, transaction) = self.find_transaction(&dataset, &id).await?;
4330
4331 let sidecar = self
4334 .load_transaction_alteration(&table_uri, &transaction.uuid)
4335 .await?;
4336
4337 Ok(Self::transaction_response(version, &transaction, sidecar))
4338 }
4339
4340 async fn alter_transaction(
4341 &self,
4342 request: AlterTransactionRequest,
4343 ) -> Result<AlterTransactionResponse> {
4344 self.record_op("alter_transaction");
4345
4346 let mut request_id = request.id.ok_or_else(|| {
4348 lance_core::Error::from(NamespaceError::InvalidInput {
4349 message: "Transaction id must include table id and transaction identifier"
4350 .to_string(),
4351 })
4352 })?;
4353 if request_id.len() < 2 {
4354 return Err(NamespaceError::InvalidInput {
4355 message: format!(
4356 "Transaction request id must include table id and transaction identifier, got {:?}",
4357 request_id
4358 ),
4359 }
4360 .into());
4361 }
4362
4363 let txn_id = request_id.pop().expect("request_id len checked above");
4364 let table_id = Some(request_id);
4365 let table_uri = self.resolve_table_location(&table_id).await?;
4366 let dataset = self
4367 .load_dataset(&table_uri, None, "alter_transaction")
4368 .await?;
4369 let (version, transaction) = self.find_transaction(&dataset, &txn_id).await?;
4370
4371 const RESERVED_KEYS: &[&str] = &["uuid", "version", "read_version", "operation", "tag"];
4375 let is_reserved = |key: &str| RESERVED_KEYS.contains(&key);
4376
4377 let mut sidecar = self
4380 .load_transaction_alteration(&table_uri, &transaction.uuid)
4381 .await?
4382 .unwrap_or_default();
4383
4384 for action in &request.actions {
4385 if let Some(ref set_status) = action.set_status_action
4386 && let Some(ref status) = set_status.status
4387 {
4388 let normalized = status.to_lowercase().replace('_', "");
4390 match normalized.as_str() {
4391 "queued" | "running" | "succeeded" | "failed" | "canceled" => {
4392 sidecar.status = Some(status.clone());
4393 }
4394 _ => {
4395 return Err(NamespaceError::InvalidInput {
4396 message: format!(
4397 "Invalid transaction status '{}'. Valid values are: Queued, Running, Succeeded, Failed, Canceled",
4398 status
4399 ),
4400 }
4401 .into());
4402 }
4403 }
4404 }
4405
4406 if let Some(ref set_property) = action.set_property_action
4407 && let (Some(key), Some(value)) = (&set_property.key, &set_property.value)
4408 {
4409 if is_reserved(key) {
4410 return Err(NamespaceError::InvalidInput {
4411 message: format!("Property '{}' is reserved and cannot be modified", key),
4412 }
4413 .into());
4414 }
4415 let mode = set_property
4416 .mode
4417 .as_deref()
4418 .unwrap_or("Overwrite")
4419 .to_lowercase();
4420 match mode.as_str() {
4421 "overwrite" => {
4422 sidecar.properties.insert(key.clone(), value.clone());
4423 }
4424 "fail" => {
4425 let exists = sidecar.properties.contains_key(key)
4428 || transaction
4429 .transaction_properties
4430 .as_ref()
4431 .is_some_and(|props| props.contains_key(key));
4432 if exists {
4433 return Err(NamespaceError::ConcurrentModification {
4434 message: format!(
4435 "Property '{}' already exists and mode is 'Fail'",
4436 key
4437 ),
4438 }
4439 .into());
4440 }
4441 sidecar.properties.insert(key.clone(), value.clone());
4442 }
4443 "skip" => {
4444 let exists = sidecar.properties.contains_key(key)
4445 || transaction
4446 .transaction_properties
4447 .as_ref()
4448 .is_some_and(|props| props.contains_key(key));
4449 if !exists {
4450 sidecar.properties.insert(key.clone(), value.clone());
4451 }
4452 }
4453 _ => {
4454 return Err(NamespaceError::InvalidInput {
4455 message: format!(
4456 "Invalid set_property mode '{}'. Valid values are: Overwrite, Fail, Skip",
4457 mode
4458 ),
4459 }
4460 .into());
4461 }
4462 }
4463 }
4464
4465 if let Some(ref unset_property) = action.unset_property_action
4466 && let Some(ref key) = unset_property.key
4467 {
4468 if is_reserved(key) {
4469 return Err(NamespaceError::InvalidInput {
4470 message: format!("Property '{}' is reserved and cannot be modified", key),
4471 }
4472 .into());
4473 }
4474 let mode = unset_property
4475 .mode
4476 .as_deref()
4477 .unwrap_or("Skip")
4478 .to_lowercase();
4479 let exists_in_transaction = transaction
4480 .transaction_properties
4481 .as_ref()
4482 .is_some_and(|props| props.contains_key(key));
4483 match mode.as_str() {
4484 "skip" => {
4485 sidecar.properties.remove(key);
4486 if exists_in_transaction {
4487 sidecar.removed_properties.insert(key.clone());
4490 }
4491 }
4492 "fail" => {
4493 if !sidecar.properties.contains_key(key) && !exists_in_transaction {
4494 return Err(NamespaceError::InvalidInput {
4495 message: format!(
4496 "Property '{}' does not exist and mode is 'Fail'",
4497 key
4498 ),
4499 }
4500 .into());
4501 }
4502 sidecar.properties.remove(key);
4503 if exists_in_transaction {
4504 sidecar.removed_properties.insert(key.clone());
4505 }
4506 }
4507 _ => {
4508 return Err(NamespaceError::InvalidInput {
4509 message: format!(
4510 "Invalid unset_property mode '{}'. Valid values are: Skip, Fail",
4511 mode
4512 ),
4513 }
4514 .into());
4515 }
4516 }
4517 }
4518 }
4519
4520 self.save_transaction_alteration(&table_uri, &transaction.uuid, &sidecar)
4524 .await?;
4525
4526 let final_status = sidecar
4529 .status
4530 .clone()
4531 .unwrap_or_else(|| "SUCCEEDED".to_string());
4532 let response = Self::transaction_response(version, &transaction, Some(sidecar));
4533 Ok(AlterTransactionResponse {
4534 status: final_status,
4535 properties: response.properties,
4536 })
4537 }
4538
4539 async fn create_table_scalar_index(
4540 &self,
4541 request: CreateTableIndexRequest,
4542 ) -> Result<CreateTableScalarIndexResponse> {
4543 self.record_op("create_table_scalar_index");
4544 let index_type = Self::parse_index_type(&request.index_type)?;
4545 if !index_type.is_scalar() {
4546 return Err(NamespaceError::InvalidInput {
4547 message: format!(
4548 "create_table_scalar_index only supports scalar index types, got {}",
4549 request.index_type
4550 ),
4551 }
4552 .into());
4553 }
4554
4555 let response = self.create_table_index(request).await?;
4556 Ok(CreateTableScalarIndexResponse {
4557 transaction_id: response.transaction_id,
4558 })
4559 }
4560
4561 async fn drop_table_index(
4562 &self,
4563 request: DropTableIndexRequest,
4564 ) -> Result<DropTableIndexResponse> {
4565 self.record_op("drop_table_index");
4566 let table_uri = self.resolve_table_location(&request.id).await?;
4567 let index_name = request.index_name.as_deref().ok_or_else(|| {
4568 lance_core::Error::from(NamespaceError::InvalidInput {
4569 message: "Index name is required for drop_table_index".to_string(),
4570 })
4571 })?;
4572 let mut dataset = self
4573 .load_dataset(&table_uri, None, "drop_table_index")
4574 .await?;
4575 let metadatas = dataset
4576 .load_indices_by_name(index_name)
4577 .await
4578 .map_err(|e| {
4579 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4580 message: format!(
4581 "Failed to load index '{}' before dropping it from table '{}': {}",
4582 index_name, table_uri, e
4583 ),
4584 })
4585 })?;
4586 if metadatas.first().is_some_and(is_system_index) {
4587 return Err(NamespaceError::Unsupported {
4588 message: format!(
4589 "System index '{}' cannot be dropped via this API",
4590 index_name
4591 ),
4592 }
4593 .into());
4594 }
4595
4596 dataset.drop_index(index_name).await.map_err(|e| {
4597 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4598 message: format!(
4599 "Failed to drop index '{}' from table '{}': {}",
4600 index_name, table_uri, e
4601 ),
4602 })
4603 })?;
4604
4605 let transaction_id = dataset
4606 .read_transaction()
4607 .await
4608 .map_err(|e| {
4609 lance_core::Error::from(NamespaceError::Internal {
4610 message: format!(
4611 "Failed to read committed transaction after dropping index '{}' from '{}': {}",
4612 index_name, table_uri, e
4613 ),
4614 })
4615 })?
4616 .map(|transaction| transaction.uuid);
4617
4618 Ok(DropTableIndexResponse { transaction_id })
4619 }
4620
4621 async fn list_all_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
4622 let mut tables = self.list_directory_tables().await?;
4625 tables = self
4626 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
4627 .await?;
4628 Self::apply_pagination(&mut tables, request.page_token, request.limit);
4629 Ok(ListTablesResponse::new(tables))
4630 }
4631
4632 async fn restore_table(&self, request: RestoreTableRequest) -> Result<RestoreTableResponse> {
4633 let version = request.version;
4634 if version < 0 {
4635 return Err(Error::invalid_input_source(
4636 format!(
4637 "Table version for restore_table must be non-negative, got {}",
4638 version
4639 )
4640 .into(),
4641 ));
4642 }
4643
4644 let branch = Self::normalized_branch(request.branch.as_deref())?;
4645 let table_uri = self.resolve_table_location(&request.id).await?;
4646 let mut dataset = match branch {
4647 Some(branch) => self.open_validated_branch(&table_uri, branch).await?,
4648 None => self.load_dataset(&table_uri, None, "restore_table").await?,
4649 };
4650
4651 dataset = dataset
4652 .checkout_version(version as u64)
4653 .await
4654 .map_err(|e| {
4655 Error::namespace_source(
4656 format!(
4657 "Failed to checkout version {} for restore at '{}': {}",
4658 version, table_uri, e
4659 )
4660 .into(),
4661 )
4662 })?;
4663
4664 dataset.restore().await.map_err(|e| {
4665 Error::namespace_source(
4666 format!(
4667 "Failed to restore table at '{}' to version {}: {}",
4668 table_uri, version, e
4669 )
4670 .into(),
4671 )
4672 })?;
4673
4674 let transaction_id = dataset
4675 .read_transaction()
4676 .await
4677 .map_err(|e| {
4678 Error::namespace_source(
4679 format!(
4680 "Failed to read transaction after restoring '{}': {}",
4681 table_uri, e
4682 )
4683 .into(),
4684 )
4685 })?
4686 .map(|t| t.uuid);
4687
4688 Ok(RestoreTableResponse { transaction_id })
4689 }
4690
4691 async fn update_table_schema_metadata(
4692 &self,
4693 request: UpdateTableSchemaMetadataRequest,
4694 ) -> Result<UpdateTableSchemaMetadataResponse> {
4695 let table_uri = self.resolve_table_location(&request.id).await?;
4696 let mut dataset = self
4697 .load_dataset(&table_uri, None, "update_table_schema_metadata")
4698 .await?;
4699
4700 let new_metadata = request.metadata.unwrap_or_default();
4701 let updated_metadata = dataset
4702 .update_schema_metadata(new_metadata.iter().map(|(k, v)| (k.as_str(), v.as_str())))
4703 .await
4704 .map_err(|e| {
4705 Error::namespace_source(
4706 format!(
4707 "Failed to update schema metadata for table at '{}': {}",
4708 table_uri, e
4709 )
4710 .into(),
4711 )
4712 })?;
4713
4714 let transaction_id = dataset
4715 .read_transaction()
4716 .await
4717 .map_err(|e| {
4718 Error::namespace_source(
4719 format!(
4720 "Failed to read transaction after updating metadata for '{}': {}",
4721 table_uri, e
4722 )
4723 .into(),
4724 )
4725 })?
4726 .map(|t| t.uuid);
4727
4728 Ok(UpdateTableSchemaMetadataResponse {
4729 metadata: Some(updated_metadata),
4730 transaction_id,
4731 })
4732 }
4733
4734 async fn get_table_stats(
4735 &self,
4736 request: GetTableStatsRequest,
4737 ) -> Result<GetTableStatsResponse> {
4738 let table_uri = self.resolve_table_location(&request.id).await?;
4739 let dataset = Arc::new(
4740 self.load_dataset(&table_uri, None, "get_table_stats")
4741 .await?,
4742 );
4743
4744 let data_stats = dataset.calculate_data_stats().await.map_err(|e| {
4746 Error::namespace_source(
4747 format!(
4748 "Failed to calculate data statistics for table at '{}': {}",
4749 table_uri, e
4750 )
4751 .into(),
4752 )
4753 })?;
4754 let total_bytes: i64 = data_stats
4755 .fields
4756 .iter()
4757 .map(|f| f.bytes_on_disk as i64)
4758 .sum();
4759
4760 let fragment_row_futures: Vec<_> = dataset
4762 .get_fragments()
4763 .into_iter()
4764 .map(|f| async move { f.physical_rows().await })
4765 .collect();
4766 let fragment_row_results = futures::future::join_all(fragment_row_futures).await;
4767 let mut fragment_row_counts: Vec<i64> = fragment_row_results
4768 .into_iter()
4769 .filter_map(|r| r.ok())
4770 .map(|r| r as i64)
4771 .collect();
4772
4773 let num_fragments = fragment_row_counts.len() as i64;
4774 let num_rows: i64 = fragment_row_counts.iter().sum();
4775
4776 const SMALL_FRAGMENT_THRESHOLD: i64 = 1024 * 1024;
4779 let num_small_fragments = fragment_row_counts
4780 .iter()
4781 .filter(|&&r| r < SMALL_FRAGMENT_THRESHOLD)
4782 .count() as i64;
4783
4784 fragment_row_counts.sort_unstable();
4786 let lengths = if fragment_row_counts.is_empty() {
4787 FragmentSummary::new(0, 0, 0, 0, 0, 0, 0)
4788 } else {
4789 let len = fragment_row_counts.len();
4790 let min = fragment_row_counts[0];
4791 let max = fragment_row_counts[len - 1];
4792 let mean = num_rows / num_fragments;
4793 let pct = |p: f64| fragment_row_counts[((len - 1) as f64 * p) as usize];
4794 FragmentSummary::new(min, max, mean, pct(0.25), pct(0.50), pct(0.75), pct(0.99))
4795 };
4796
4797 let indices = dataset.load_indices().await.map_err(|e| {
4799 Error::namespace_source(
4800 format!("Failed to load indices for table at '{}': {}", table_uri, e).into(),
4801 )
4802 })?;
4803 let num_indices = indices.iter().filter(|m| !is_system_index(m)).count() as i64;
4804
4805 let fragment_stats = FragmentStats::new(num_fragments, num_small_fragments, lengths);
4806 Ok(GetTableStatsResponse::new(
4807 total_bytes,
4808 num_rows,
4809 num_indices,
4810 fragment_stats,
4811 ))
4812 }
4813
4814 async fn explain_table_query_plan(
4815 &self,
4816 request: ExplainTableQueryPlanRequest,
4817 ) -> Result<String> {
4818 let table_uri = self.resolve_table_location(&request.id).await?;
4819 let dataset = self
4820 .load_dataset(
4821 &table_uri,
4822 request.query.version,
4823 "explain_table_query_plan",
4824 )
4825 .await?;
4826 let verbose = request.verbose.unwrap_or(false);
4827
4828 let mut scanner = dataset.scan();
4829 Self::apply_query_params_to_scanner(
4830 &mut scanner,
4831 request.query.filter.as_deref(),
4832 request.query.columns.as_deref(),
4833 request.query.vector_column.as_deref(),
4834 &request.query.vector,
4835 request.query.k,
4836 request.query.offset,
4837 request.query.prefilter,
4838 request.query.bypass_vector_index,
4839 request.query.nprobes,
4840 request.query.ef,
4841 request.query.refine_factor,
4842 request.query.distance_type.as_deref(),
4843 request.query.fast_search,
4844 request.query.with_row_id,
4845 request.query.lower_bound,
4846 request.query.upper_bound,
4847 "explain_table_query_plan",
4848 )?;
4849
4850 scanner.explain_plan(verbose).await.map_err(|e| {
4851 Error::namespace_source(
4852 format!(
4853 "Failed to explain query plan for table at '{}': {}",
4854 table_uri, e
4855 )
4856 .into(),
4857 )
4858 })
4859 }
4860
4861 async fn analyze_table_query_plan(
4862 &self,
4863 request: AnalyzeTableQueryPlanRequest,
4864 ) -> Result<String> {
4865 let table_uri = self.resolve_table_location(&request.id).await?;
4866 let dataset = self
4867 .load_dataset(&table_uri, request.version, "analyze_table_query_plan")
4868 .await?;
4869
4870 let mut scanner = dataset.scan();
4871 Self::apply_query_params_to_scanner(
4872 &mut scanner,
4873 request.filter.as_deref(),
4874 request.columns.as_deref(),
4875 request.vector_column.as_deref(),
4876 &request.vector,
4877 request.k,
4878 request.offset,
4879 request.prefilter,
4880 request.bypass_vector_index,
4881 request.nprobes,
4882 request.ef,
4883 request.refine_factor,
4884 request.distance_type.as_deref(),
4885 request.fast_search,
4886 request.with_row_id,
4887 request.lower_bound,
4888 request.upper_bound,
4889 "analyze_table_query_plan",
4890 )?;
4891
4892 scanner.analyze_plan().await.map_err(|e| {
4893 Error::namespace_source(
4894 format!(
4895 "Failed to analyze query plan for table at '{}': {}",
4896 table_uri, e
4897 )
4898 .into(),
4899 )
4900 })
4901 }
4902
4903 async fn count_table_rows(&self, request: CountTableRowsRequest) -> Result<i64> {
4904 self.record_op("count_table_rows");
4905 let table_uri = self.resolve_table_location(&request.id).await?;
4906 let dataset = self
4907 .load_dataset(&table_uri, request.version, "count_table_rows")
4908 .await?;
4909
4910 let count =
4911 dataset
4912 .count_rows(request.predicate)
4913 .await
4914 .map_err(|e| NamespaceError::Internal {
4915 message: format!("Failed to count rows for table at '{}': {:?}", table_uri, e),
4916 })?;
4917
4918 Ok(count as i64)
4919 }
4920
4921 async fn insert_into_table(
4922 &self,
4923 request: InsertIntoTableRequest,
4924 request_data: Bytes,
4925 ) -> Result<InsertIntoTableResponse> {
4926 self.record_op("insert_into_table");
4927 let table_uri = self.resolve_table_location(&request.id).await?;
4928 let (reader, _num_rows) =
4929 Self::ipc_reader_from_request_data(&request_data, "insert_into_table")?;
4930
4931 let mode = match request.mode.as_deref() {
4932 Some(m) if m.eq_ignore_ascii_case("overwrite") => WriteMode::Overwrite,
4933 Some(m) if m.eq_ignore_ascii_case("append") => WriteMode::Append,
4934 None => WriteMode::Append,
4935 Some(m) => {
4936 return Err(lance_namespace::error::NamespaceError::InvalidInput {
4937 message: format!(
4938 "Unsupported write mode '{}'. Supported modes are: 'append', 'overwrite'",
4939 m
4940 ),
4941 }
4942 .into());
4943 }
4944 };
4945
4946 if !self.table_uri_has_actual_manifests(&table_uri).await? {
4947 self.write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
4948 .await?;
4949 } else {
4950 self.write_reader_to_table(&table_uri, reader, mode, None)
4951 .await?;
4952 }
4953
4954 Ok(InsertIntoTableResponse {
4955 transaction_id: None,
4956 })
4957 }
4958
4959 async fn merge_insert_into_table(
4960 &self,
4961 request: MergeInsertIntoTableRequest,
4962 request_data: Bytes,
4963 ) -> Result<MergeInsertIntoTableResponse> {
4964 self.record_op("merge_insert_into_table");
4965 let table_uri = self.resolve_table_location(&request.id).await?;
4966 let on = request.on.as_ref().ok_or_else(|| {
4967 lance_core::Error::from(NamespaceError::InvalidInput {
4968 message: "'on' field is required for merge_insert_into_table".to_string(),
4969 })
4970 })?;
4971
4972 let table_has_manifests = self.table_uri_has_actual_manifests(&table_uri).await?;
4973 let (reader, num_rows) =
4974 Self::ipc_reader_from_request_data(&request_data, "merge_insert_into_table")?;
4975
4976 if !table_has_manifests {
4977 let dataset = self
4978 .write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
4979 .await?;
4980 let version = dataset.version().version as i64;
4981 return Ok(MergeInsertIntoTableResponse {
4982 transaction_id: None,
4983 num_updated_rows: Some(0),
4984 num_inserted_rows: Some(num_rows as i64),
4985 num_deleted_rows: Some(0),
4986 version: Some(version),
4987 });
4988 }
4989
4990 let dataset = Arc::new(
4991 self.load_dataset(&table_uri, None, "merge_insert_into_table")
4992 .await?,
4993 );
4994
4995 let mut merge_builder = MergeInsertBuilder::try_new(dataset.clone(), vec![on.clone()])
4996 .map_err(|e| {
4997 lance_core::Error::from(NamespaceError::InvalidInput {
4998 message: format!("Failed to create merge_insert_into_table builder: {}", e),
4999 })
5000 })?;
5001
5002 if let Some(filter) = request.when_matched_update_all_filt.as_deref() {
5003 let behavior = WhenMatched::update_if(dataset.as_ref(), filter).map_err(|e| {
5004 lance_core::Error::from(NamespaceError::InvalidInput {
5005 message: format!(
5006 "Invalid when_matched_update_all_filt for merge_insert_into_table: {}",
5007 e
5008 ),
5009 })
5010 })?;
5011 merge_builder.when_matched(behavior);
5012 } else if request.when_matched_update_all.unwrap_or(false) {
5013 merge_builder.when_matched(WhenMatched::UpdateAll);
5014 }
5015
5016 if matches!(request.when_not_matched_insert_all, Some(false)) {
5017 merge_builder.when_not_matched(WhenNotMatched::DoNothing);
5018 } else {
5019 merge_builder.when_not_matched(WhenNotMatched::InsertAll);
5020 }
5021
5022 if let Some(filter) = request.when_not_matched_by_source_delete_filt.as_deref() {
5023 let behavior = WhenNotMatchedBySource::delete_if(dataset.as_ref(), filter).map_err(|e| {
5024 lance_core::Error::from(NamespaceError::InvalidInput {
5025 message: format!(
5026 "Invalid when_not_matched_by_source_delete_filt for merge_insert_into_table: {}",
5027 e
5028 ),
5029 })
5030 })?;
5031 merge_builder.when_not_matched_by_source(behavior);
5032 } else if request.when_not_matched_by_source_delete.unwrap_or(false) {
5033 merge_builder.when_not_matched_by_source(WhenNotMatchedBySource::Delete);
5034 }
5035
5036 if let Some(use_index) = request.use_index {
5037 merge_builder.use_index(use_index);
5038 }
5039
5040 let (dataset, stats) = merge_builder
5041 .try_build()
5042 .map_err(|e| {
5043 lance_core::Error::from(NamespaceError::InvalidInput {
5044 message: format!("Failed to build merge_insert_into_table job: {}", e),
5045 })
5046 })?
5047 .execute_reader(reader)
5048 .await
5049 .map_err(|e| Self::map_mutation_error(e, "merge_insert_into_table", &table_uri))?;
5050
5051 Ok(MergeInsertIntoTableResponse {
5052 transaction_id: None,
5053 num_updated_rows: Some(stats.num_updated_rows as i64),
5054 num_inserted_rows: Some(stats.num_inserted_rows as i64),
5055 num_deleted_rows: Some(stats.num_deleted_rows as i64),
5056 version: Some(dataset.version().version as i64),
5057 })
5058 }
5059
5060 async fn update_table(&self, request: UpdateTableRequest) -> Result<UpdateTableResponse> {
5061 self.record_op("update_table");
5062
5063 if request.updates.is_empty() {
5064 return Err(NamespaceError::InvalidInput {
5065 message: "update_table requires at least one [column, expression] pair".to_string(),
5066 }
5067 .into());
5068 }
5069
5070 let mut seen_columns: HashMap<String, ()> = HashMap::with_capacity(request.updates.len());
5073 for (idx, pair) in request.updates.iter().enumerate() {
5074 if pair.len() != 2 {
5075 return Err(NamespaceError::InvalidInput {
5076 message: format!(
5077 "update_table updates[{}] must be a [column, expression] pair, got {} elements",
5078 idx,
5079 pair.len()
5080 ),
5081 }
5082 .into());
5083 }
5084 let column = &pair[0];
5085 if column.trim().is_empty() {
5086 return Err(NamespaceError::InvalidInput {
5087 message: format!("update_table updates[{}] has an empty column name", idx),
5088 }
5089 .into());
5090 }
5091 if seen_columns.insert(column.clone(), ()).is_some() {
5092 return Err(NamespaceError::InvalidInput {
5093 message: format!(
5094 "update_table cannot update column '{}' more than once",
5095 column
5096 ),
5097 }
5098 .into());
5099 }
5100 }
5101
5102 let table_uri = self.resolve_table_location(&request.id).await?;
5103 let dataset = Arc::new(self.load_dataset(&table_uri, None, "update_table").await?);
5104
5105 let mut builder = UpdateBuilder::new(dataset);
5106 for pair in &request.updates {
5107 builder = builder.set(&pair[0], &pair[1]).map_err(|e| {
5109 lance_core::Error::from(NamespaceError::InvalidInput {
5110 message: format!("Invalid update expression for column '{}': {}", pair[0], e),
5111 })
5112 })?;
5113 }
5114 if let Some(predicate) = request.predicate.as_deref()
5115 && !predicate.trim().is_empty()
5116 {
5117 builder = builder.update_where(predicate).map_err(|e| {
5118 lance_core::Error::from(NamespaceError::InvalidInput {
5119 message: format!("Invalid update_table predicate '{}': {}", predicate, e),
5120 })
5121 })?;
5122 }
5123
5124 let job = builder.build().map_err(|e| {
5125 lance_core::Error::from(NamespaceError::InvalidInput {
5126 message: format!("Failed to build update_table job: {}", e),
5127 })
5128 })?;
5129
5130 let result = job
5131 .execute()
5132 .await
5133 .map_err(|e| Self::map_mutation_error(e, "update_table", &table_uri))?;
5134
5135 let version = result.new_dataset.version().version as i64;
5136 Ok(UpdateTableResponse {
5137 transaction_id: None,
5138 updated_rows: result.rows_updated as i64,
5139 version,
5140 properties: None,
5141 })
5142 }
5143
5144 async fn delete_from_table(
5145 &self,
5146 request: DeleteFromTableRequest,
5147 ) -> Result<DeleteFromTableResponse> {
5148 self.record_op("delete_from_table");
5149
5150 if request.predicate.trim().is_empty() {
5151 return Err(NamespaceError::InvalidInput {
5152 message: "delete_from_table requires a non-empty predicate".to_string(),
5153 }
5154 .into());
5155 }
5156
5157 let table_uri = self.resolve_table_location(&request.id).await?;
5158 let mut dataset = self
5159 .load_dataset(&table_uri, None, "delete_from_table")
5160 .await?;
5161
5162 let result = dataset
5163 .delete(&request.predicate)
5164 .await
5165 .map_err(|e| Self::map_mutation_error(e, "delete_from_table", &table_uri))?;
5166
5167 Ok(DeleteFromTableResponse {
5168 transaction_id: None,
5169 version: Some(result.new_dataset.version().version as i64),
5170 })
5171 }
5172
5173 async fn query_table(&self, request: QueryTableRequest) -> Result<Bytes> {
5174 use arrow::ipc::writer::FileWriter;
5175
5176 self.record_op("query_table");
5177 let table_uri = self.resolve_table_location(&request.id).await?;
5178 let dataset = self
5179 .load_dataset(&table_uri, request.version, "query_table")
5180 .await?;
5181
5182 let mut scanner = dataset.scan();
5184
5185 let has_vector_query = request
5188 .vector
5189 .single_vector
5190 .as_ref()
5191 .map(|sv| !sv.is_empty())
5192 .unwrap_or(false)
5193 || request
5194 .vector
5195 .multi_vector
5196 .as_ref()
5197 .map(|mv| !mv.is_empty())
5198 .unwrap_or(false);
5199
5200 if let Some(prefilter) = request.prefilter {
5202 scanner.prefilter(prefilter);
5203 }
5204
5205 if has_vector_query {
5207 let vector_column = request.vector_column.as_deref().unwrap_or("vector");
5208
5209 let query_vector: Vec<f32> = request
5211 .vector
5212 .single_vector
5213 .clone()
5214 .or_else(|| {
5215 request
5216 .vector
5217 .multi_vector
5218 .as_ref()
5219 .and_then(|mv| mv.first().cloned())
5220 })
5221 .unwrap_or_default();
5222
5223 if !query_vector.is_empty() {
5224 let k = if request.k > 0 {
5225 request.k as usize
5226 } else {
5227 10
5228 };
5229 let query_array = Float32Array::from(query_vector);
5230 scanner
5231 .nearest(vector_column, &query_array, k)
5232 .map_err(|e| NamespaceError::InvalidInput {
5233 message: format!("Invalid vector search: {:?}", e),
5234 })?;
5235
5236 if let Some(ref distance_type) = request.distance_type {
5238 let metric = match distance_type.to_lowercase().as_str() {
5239 "l2" | "euclidean" => MetricType::L2,
5240 "cosine" => MetricType::Cosine,
5241 "dot" | "inner_product" => MetricType::Dot,
5242 "hamming" => MetricType::Hamming,
5243 _ => {
5244 return Err(NamespaceError::InvalidInput {
5245 message: format!("Unknown distance type: {}", distance_type),
5246 }
5247 .into());
5248 }
5249 };
5250 scanner.distance_metric(metric);
5251 }
5252
5253 if let Some(nprobes) = request.nprobes {
5255 scanner.minimum_nprobes(nprobes as usize);
5256 }
5257
5258 if let Some(ef) = request.ef {
5260 scanner.ef(ef as usize);
5261 }
5262
5263 if let Some(refine_factor) = request.refine_factor {
5265 scanner.refine(refine_factor as u32);
5266 }
5267
5268 if request.lower_bound.is_some() || request.upper_bound.is_some() {
5270 scanner.distance_range(request.lower_bound, request.upper_bound);
5271 }
5272
5273 if let Some(bypass) = request.bypass_vector_index {
5275 scanner.use_index(!bypass);
5276 }
5277
5278 if request.fast_search == Some(true) {
5280 scanner.fast_search();
5281 }
5282 }
5283 }
5284
5285 if let Some(ref fts_query) = request.full_text_query {
5287 if let Some(ref string_query) = fts_query.string_query {
5289 let mut fts = FullTextSearchQuery::new(string_query.query.clone());
5290
5291 if let Some(ref columns) = string_query.columns
5293 && !columns.is_empty()
5294 {
5295 fts = fts
5296 .with_columns(columns)
5297 .map_err(|e| NamespaceError::InvalidInput {
5298 message: format!("Invalid FTS columns: {:?}", e),
5299 })?;
5300 }
5301
5302 scanner
5303 .full_text_search(fts)
5304 .map_err(|e| NamespaceError::InvalidInput {
5305 message: format!("Invalid full text search: {:?}", e),
5306 })?;
5307 } else if let Some(ref structured_query) = fts_query.structured_query {
5308 let engine_query = build_engine_fts_query(&structured_query.query)?;
5310 let fts = FullTextSearchQuery::new_query(engine_query);
5311 scanner
5312 .full_text_search(fts)
5313 .map_err(|e| NamespaceError::InvalidInput {
5314 message: format!("Invalid full text search: {:?}", e),
5315 })?;
5316 }
5317 }
5318
5319 if let Some(ref columns) = request.columns {
5321 if let Some(ref column_names) = columns.column_names
5322 && !column_names.is_empty()
5323 {
5324 scanner
5325 .project(column_names)
5326 .map_err(|e| NamespaceError::InvalidInput {
5327 message: format!("Invalid column projection: {:?}", e),
5328 })?;
5329 } else if let Some(ref column_aliases) = columns.column_aliases
5330 && !column_aliases.is_empty()
5331 {
5332 let transform_pairs: Vec<(String, String)> = column_aliases
5334 .iter()
5335 .map(|(alias, sql)| (alias.clone(), sql.clone()))
5336 .collect();
5337 scanner
5338 .project_with_transform(
5339 &transform_pairs
5340 .iter()
5341 .map(|(a, s)| (a.as_str(), s.as_str()))
5342 .collect::<Vec<_>>(),
5343 )
5344 .map_err(|e| NamespaceError::InvalidInput {
5345 message: format!("Invalid column alias expression: {:?}", e),
5346 })?;
5347 }
5348 }
5349
5350 if let Some(ref filter) = request.filter
5352 && !filter.is_empty()
5353 {
5354 scanner
5355 .filter(filter)
5356 .map_err(|e| NamespaceError::InvalidInput {
5357 message: format!("Invalid filter expression: {:?}", e),
5358 })?;
5359 }
5360
5361 if request.with_row_id == Some(true) {
5363 scanner.with_row_id();
5364 }
5365
5366 if !has_vector_query && request.k > 0 {
5370 let offset = request.offset.map(|o| o as i64);
5371 scanner.limit(Some(request.k as i64), offset).map_err(|e| {
5372 NamespaceError::InvalidInput {
5373 message: format!("Invalid limit/offset: {:?}", e),
5374 }
5375 })?;
5376 } else if has_vector_query && request.offset.is_some() {
5377 let offset = request.offset.map(|o| o as i64);
5379 scanner
5380 .limit(None, offset)
5381 .map_err(|e| NamespaceError::InvalidInput {
5382 message: format!("Invalid offset: {:?}", e),
5383 })?;
5384 }
5385
5386 let batch = scanner
5388 .try_into_batch()
5389 .await
5390 .map_err(|e| NamespaceError::Internal {
5391 message: format!("Failed to execute query: {:?}", e),
5392 })?;
5393
5394 let schema = batch.schema();
5396 let mut buffer = Vec::new();
5397 {
5398 let mut writer = FileWriter::try_new(&mut buffer, &schema).map_err(|e| {
5399 NamespaceError::Internal {
5400 message: format!("Failed to create IPC writer: {:?}", e),
5401 }
5402 })?;
5403 writer.write(&batch).map_err(|e| NamespaceError::Internal {
5404 message: format!("Failed to write batch to IPC: {:?}", e),
5405 })?;
5406 writer.finish().map_err(|e| NamespaceError::Internal {
5407 message: format!("Failed to finish IPC writer: {:?}", e),
5408 })?;
5409 }
5410
5411 Ok(Bytes::from(buffer))
5412 }
5413
5414 async fn list_table_tags(
5415 &self,
5416 request: ListTableTagsRequest,
5417 ) -> Result<ListTableTagsResponse> {
5418 self.record_op("list_table_tags");
5419 let table_uri = self.resolve_table_location(&request.id).await?;
5420 let dataset = self
5421 .load_dataset(&table_uri, None, "list_table_tags")
5422 .await?;
5423
5424 let raw_tags = dataset.tags().list().await.map_err(|e| {
5425 lance_core::Error::from(NamespaceError::Internal {
5426 message: format!("Failed to list tags for table at '{}': {}", table_uri, e),
5427 })
5428 })?;
5429
5430 let tags = raw_tags
5431 .into_iter()
5432 .map(|(name, contents)| {
5433 let mut tag_model =
5434 ModelTagContents::new(contents.version as i64, contents.manifest_size as i64);
5435 tag_model.branch = contents.branch;
5436 (name, tag_model)
5437 })
5438 .collect();
5439
5440 Ok(ListTableTagsResponse {
5441 tags,
5442 page_token: None,
5443 })
5444 }
5445
5446 async fn get_table_tag_version(
5447 &self,
5448 request: GetTableTagVersionRequest,
5449 ) -> Result<GetTableTagVersionResponse> {
5450 self.record_op("get_table_tag_version");
5451 if request.tag.is_empty() {
5452 return Err(NamespaceError::InvalidInput {
5453 message: "tag name must not be empty for get_table_tag_version".to_string(),
5454 }
5455 .into());
5456 }
5457
5458 let table_uri = self.resolve_table_location(&request.id).await?;
5459 let dataset = self
5460 .load_dataset(&table_uri, None, "get_table_tag_version")
5461 .await?;
5462
5463 let contents = dataset
5464 .tags()
5465 .get(&request.tag)
5466 .await
5467 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5468
5469 Ok(GetTableTagVersionResponse {
5470 version: contents.version as i64,
5471 branch: contents.branch,
5472 })
5473 }
5474
5475 async fn create_table_tag(
5476 &self,
5477 request: CreateTableTagRequest,
5478 ) -> Result<CreateTableTagResponse> {
5479 self.record_op("create_table_tag");
5480 if request.tag.is_empty() {
5481 return Err(NamespaceError::InvalidInput {
5482 message: "tag name must not be empty for create_table_tag".to_string(),
5483 }
5484 .into());
5485 }
5486 if request.version <= 0 {
5487 return Err(NamespaceError::InvalidInput {
5488 message: format!(
5489 "tag version must be a positive integer, got {} for create_table_tag",
5490 request.version
5491 ),
5492 }
5493 .into());
5494 }
5495
5496 let table_uri = self.resolve_table_location(&request.id).await?;
5497 let dataset = self
5498 .load_dataset(&table_uri, None, "create_table_tag")
5499 .await?;
5500
5501 dataset
5502 .tags()
5503 .create(&request.tag, request.version as u64)
5504 .await
5505 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5506
5507 Ok(CreateTableTagResponse {
5508 transaction_id: None,
5509 })
5510 }
5511
5512 async fn delete_table_tag(
5513 &self,
5514 request: DeleteTableTagRequest,
5515 ) -> Result<DeleteTableTagResponse> {
5516 self.record_op("delete_table_tag");
5517 if request.tag.is_empty() {
5518 return Err(NamespaceError::InvalidInput {
5519 message: "tag name must not be empty for delete_table_tag".to_string(),
5520 }
5521 .into());
5522 }
5523
5524 let table_uri = self.resolve_table_location(&request.id).await?;
5525 let dataset = self
5526 .load_dataset(&table_uri, None, "delete_table_tag")
5527 .await?;
5528
5529 dataset
5530 .tags()
5531 .delete(&request.tag)
5532 .await
5533 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5534
5535 Ok(DeleteTableTagResponse {
5536 transaction_id: None,
5537 })
5538 }
5539
5540 async fn update_table_tag(
5541 &self,
5542 request: UpdateTableTagRequest,
5543 ) -> Result<UpdateTableTagResponse> {
5544 self.record_op("update_table_tag");
5545 if request.tag.is_empty() {
5546 return Err(NamespaceError::InvalidInput {
5547 message: "tag name must not be empty for update_table_tag".to_string(),
5548 }
5549 .into());
5550 }
5551 if request.version <= 0 {
5552 return Err(NamespaceError::InvalidInput {
5553 message: format!(
5554 "tag version must be a positive integer, got {} for update_table_tag",
5555 request.version
5556 ),
5557 }
5558 .into());
5559 }
5560
5561 let table_uri = self.resolve_table_location(&request.id).await?;
5562 let dataset = self
5563 .load_dataset(&table_uri, None, "update_table_tag")
5564 .await?;
5565
5566 dataset
5567 .tags()
5568 .update(&request.tag, request.version as u64)
5569 .await
5570 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5571
5572 Ok(UpdateTableTagResponse {
5573 transaction_id: None,
5574 })
5575 }
5576
5577 async fn create_table_branch(
5578 &self,
5579 request: CreateTableBranchRequest,
5580 ) -> Result<CreateTableBranchResponse> {
5581 self.record_op("create_table_branch");
5582 if request.name.is_empty() {
5583 return Err(NamespaceError::InvalidInput {
5584 message: "branch name must not be empty for create_table_branch".to_string(),
5585 }
5586 .into());
5587 }
5588 let from_version = match request.from_version {
5589 Some(v) if v <= 0 => {
5590 return Err(NamespaceError::InvalidInput {
5591 message: format!(
5592 "from_version must be a positive integer, got {} for create_table_branch",
5593 v
5594 ),
5595 }
5596 .into());
5597 }
5598 Some(v) => Some(v as u64),
5599 None => None,
5600 };
5601
5602 let table_uri = self.resolve_table_location(&request.id).await?;
5603 let mut dataset = self
5604 .load_dataset(&table_uri, None, "create_table_branch")
5605 .await?;
5606
5607 if dataset.branches().get(&request.name).await.is_ok() {
5612 return Err(NamespaceError::TableBranchAlreadyExists {
5613 message: format!("branch '{}' for table at '{}'", request.name, table_uri),
5614 }
5615 .into());
5616 }
5617
5618 dataset
5619 .create_branch(
5620 &request.name,
5621 (request.from_branch.as_deref(), from_version),
5622 None,
5623 )
5624 .await
5625 .map_err(|e| {
5626 if matches!(e, lance_core::Error::DatasetNotFound { .. }) {
5629 NamespaceError::InvalidInput {
5630 message: format!(
5631 "from_branch/from_version for branch '{}' refers to a source that does not exist: {}",
5632 request.name, e
5633 ),
5634 }
5635 .into()
5636 } else {
5637 Self::map_branch_error(e, &request.name, &table_uri)
5638 }
5639 })?;
5640
5641 Ok(CreateTableBranchResponse {
5642 transaction_id: None,
5643 })
5644 }
5645
5646 async fn list_table_branches(
5647 &self,
5648 request: ListTableBranchesRequest,
5649 ) -> Result<ListTableBranchesResponse> {
5650 self.record_op("list_table_branches");
5651 let table_uri = self.resolve_table_location(&request.id).await?;
5652 let dataset = self
5653 .load_dataset(&table_uri, None, "list_table_branches")
5654 .await?;
5655
5656 let raw_branches = dataset.list_branches().await.map_err(|e| {
5657 lance_core::Error::from(NamespaceError::Internal {
5658 message: format!(
5659 "Failed to list branches for table at '{}': {}",
5660 table_uri, e
5661 ),
5662 })
5663 })?;
5664
5665 let branches = raw_branches
5666 .into_iter()
5667 .map(|(name, contents)| {
5668 let mut branch_model = ModelBranchContents::new(
5671 contents.parent_version as i64,
5672 contents.create_at as i64,
5673 contents.manifest_size as i64,
5674 );
5675 branch_model.parent_branch = contents.parent_branch;
5676 branch_model.metadata = if contents.metadata.is_empty() {
5677 None
5678 } else {
5679 Some(contents.metadata)
5680 };
5681 (name, branch_model)
5682 })
5683 .collect();
5684
5685 Ok(ListTableBranchesResponse {
5686 branches,
5687 page_token: None,
5688 })
5689 }
5690
5691 async fn delete_table_branch(
5692 &self,
5693 request: DeleteTableBranchRequest,
5694 ) -> Result<DeleteTableBranchResponse> {
5695 self.record_op("delete_table_branch");
5696 if request.name.is_empty() {
5697 return Err(NamespaceError::InvalidInput {
5698 message: "branch name must not be empty for delete_table_branch".to_string(),
5699 }
5700 .into());
5701 }
5702
5703 let table_uri = self.resolve_table_location(&request.id).await?;
5704 let mut dataset = self
5705 .load_dataset(&table_uri, None, "delete_table_branch")
5706 .await?;
5707
5708 dataset
5709 .delete_branch(&request.name)
5710 .await
5711 .map_err(|e| match e {
5712 lance_core::Error::RefConflict { message } => NamespaceError::InvalidInput {
5713 message: format!(
5714 "branch '{}' for table at '{}': {}",
5715 request.name, table_uri, message
5716 ),
5717 }
5718 .into(),
5719 other => Self::map_branch_error(other, &request.name, &table_uri),
5720 })?;
5721
5722 Ok(DeleteTableBranchResponse {
5723 transaction_id: None,
5724 })
5725 }
5726
5727 fn namespace_id(&self) -> String {
5728 format!("DirectoryNamespace {{ root: {:?} }}", self.root)
5729 }
5730}
5731
5732#[derive(Debug)]
5734pub(crate) enum MarkerFileError {
5735 AlreadyExists { description: String },
5737 Other { message: String },
5739}
5740
5741impl std::fmt::Display for MarkerFileError {
5742 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5743 match self {
5744 Self::AlreadyExists { description } => {
5745 write!(f, "{} already exists", description)
5746 }
5747 Self::Other { message } => write!(f, "{}", message),
5748 }
5749 }
5750}
5751
5752pub(crate) async fn put_marker_file_atomic(
5773 object_store: &ObjectStore,
5774 path: &Path,
5775 file_description: &str,
5776) -> std::result::Result<(), MarkerFileError> {
5777 let staging_name = format!("lance-marker.staging.{}", uuid::Uuid::new_v4().simple());
5778 let path_str = path.as_ref();
5779 let staging_path = match path_str.rfind('/') {
5780 Some(idx) => Path::from(format!("{}/{}", &path_str[..idx], staging_name)),
5781 None => Path::from(staging_name.as_str()),
5782 };
5783
5784 object_store
5785 .inner
5786 .put(&staging_path, bytes::Bytes::from_static(b"reserved").into())
5787 .await
5788 .map_err(|e| MarkerFileError::Other {
5789 message: format!("Failed to stage {}: {:?}", file_description, e),
5790 })?;
5791
5792 let mut staging_consumed = false;
5795 let publish_result = match object_store
5796 .inner
5797 .rename_if_not_exists(&staging_path, path)
5798 .await
5799 {
5800 Ok(()) => {
5801 staging_consumed = true;
5802 Ok(())
5803 }
5804 Err(ObjectStoreError::NotImplemented { .. })
5805 | Err(ObjectStoreError::NotSupported { .. }) => {
5806 match object_store
5807 .inner
5808 .copy_if_not_exists(&staging_path, path)
5809 .await
5810 {
5811 Ok(()) => Ok(()),
5812 Err(ObjectStoreError::NotImplemented { .. })
5813 | Err(ObjectStoreError::NotSupported { .. }) => object_store
5814 .inner
5815 .put_opts(
5816 path,
5817 bytes::Bytes::from_static(b"reserved").into(),
5818 PutOptions {
5819 mode: PutMode::Create,
5820 ..Default::default()
5821 },
5822 )
5823 .await
5824 .map(|_| ()),
5825 Err(e) => Err(e),
5826 }
5827 }
5828 Err(e) => Err(e),
5829 };
5830
5831 if !staging_consumed {
5832 delete_staging_marker_best_effort(object_store, &staging_path).await;
5833 }
5834
5835 match publish_result {
5836 Ok(()) => Ok(()),
5837 Err(ObjectStoreError::AlreadyExists { .. })
5838 | Err(ObjectStoreError::Precondition { .. }) => Err(MarkerFileError::AlreadyExists {
5839 description: file_description.to_string(),
5840 }),
5841 Err(e) => Err(MarkerFileError::Other {
5842 message: format!("Failed to create {}: {:?}", file_description, e),
5843 }),
5844 }
5845}
5846
5847async fn delete_staging_marker_best_effort(object_store: &ObjectStore, staging_path: &Path) {
5851 const MAX_ATTEMPTS: u32 = 3;
5852 const BACKOFF_MS: [u64; 2] = [20, 50];
5853
5854 let mut last_err: Option<ObjectStoreError> = None;
5855 for attempt in 0..MAX_ATTEMPTS {
5856 match object_store.inner.delete(staging_path).await {
5857 Ok(()) => return,
5858 Err(ObjectStoreError::NotFound { .. }) => return,
5859 Err(e) => {
5860 last_err = Some(e);
5861 if let Some(&delay_ms) = BACKOFF_MS.get(attempt as usize) {
5862 tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
5863 }
5864 }
5865 }
5866 }
5867 if let Some(del_err) = last_err {
5868 log::warn!(
5869 "Failed to delete staging marker at '{}' after {} attempts: {:?}",
5870 staging_path,
5871 MAX_ATTEMPTS,
5872 del_err
5873 );
5874 }
5875}
5876
5877fn build_engine_fts_query(
5881 query: &lance_namespace::models::FtsQuery,
5882) -> std::result::Result<FtsQuery, NamespaceError> {
5883 if let Some(ref m) = query.r#match {
5884 Ok(FtsQuery::Match(build_engine_match_query(m)?))
5885 } else if let Some(ref p) = query.phrase {
5886 let mut phrase = PhraseQuery::new(p.terms.clone());
5887 if let Some(ref column) = p.column {
5888 phrase = phrase.with_column(Some(column.clone()));
5889 }
5890 if let Some(slop) = p.slop {
5891 phrase = phrase.with_slop(slop as u32);
5892 }
5893 Ok(FtsQuery::Phrase(phrase))
5894 } else if let Some(ref mm) = query.multi_match {
5895 let match_queries = mm
5896 .match_queries
5897 .iter()
5898 .map(build_engine_match_query)
5899 .collect::<std::result::Result<Vec<_>, _>>()?;
5900 Ok(FtsQuery::MultiMatch(MultiMatchQuery { match_queries }))
5901 } else if let Some(ref b) = query.boolean {
5902 let mut clauses: Vec<(Occur, FtsQuery)> = Vec::new();
5903 for clause in &b.must {
5904 clauses.push((Occur::Must, build_engine_fts_query(clause)?));
5905 }
5906 for clause in &b.should {
5907 clauses.push((Occur::Should, build_engine_fts_query(clause)?));
5908 }
5909 for clause in &b.must_not {
5910 clauses.push((Occur::MustNot, build_engine_fts_query(clause)?));
5911 }
5912 Ok(FtsQuery::Boolean(BooleanQuery::new(clauses)))
5913 } else if let Some(ref boost) = query.boost {
5914 let positive = build_engine_fts_query(&boost.positive)?;
5915 let negative = build_engine_fts_query(&boost.negative)?;
5916 Ok(FtsQuery::Boost(BoostQuery::new(
5917 positive,
5918 negative,
5919 boost.negative_boost,
5920 )))
5921 } else {
5922 Err(NamespaceError::InvalidInput {
5923 message: "structured_query.query must set exactly one of match, phrase, multi_match, \
5924 boolean, or boost"
5925 .to_string(),
5926 })
5927 }
5928}
5929
5930fn build_engine_match_query(
5931 m: &lance_namespace::models::MatchQuery,
5932) -> std::result::Result<MatchQuery, NamespaceError> {
5933 let mut match_query = MatchQuery::new(m.terms.clone());
5934 if let Some(ref column) = m.column {
5935 match_query = match_query.with_column(Some(column.clone()));
5936 }
5937 if let Some(boost) = m.boost {
5938 match_query = match_query.with_boost(boost);
5939 }
5940 if let Some(fuzziness) = m.fuzziness {
5941 match_query = match_query.with_fuzziness(Some(fuzziness as u32));
5942 }
5943 if let Some(max_expansions) = m.max_expansions {
5944 match_query = match_query.with_max_expansions(max_expansions as usize);
5945 }
5946 if let Some(ref operator) = m.operator {
5947 let op =
5948 Operator::try_from(operator.as_str()).map_err(|e| NamespaceError::InvalidInput {
5949 message: format!("Invalid FTS operator: {:?}", e),
5950 })?;
5951 match_query = match_query.with_operator(op);
5952 }
5953 if let Some(prefix_length) = m.prefix_length {
5954 match_query = match_query.with_prefix_length(prefix_length as u32);
5955 }
5956 Ok(match_query)
5957}
5958
5959#[cfg(test)]
5960mod tests {
5961 use super::*;
5962 use arrow_ipc::reader::{FileReader, StreamReader};
5963
5964 #[test]
5965 fn test_build_engine_fts_query_match() {
5966 let mut ns_match = lance_namespace::models::MatchQuery::new("hello world".to_string());
5967 ns_match.column = Some("body".to_string());
5968 ns_match.operator = Some("AND".to_string());
5969 ns_match.fuzziness = Some(1);
5970 ns_match.max_expansions = Some(30);
5971 ns_match.boost = Some(2.0);
5972 ns_match.prefix_length = Some(2);
5973
5974 let mut ns_query = lance_namespace::models::FtsQuery::new();
5975 ns_query.r#match = Some(Box::new(ns_match));
5976
5977 match build_engine_fts_query(&ns_query).unwrap() {
5978 FtsQuery::Match(m) => {
5979 assert_eq!(m.terms, "hello world");
5980 assert_eq!(m.column, Some("body".to_string()));
5981 assert_eq!(m.operator, Operator::And);
5982 assert_eq!(m.fuzziness, Some(1));
5983 assert_eq!(m.max_expansions, 30);
5984 assert_eq!(m.boost, 2.0);
5985 assert_eq!(m.prefix_length, 2);
5986 }
5987 other => panic!("expected Match, got {:?}", other),
5988 }
5989 }
5990
5991 fn ns_match_query(terms: &str, column: &str) -> lance_namespace::models::FtsQuery {
5994 let mut m = lance_namespace::models::MatchQuery::new(terms.to_string());
5995 m.column = Some(column.to_string());
5996 let mut q = lance_namespace::models::FtsQuery::new();
5997 q.r#match = Some(Box::new(m));
5998 q
5999 }
6000
6001 #[test]
6002 fn test_build_engine_fts_query_phrase() {
6003 let mut ns_phrase = lance_namespace::models::PhraseQuery::new("hello world".to_string());
6004 ns_phrase.column = Some("body".to_string());
6005 ns_phrase.slop = Some(2);
6006
6007 let mut ns_query = lance_namespace::models::FtsQuery::new();
6008 ns_query.phrase = Some(Box::new(ns_phrase));
6009
6010 match build_engine_fts_query(&ns_query).unwrap() {
6011 FtsQuery::Phrase(p) => {
6012 assert_eq!(p.terms, "hello world");
6013 assert_eq!(p.column, Some("body".to_string()));
6014 assert_eq!(p.slop, 2);
6015 }
6016 other => panic!("expected Phrase, got {:?}", other),
6017 }
6018 }
6019
6020 #[test]
6021 fn test_build_engine_fts_query_multi_match() {
6022 let mut m1 = lance_namespace::models::MatchQuery::new("hello".to_string());
6023 m1.column = Some("title".to_string());
6024 let mut m2 = lance_namespace::models::MatchQuery::new("hello".to_string());
6025 m2.column = Some("body".to_string());
6026 m2.boost = Some(2.0);
6027
6028 let ns_multi = lance_namespace::models::MultiMatchQuery::new(vec![m1, m2]);
6029 let mut ns_query = lance_namespace::models::FtsQuery::new();
6030 ns_query.multi_match = Some(Box::new(ns_multi));
6031
6032 match build_engine_fts_query(&ns_query).unwrap() {
6033 FtsQuery::MultiMatch(mm) => {
6034 assert_eq!(mm.match_queries.len(), 2);
6035 assert_eq!(mm.match_queries[0].terms, "hello");
6036 assert_eq!(mm.match_queries[0].column, Some("title".to_string()));
6037 assert_eq!(mm.match_queries[1].column, Some("body".to_string()));
6038 assert_eq!(mm.match_queries[1].boost, 2.0);
6039 }
6040 other => panic!("expected MultiMatch, got {:?}", other),
6041 }
6042 }
6043
6044 #[test]
6045 fn test_build_engine_fts_query_boolean() {
6046 let ns_boolean = lance_namespace::models::BooleanQuery::new(
6048 vec![ns_match_query("must-term", "body")],
6049 vec![ns_match_query("must-not-term", "body")],
6050 vec![ns_match_query("should-term", "body")],
6051 );
6052 let mut ns_query = lance_namespace::models::FtsQuery::new();
6053 ns_query.boolean = Some(Box::new(ns_boolean));
6054
6055 match build_engine_fts_query(&ns_query).unwrap() {
6056 FtsQuery::Boolean(b) => {
6057 assert!(matches!(&b.must[..], [FtsQuery::Match(m)] if m.terms == "must-term"));
6058 assert!(
6059 matches!(&b.must_not[..], [FtsQuery::Match(m)] if m.terms == "must-not-term")
6060 );
6061 assert!(matches!(&b.should[..], [FtsQuery::Match(m)] if m.terms == "should-term"));
6062 }
6063 other => panic!("expected Boolean, got {:?}", other),
6064 }
6065 }
6066
6067 #[test]
6068 fn test_build_engine_fts_query_boost() {
6069 let mut ns_boost = lance_namespace::models::BoostQuery::new(
6070 ns_match_query("positive-term", "body"),
6071 ns_match_query("negative-term", "body"),
6072 );
6073 ns_boost.negative_boost = Some(0.25);
6074
6075 let mut ns_query = lance_namespace::models::FtsQuery::new();
6076 ns_query.boost = Some(Box::new(ns_boost));
6077
6078 match build_engine_fts_query(&ns_query).unwrap() {
6079 FtsQuery::Boost(b) => {
6080 assert!(
6081 matches!(b.positive.as_ref(), FtsQuery::Match(m) if m.terms == "positive-term")
6082 );
6083 assert!(
6084 matches!(b.negative.as_ref(), FtsQuery::Match(m) if m.terms == "negative-term")
6085 );
6086 assert_eq!(b.negative_boost, 0.25);
6087 }
6088 other => panic!("expected Boost, got {:?}", other),
6089 }
6090 }
6091
6092 #[test]
6093 fn test_build_engine_fts_query_requires_a_variant() {
6094 let empty = lance_namespace::models::FtsQuery::new();
6096 assert!(build_engine_fts_query(&empty).is_err());
6097 }
6098 use lance::dataset::Dataset;
6099 use lance::index::DatasetIndexExt;
6100 use lance_core::utils::tempfile::{TempStdDir, TempStrDir};
6101 use lance_core::utils::testing::CountingObjectStore;
6102 use lance_io::object_store::{providers::local::FileStoreProvider, uri_to_url};
6103 use lance_namespace::error::ErrorCode;
6104 use lance_namespace::models::{
6105 CreateTableRequest, JsonArrowDataType, JsonArrowField, JsonArrowSchema, ListTablesRequest,
6106 QueryTableRequestColumns,
6107 };
6108 use lance_namespace::schema::convert_json_arrow_schema;
6109 use std::io::Cursor;
6110 use std::sync::{
6111 Arc,
6112 atomic::{AtomicUsize, Ordering},
6113 };
6114 use url::Url;
6115
6116 fn assert_plan_contains_all(plan: &str, expected_fragments: &[&str], context: &str) {
6117 for expected_fragment in expected_fragments {
6118 assert!(
6119 plan.contains(expected_fragment),
6120 "{}. Missing fragment: '{}'. Plan:\n{}",
6121 context,
6122 expected_fragment,
6123 plan
6124 );
6125 }
6126 }
6127
6128 fn mutation_error_code(err: lance_core::Error) -> ErrorCode {
6129 match err {
6130 lance_core::Error::Namespace { source, .. } => source
6131 .downcast_ref::<NamespaceError>()
6132 .expect("mutation error should wrap a NamespaceError")
6133 .code(),
6134 other => panic!("expected Namespace error, got: {other:?}"),
6135 }
6136 }
6137
6138 #[test]
6143 fn test_map_mutation_error_commit_conflict_alignment() {
6144 let boxed = || -> Box<dyn std::error::Error + Send + Sync + 'static> {
6145 Box::<dyn std::error::Error + Send + Sync>::from("inner conflict")
6146 };
6147
6148 let throttling_cases = vec![lance_core::Error::commit_conflict_source(1, boxed())];
6149 for err in throttling_cases {
6150 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6151 err,
6152 "update",
6153 "memory://t",
6154 ));
6155 assert_eq!(code, ErrorCode::Throttling);
6156 }
6157
6158 let concurrent_cases = vec![
6159 lance_core::Error::too_much_write_contention("contention"),
6160 lance_core::Error::retryable_commit_conflict_source(1, boxed()),
6161 lance_core::Error::incompatible_transaction_source(boxed()),
6162 lance_core::Error::version_conflict("conflict", 0, 3),
6163 ];
6164 for err in concurrent_cases {
6165 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6166 err,
6167 "update",
6168 "memory://t",
6169 ));
6170 assert_eq!(code, ErrorCode::ConcurrentModification);
6171 }
6172 }
6173
6174 async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) {
6176 let temp_dir = TempStdDir::default();
6177
6178 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
6179 .build()
6180 .await
6181 .unwrap();
6182 (namespace, temp_dir)
6183 }
6184
6185 #[derive(Debug)]
6186 #[allow(dead_code)]
6187 struct CountingFileStoreProvider {
6188 listing_count: Arc<AtomicUsize>,
6189 }
6190
6191 #[async_trait]
6192 impl lance_io::object_store::ObjectStoreProvider for CountingFileStoreProvider {
6193 async fn new_store(
6194 &self,
6195 base_path: Url,
6196 params: &ObjectStoreParams,
6197 ) -> Result<ObjectStore> {
6198 let provider = FileStoreProvider;
6199 let mut store = provider.new_store(base_path, params).await?;
6200 store.inner = Arc::new(CountingObjectStore::new(
6201 store.inner.clone(),
6202 self.listing_count.clone(),
6203 ));
6204 Ok(store)
6205 }
6206
6207 fn extract_path(&self, url: &Url) -> Result<Path> {
6208 let provider = FileStoreProvider;
6209 provider.extract_path(url)
6210 }
6211
6212 fn calculate_object_store_prefix(
6213 &self,
6214 url: &Url,
6215 storage_options: Option<&HashMap<String, String>>,
6216 ) -> Result<String> {
6217 let provider = FileStoreProvider;
6218 provider.calculate_object_store_prefix(url, storage_options)
6219 }
6220 }
6221
6222 #[allow(dead_code)]
6223 fn file_object_store_uri(path: &str) -> String {
6224 let file_url = uri_to_url(path).unwrap();
6225 let mut url = Url::parse("file-object-store:///").unwrap();
6226 url.set_path(file_url.path());
6227 url.to_string()
6228 }
6229
6230 #[allow(dead_code)]
6231 fn build_listing_counting_session(listing_count: Arc<AtomicUsize>) -> Arc<Session> {
6232 let registry = Arc::new(ObjectStoreRegistry::default());
6233 registry.insert(
6234 "file-object-store",
6235 Arc::new(CountingFileStoreProvider { listing_count }),
6236 );
6237 Arc::new(Session::new(0, 0, registry))
6238 }
6239
6240 use futures::stream::BoxStream;
6244 use object_store::{
6245 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
6246 PutMultipartOptions, PutPayload, PutResult, Result as OSResult,
6247 };
6248 use std::ops::Range;
6249
6250 #[derive(Debug, Clone, Copy)]
6251 enum ListBehavior {
6252 Throttle,
6253 ServiceUnavailable,
6254 Internal,
6255 NotFound,
6256 EmptyListing,
6257 }
6258
6259 #[derive(Debug)]
6260 struct FailingListStore {
6261 target: Arc<dyn OSObjectStore>,
6262 behavior: Arc<Mutex<Option<ListBehavior>>>,
6263 }
6264
6265 impl std::fmt::Display for FailingListStore {
6266 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6267 write!(f, "FailingListStore({})", self.target)
6268 }
6269 }
6270
6271 #[async_trait]
6272 impl OSObjectStore for FailingListStore {
6273 async fn put_opts(
6274 &self,
6275 location: &Path,
6276 bytes: PutPayload,
6277 opts: PutOptions,
6278 ) -> OSResult<PutResult> {
6279 self.target.put_opts(location, bytes, opts).await
6280 }
6281
6282 async fn put_multipart_opts(
6283 &self,
6284 location: &Path,
6285 opts: PutMultipartOptions,
6286 ) -> OSResult<Box<dyn MultipartUpload>> {
6287 self.target.put_multipart_opts(location, opts).await
6288 }
6289
6290 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
6291 self.target.get_opts(location, options).await
6292 }
6293
6294 async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> OSResult<Vec<Bytes>> {
6295 self.target.get_ranges(location, ranges).await
6296 }
6297
6298 fn delete_stream(
6299 &self,
6300 locations: BoxStream<'static, OSResult<Path>>,
6301 ) -> BoxStream<'static, OSResult<Path>> {
6302 self.target.delete_stream(locations)
6303 }
6304
6305 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
6306 self.target.list(prefix)
6307 }
6308
6309 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
6310 let behavior = *self.behavior.lock().unwrap();
6311 match behavior {
6312 None => self.target.list_with_delimiter(prefix).await,
6313 Some(ListBehavior::EmptyListing) => Ok(ListResult {
6314 common_prefixes: Vec::new(),
6315 objects: Vec::new(),
6316 }),
6317 Some(ListBehavior::Throttle) => Err(ObjectStoreError::Generic {
6320 store: "test",
6321 source: "Error performing list request: response error, after 3 retries, \
6322 max_retries: 3, retry_timeout: 180s - HTTP status server error \
6323 (503 Service Unavailable): ServerBusy: The server is busy"
6324 .into(),
6325 }),
6326 Some(ListBehavior::ServiceUnavailable) => Err(ObjectStoreError::Generic {
6327 store: "test",
6328 source: "Error performing list request: 503 Service Unavailable".into(),
6329 }),
6330 Some(ListBehavior::Internal) => Err(ObjectStoreError::Generic {
6331 store: "test",
6332 source: "Error performing list request: catastrophic unclassified failure"
6333 .into(),
6334 }),
6335 Some(ListBehavior::NotFound) => Err(ObjectStoreError::NotFound {
6336 path: "test_table.lance".to_string(),
6337 source: "entity not found".into(),
6338 }),
6339 }
6340 }
6341
6342 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
6343 self.target.copy_opts(from, to, opts).await
6344 }
6345 }
6346
6347 #[derive(Debug)]
6348 struct FailingListStoreProvider {
6349 behavior: Arc<Mutex<Option<ListBehavior>>>,
6350 }
6351
6352 #[async_trait]
6353 impl lance_io::object_store::ObjectStoreProvider for FailingListStoreProvider {
6354 async fn new_store(
6355 &self,
6356 base_path: Url,
6357 params: &ObjectStoreParams,
6358 ) -> Result<ObjectStore> {
6359 let mut store = FileStoreProvider.new_store(base_path, params).await?;
6360 store.inner = Arc::new(FailingListStore {
6361 target: store.inner.clone(),
6362 behavior: self.behavior.clone(),
6363 });
6364 Ok(store)
6365 }
6366
6367 fn extract_path(&self, url: &Url) -> Result<Path> {
6368 FileStoreProvider.extract_path(url)
6369 }
6370
6371 fn calculate_object_store_prefix(
6372 &self,
6373 url: &Url,
6374 storage_options: Option<&HashMap<String, String>>,
6375 ) -> Result<String> {
6376 FileStoreProvider.calculate_object_store_prefix(url, storage_options)
6377 }
6378 }
6379
6380 fn build_failing_list_session(behavior: Arc<Mutex<Option<ListBehavior>>>) -> Arc<Session> {
6381 let registry = Arc::new(ObjectStoreRegistry::default());
6382 registry.insert(
6383 "file-object-store",
6384 Arc::new(FailingListStoreProvider { behavior }),
6385 );
6386 Arc::new(Session::new(0, 0, registry))
6387 }
6388
6389 async fn failing_list_namespace() -> (
6393 DirectoryNamespace,
6394 TempStdDir,
6395 Arc<Mutex<Option<ListBehavior>>>,
6396 ) {
6397 let temp_dir = TempStdDir::default();
6398 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
6399 let behavior = Arc::new(Mutex::new(None));
6400 let session = build_failing_list_session(behavior.clone());
6401 let namespace = DirectoryNamespaceBuilder::new(root_uri)
6402 .session(session)
6403 .manifest_enabled(false)
6404 .dir_listing_enabled(true)
6405 .build()
6406 .await
6407 .unwrap();
6408 (namespace, temp_dir, behavior)
6409 }
6410
6411 async fn create_named_dir_table(namespace: &DirectoryNamespace, name: &str) {
6412 let schema = create_test_schema();
6413 let ipc_data = create_test_ipc_data(&schema);
6414 let mut create_req = CreateTableRequest::new();
6415 create_req.id = Some(vec![name.to_string()]);
6416 namespace
6417 .create_table(create_req, Bytes::from(ipc_data))
6418 .await
6419 .unwrap();
6420 }
6421
6422 #[tokio::test]
6427 async fn test_table_resolution_propagates_storage_errors_not_table_not_found() {
6428 for (behavior, expected_code, evidence) in [
6429 (ListBehavior::Throttle, ErrorCode::Throttling, "serverbusy"),
6430 (
6431 ListBehavior::ServiceUnavailable,
6432 ErrorCode::ServiceUnavailable,
6433 "503 service unavailable",
6434 ),
6435 (ListBehavior::Internal, ErrorCode::Internal, "catastrophic"),
6436 ] {
6437 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6438 create_named_dir_table(&namespace, "checkpoint").await;
6439 *toggle.lock().unwrap() = Some(behavior);
6440
6441 let mut describe_req = DescribeTableRequest::new();
6442 describe_req.id = Some(vec!["checkpoint".to_string()]);
6443 let err = namespace.describe_table(describe_req).await.unwrap_err();
6444 let msg = err.to_string();
6445 assert_eq!(
6446 mutation_error_code(err),
6447 expected_code,
6448 "describe_table under {behavior:?}; msg: {msg}"
6449 );
6450 assert!(
6451 msg.to_ascii_lowercase().contains(evidence),
6452 "describe_table message must carry storage evidence '{evidence}', got: {msg}"
6453 );
6454
6455 let mut exists_req = TableExistsRequest::new();
6456 exists_req.id = Some(vec!["checkpoint".to_string()]);
6457 let err = namespace.table_exists(exists_req).await.unwrap_err();
6458 let msg = err.to_string();
6459 assert_eq!(
6460 mutation_error_code(err),
6461 expected_code,
6462 "table_exists under {behavior:?}; msg: {msg}"
6463 );
6464 assert!(
6465 msg.to_ascii_lowercase().contains(evidence),
6466 "table_exists message must carry storage evidence '{evidence}', got: {msg}"
6467 );
6468 }
6469 }
6470
6471 #[tokio::test]
6474 async fn test_table_resolution_missing_table_yields_table_not_found() {
6475 for behavior in [ListBehavior::NotFound, ListBehavior::EmptyListing] {
6476 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6477 *toggle.lock().unwrap() = Some(behavior);
6478
6479 let mut describe_req = DescribeTableRequest::new();
6480 describe_req.id = Some(vec!["missing".to_string()]);
6481 let err = namespace.describe_table(describe_req).await.unwrap_err();
6482 assert_eq!(
6483 mutation_error_code(err),
6484 ErrorCode::TableNotFound,
6485 "describe_table under {behavior:?} should be TableNotFound"
6486 );
6487
6488 let mut exists_req = TableExistsRequest::new();
6489 exists_req.id = Some(vec!["missing".to_string()]);
6490 let err = namespace.table_exists(exists_req).await.unwrap_err();
6491 assert_eq!(
6492 mutation_error_code(err),
6493 ErrorCode::TableNotFound,
6494 "table_exists under {behavior:?} should be TableNotFound"
6495 );
6496 }
6497 }
6498
6499 #[tokio::test]
6509 async fn test_hybrid_resolution_falls_through_and_does_not_mask_throttle() {
6510 let temp_dir = TempStdDir::default();
6511 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
6512 let behavior = Arc::new(Mutex::new(None));
6513 let session = build_failing_list_session(behavior.clone());
6514
6515 let manifest_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
6517 .session(session.clone())
6518 .manifest_enabled(true)
6519 .dir_listing_enabled(true)
6520 .build()
6521 .await
6522 .unwrap();
6523 create_named_dir_table(&manifest_ns, "seed").await;
6524
6525 let dir_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
6527 .session(session.clone())
6528 .manifest_enabled(false)
6529 .dir_listing_enabled(true)
6530 .build()
6531 .await
6532 .unwrap();
6533 create_named_dir_table(&dir_ns, "checkpoint").await;
6534
6535 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
6538 .session(session)
6539 .manifest_enabled(true)
6540 .dir_listing_enabled(true)
6541 .dir_listing_to_manifest_migration_enabled(true)
6542 .build()
6543 .await
6544 .unwrap();
6545
6546 let mut describe_req = DescribeTableRequest::new();
6549 describe_req.id = Some(vec!["checkpoint".to_string()]);
6550 hybrid_ns
6551 .describe_table(describe_req)
6552 .await
6553 .expect("unregistered on-disk table should resolve via the manifest fall-through");
6554
6555 *behavior.lock().unwrap() = Some(ListBehavior::Throttle);
6559 let mut describe_req = DescribeTableRequest::new();
6560 describe_req.id = Some(vec!["checkpoint".to_string()]);
6561 let err = hybrid_ns.describe_table(describe_req).await.unwrap_err();
6562 let code = mutation_error_code(err);
6563 assert_ne!(
6564 code,
6565 ErrorCode::TableNotFound,
6566 "hybrid resolution masked a throttle as TableNotFound"
6567 );
6568 assert!(
6569 matches!(
6570 code,
6571 ErrorCode::Throttling | ErrorCode::ServiceUnavailable | ErrorCode::Internal
6572 ),
6573 "hybrid resolution should surface a storage error, got {code:?}"
6574 );
6575 }
6576
6577 #[test]
6578 fn test_classify_storage_error_maps_variants_and_preserves_evidence() {
6579 let throttle: Error = ObjectStoreError::Generic {
6580 store: "test",
6581 source: "list request failed, after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
6582 }
6583 .into();
6584 assert!(matches!(throttle, Error::IO { .. }));
6585 let classified = DirectoryNamespace::classify_storage_error(throttle);
6586 let msg = classified.to_string();
6587 assert_eq!(mutation_error_code(classified), ErrorCode::Throttling);
6588 assert!(
6589 msg.to_ascii_lowercase().contains("serverbusy"),
6590 "throttle evidence lost: {msg}"
6591 );
6592
6593 let service: Error = ObjectStoreError::Generic {
6594 store: "test",
6595 source: "504 Gateway Timeout".into(),
6596 }
6597 .into();
6598 assert_eq!(
6599 mutation_error_code(DirectoryNamespace::classify_storage_error(service)),
6600 ErrorCode::ServiceUnavailable
6601 );
6602
6603 let internal: Error = ObjectStoreError::Generic {
6604 store: "test",
6605 source: "disk caught fire".into(),
6606 }
6607 .into();
6608 assert_eq!(
6609 mutation_error_code(DirectoryNamespace::classify_storage_error(internal)),
6610 ErrorCode::Internal
6611 );
6612
6613 let preexisting: Error = NamespaceError::TableAlreadyExists {
6615 message: "t".to_string(),
6616 }
6617 .into();
6618 assert_eq!(
6619 mutation_error_code(DirectoryNamespace::classify_storage_error(preexisting)),
6620 ErrorCode::TableAlreadyExists
6621 );
6622 }
6623
6624 #[test]
6625 fn test_is_manifest_table_absent_error() {
6626 let table_not_found: Error = NamespaceError::TableNotFound {
6627 message: "t".to_string(),
6628 }
6629 .into();
6630 assert!(DirectoryNamespace::is_manifest_table_absent_error(
6631 &table_not_found
6632 ));
6633 let raw_not_found: Error = ObjectStoreError::NotFound {
6634 path: "t".to_string(),
6635 source: "x".into(),
6636 }
6637 .into();
6638 assert!(DirectoryNamespace::is_manifest_table_absent_error(
6639 &raw_not_found
6640 ));
6641
6642 let throttle: Error = ObjectStoreError::Generic {
6643 store: "test",
6644 source: "after 3 retries, max_retries: 3 ServerBusy".into(),
6645 }
6646 .into();
6647 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
6648 &throttle
6649 ));
6650 let internal: Error = NamespaceError::Internal {
6651 message: "boom".to_string(),
6652 }
6653 .into();
6654 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
6655 &internal
6656 ));
6657 }
6658
6659 #[test]
6660 fn test_map_open_error() {
6661 let not_found = || NamespaceError::TableNotFound {
6662 message: "table at 'x' not found: ...".to_string(),
6663 };
6664
6665 let throttle: Error = ObjectStoreError::Generic {
6666 store: "test",
6667 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
6668 }
6669 .into();
6670 assert_eq!(
6671 mutation_error_code(DirectoryNamespace::map_open_error(throttle, not_found())),
6672 ErrorCode::Throttling
6673 );
6674
6675 let generic_io: Error = ObjectStoreError::Generic {
6676 store: "test",
6677 source: "connection reset".into(),
6678 }
6679 .into();
6680 assert_eq!(
6681 mutation_error_code(DirectoryNamespace::map_open_error(generic_io, not_found())),
6682 ErrorCode::Internal
6683 );
6684
6685 let io_not_found: Error = ObjectStoreError::NotFound {
6686 path: "x".to_string(),
6687 source: "missing".into(),
6688 }
6689 .into();
6690 assert_eq!(
6691 mutation_error_code(DirectoryNamespace::map_open_error(
6692 io_not_found,
6693 not_found()
6694 )),
6695 ErrorCode::TableNotFound
6696 );
6697
6698 let dataset_not_found = Error::dataset_not_found("x".to_string(), "missing".into());
6699 assert_eq!(
6700 mutation_error_code(DirectoryNamespace::map_open_error(
6701 dataset_not_found,
6702 not_found()
6703 )),
6704 ErrorCode::TableNotFound
6705 );
6706
6707 let ref_not_found = Error::RefNotFound {
6709 message: "branch 'b' does not exist".to_string(),
6710 };
6711 assert_eq!(
6712 mutation_error_code(DirectoryNamespace::map_open_error(
6713 ref_not_found,
6714 not_found()
6715 )),
6716 ErrorCode::TableNotFound
6717 );
6718
6719 let version_miss = Error::RefNotFound {
6721 message: "version 5 does not exist".to_string(),
6722 };
6723 assert_eq!(
6724 mutation_error_code(DirectoryNamespace::map_open_error(
6725 version_miss,
6726 NamespaceError::TableVersionNotFound {
6727 message: "version 5 not found".to_string(),
6728 },
6729 )),
6730 ErrorCode::TableVersionNotFound
6731 );
6732 let version_throttle: Error = ObjectStoreError::Generic {
6733 store: "test",
6734 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
6735 }
6736 .into();
6737 assert_eq!(
6738 mutation_error_code(DirectoryNamespace::map_open_error(
6739 version_throttle,
6740 NamespaceError::TableVersionNotFound {
6741 message: "version 5 not found".to_string(),
6742 },
6743 )),
6744 ErrorCode::Throttling
6745 );
6746 }
6747
6748 fn create_test_ipc_data(schema: &JsonArrowSchema) -> Vec<u8> {
6750 use arrow::ipc::writer::StreamWriter;
6751
6752 let arrow_schema = convert_json_arrow_schema(schema).unwrap();
6753 let arrow_schema = Arc::new(arrow_schema);
6754 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
6755 let mut buffer = Vec::new();
6756 {
6757 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
6758 writer.write(&batch).unwrap();
6759 writer.finish().unwrap();
6760 }
6761 buffer
6762 }
6763
6764 fn create_ipc_data_from_batches(
6765 schema: Arc<arrow_schema::Schema>,
6766 batches: Vec<arrow::record_batch::RecordBatch>,
6767 ) -> Vec<u8> {
6768 use arrow::ipc::writer::StreamWriter;
6769
6770 let mut buffer = Vec::new();
6771 {
6772 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
6773 for batch in &batches {
6774 writer.write(batch).unwrap();
6775 }
6776 writer.finish().unwrap();
6777 }
6778 buffer
6779 }
6780
6781 fn create_non_empty_test_ipc_data() -> Vec<u8> {
6782 use arrow::array::{Int32Array, StringArray};
6783 use arrow::record_batch::RecordBatch;
6784
6785 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
6786 let batch = RecordBatch::try_new(
6787 schema.clone(),
6788 vec![
6789 Arc::new(Int32Array::from(vec![1, 2])),
6790 Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])),
6791 ],
6792 )
6793 .unwrap();
6794 create_ipc_data_from_batches(schema, vec![batch])
6795 }
6796
6797 fn create_single_row_test_ipc_data() -> Vec<u8> {
6798 use arrow::array::{Int32Array, StringArray};
6799 use arrow::record_batch::RecordBatch;
6800
6801 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
6802 let batch = RecordBatch::try_new(
6803 schema.clone(),
6804 vec![
6805 Arc::new(Int32Array::from(vec![10])),
6806 Arc::new(StringArray::from(vec![Some("carol")])),
6807 ],
6808 )
6809 .unwrap();
6810 create_ipc_data_from_batches(schema, vec![batch])
6811 }
6812
6813 fn create_test_schema() -> JsonArrowSchema {
6815 let int_type = JsonArrowDataType::new("int32".to_string());
6816 let string_type = JsonArrowDataType::new("utf8".to_string());
6817
6818 let id_field = JsonArrowField {
6819 name: "id".to_string(),
6820 r#type: Box::new(int_type),
6821 nullable: false,
6822 metadata: None,
6823 };
6824
6825 let name_field = JsonArrowField {
6826 name: "name".to_string(),
6827 r#type: Box::new(string_type),
6828 nullable: true,
6829 metadata: None,
6830 };
6831
6832 JsonArrowSchema {
6833 fields: vec![id_field, name_field],
6834 metadata: None,
6835 }
6836 }
6837
6838 fn create_scalar_table_ipc_data() -> Vec<u8> {
6839 use arrow::array::{Int32Array, StringArray};
6840 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6841
6842 let schema = Arc::new(ArrowSchema::new(vec![
6843 Field::new("id", DataType::Int32, false),
6844 Field::new("name", DataType::Utf8, true),
6845 ]));
6846 let batch = arrow::record_batch::RecordBatch::try_new(
6847 schema.clone(),
6848 vec![
6849 Arc::new(Int32Array::from(vec![1, 2, 3])),
6850 Arc::new(StringArray::from(vec!["alice", "bob", "cory"])),
6851 ],
6852 )
6853 .unwrap();
6854 create_ipc_data_from_batches(schema, vec![batch])
6855 }
6856
6857 async fn create_legacy_manifest_without_primary_key_metadata(root: &str) {
6858 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6859 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
6860
6861 let schema = Arc::new(ArrowSchema::new(vec![
6862 Field::new("object_id", DataType::Utf8, false),
6863 Field::new("object_type", DataType::Utf8, false),
6864 Field::new("location", DataType::Utf8, true),
6865 Field::new("metadata", DataType::Utf8, true),
6866 Field::new(
6867 "base_objects",
6868 DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))),
6869 true,
6870 ),
6871 ]));
6872 let batch = RecordBatch::new_empty(schema.clone());
6873 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
6874 Dataset::write(Box::new(reader), &format!("{}/__manifest", root), None)
6875 .await
6876 .unwrap();
6877 }
6878
6879 async fn manifest_has_primary_key_metadata(root: &str) -> bool {
6880 let dataset = Dataset::open(&format!("{}/__manifest", root))
6881 .await
6882 .unwrap();
6883 dataset
6884 .schema()
6885 .field("object_id")
6886 .map(|field| {
6887 field
6888 .metadata
6889 .contains_key(lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
6890 })
6891 .unwrap_or(false)
6892 }
6893
6894 fn create_vector_table_ipc_data() -> Vec<u8> {
6895 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
6896 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
6897
6898 let schema = Arc::new(ArrowSchema::new(vec![
6899 Field::new("id", DataType::Int32, false),
6900 Field::new(
6901 "vector",
6902 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
6903 true,
6904 ),
6905 ]));
6906 let vector_field = Arc::new(Field::new("item", DataType::Float32, true));
6907 let vectors = FixedSizeListArray::try_new(
6908 vector_field,
6909 2,
6910 Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6])),
6911 None,
6912 )
6913 .unwrap();
6914 let batch = arrow::record_batch::RecordBatch::try_new(
6915 schema.clone(),
6916 vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(vectors)],
6917 )
6918 .unwrap();
6919 create_ipc_data_from_batches(schema, vec![batch])
6920 }
6921
6922 async fn create_scalar_table(namespace: &DirectoryNamespace, table_name: &str) {
6923 let mut create_table_request = CreateTableRequest::new();
6924 create_table_request.id = Some(vec![table_name.to_string()]);
6925 namespace
6926 .create_table(
6927 create_table_request,
6928 Bytes::from(create_scalar_table_ipc_data()),
6929 )
6930 .await
6931 .unwrap();
6932 }
6933
6934 async fn create_vector_table(namespace: &DirectoryNamespace, table_name: &str) {
6935 let mut create_table_request = CreateTableRequest::new();
6936 create_table_request.id = Some(vec![table_name.to_string()]);
6937 namespace
6938 .create_table(
6939 create_table_request,
6940 Bytes::from(create_vector_table_ipc_data()),
6941 )
6942 .await
6943 .unwrap();
6944 }
6945
6946 async fn open_dataset(namespace: &DirectoryNamespace, table_name: &str) -> Dataset {
6947 let mut describe_request = DescribeTableRequest::new();
6948 describe_request.id = Some(vec![table_name.to_string()]);
6949 let table_uri = namespace
6950 .describe_table(describe_request)
6951 .await
6952 .unwrap()
6953 .location
6954 .expect("table location should exist");
6955 Dataset::open(&table_uri).await.unwrap()
6956 }
6957
6958 async fn create_scalar_index(
6959 namespace: &DirectoryNamespace,
6960 table_name: &str,
6961 index_name: &str,
6962 ) -> Option<String> {
6963 use lance_namespace::models::CreateTableIndexRequest;
6964
6965 let mut create_index_request =
6966 CreateTableIndexRequest::new("id".to_string(), "BTREE".to_string());
6967 create_index_request.id = Some(vec![table_name.to_string()]);
6968 create_index_request.name = Some(index_name.to_string());
6969 namespace
6970 .create_table_scalar_index(create_index_request)
6971 .await
6972 .unwrap()
6973 .transaction_id
6974 }
6975
6976 async fn create_branch_with_commits(
6981 namespace: &DirectoryNamespace,
6982 table_name: &str,
6983 branch_name: &str,
6984 extra_versions: usize,
6985 ) -> String {
6986 let mut main = open_dataset(namespace, table_name).await;
6987 let fork_version = main.version().version;
6988 let branch = main
6989 .create_branch(branch_name, fork_version, None)
6990 .await
6991 .unwrap();
6992 let branch_uri = branch.uri().to_string();
6993 for i in 0..extra_versions {
6994 append_scalar_version(&branch_uri, (i as i32 + 1) * 100).await;
6995 }
6996 branch_uri
6997 }
6998
6999 async fn append_scalar_version(uri: &str, seed: i32) {
7002 use arrow::array::{Int32Array, StringArray};
7003 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7004 let schema = Arc::new(ArrowSchema::new(vec![
7005 Field::new("id", DataType::Int32, false),
7006 Field::new("name", DataType::Utf8, true),
7007 ]));
7008 let batch = arrow::record_batch::RecordBatch::try_new(
7009 schema.clone(),
7010 vec![
7011 Arc::new(Int32Array::from(vec![seed, seed + 1])),
7012 Arc::new(StringArray::from(vec![Some("x"), Some("y")])),
7013 ],
7014 )
7015 .unwrap();
7016 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
7017 Dataset::write(
7018 reader,
7019 uri,
7020 Some(WriteParams {
7021 mode: WriteMode::Append,
7022 ..Default::default()
7023 }),
7024 )
7025 .await
7026 .unwrap();
7027 }
7028
7029 async fn list_versions(
7031 namespace: &DirectoryNamespace,
7032 table_name: &str,
7033 branch: Option<&str>,
7034 ) -> Result<Vec<TableVersion>> {
7035 let req = ListTableVersionsRequest {
7036 id: Some(vec![table_name.to_string()]),
7037 branch: branch.map(|b| b.to_string()),
7038 ..Default::default()
7039 };
7040 namespace.list_table_versions(req).await.map(|r| r.versions)
7041 }
7042
7043 #[tokio::test]
7044 async fn test_list_table_versions_on_branch() {
7045 let (namespace, _temp_dir) = create_test_namespace().await;
7046 create_scalar_table(&namespace, "users").await;
7047 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7048
7049 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7052 .await
7053 .unwrap();
7054 assert!(branch_versions.len() >= 2);
7055 assert!(
7056 branch_versions
7057 .iter()
7058 .all(|v| v.manifest_path.contains("tree/exp")),
7059 "branch versions must resolve to branch manifests: {:?}",
7060 branch_versions
7061 );
7062
7063 let main_versions = list_versions(&namespace, "users", None).await.unwrap();
7065 let main_explicit = list_versions(&namespace, "users", Some("main"))
7066 .await
7067 .unwrap();
7068 assert_eq!(main_versions.len(), main_explicit.len());
7069 assert!(
7070 main_versions
7071 .iter()
7072 .all(|v| !v.manifest_path.contains("tree/"))
7073 );
7074
7075 let missing = list_versions(&namespace, "users", Some("does-not-exist")).await;
7077 assert!(missing.is_err());
7078 assert!(missing.unwrap_err().to_string().contains("not found"));
7079 }
7080
7081 #[tokio::test]
7082 async fn test_describe_table_version_on_branch() {
7083 let (namespace, _temp_dir) = create_test_namespace().await;
7084 create_scalar_table(&namespace, "users").await;
7085 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7086
7087 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7088 .await
7089 .unwrap();
7090 let latest = branch_versions.iter().map(|v| v.version).max().unwrap();
7091
7092 let req = DescribeTableVersionRequest {
7094 id: Some(vec!["users".to_string()]),
7095 branch: Some("exp".to_string()),
7096 ..Default::default()
7097 };
7098 let resp = namespace.describe_table_version(req).await.unwrap();
7099 assert_eq!(resp.version.version, latest);
7100 assert!(resp.version.manifest_path.contains("tree/exp"));
7101
7102 let req = DescribeTableVersionRequest {
7104 id: Some(vec!["users".to_string()]),
7105 version: Some(latest),
7106 branch: Some("exp".to_string()),
7107 ..Default::default()
7108 };
7109 assert!(namespace.describe_table_version(req).await.is_ok());
7110
7111 let req = DescribeTableVersionRequest {
7113 id: Some(vec!["users".to_string()]),
7114 version: Some(999_999),
7115 branch: Some("exp".to_string()),
7116 ..Default::default()
7117 };
7118 assert!(namespace.describe_table_version(req).await.is_err());
7119
7120 let req = DescribeTableVersionRequest {
7122 id: Some(vec!["users".to_string()]),
7123 branch: Some("nope".to_string()),
7124 ..Default::default()
7125 };
7126 let err = namespace.describe_table_version(req).await;
7127 assert!(err.is_err() && err.unwrap_err().to_string().contains("not found"));
7128 }
7129
7130 #[tokio::test]
7131 async fn test_restore_table_on_branch() {
7132 use lance_namespace::models::RestoreTableRequest;
7133
7134 let (namespace, _temp_dir) = create_test_namespace().await;
7135 create_scalar_table(&namespace, "users").await;
7136 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7137
7138 let before = list_versions(&namespace, "users", Some("exp"))
7139 .await
7140 .unwrap();
7141 let branch_latest = before.iter().map(|v| v.version).max().unwrap();
7142 let earliest = before.iter().map(|v| v.version).min().unwrap();
7143 let main_before = list_versions(&namespace, "users", None)
7144 .await
7145 .unwrap()
7146 .len();
7147
7148 let req = RestoreTableRequest {
7151 id: Some(vec!["users".to_string()]),
7152 version: earliest,
7153 branch: Some("exp".to_string()),
7154 ..Default::default()
7155 };
7156 let resp = namespace.restore_table(req).await.unwrap();
7157 assert!(resp.transaction_id.is_some());
7158
7159 let after = list_versions(&namespace, "users", Some("exp"))
7160 .await
7161 .unwrap();
7162 let new_latest = after.iter().map(|v| v.version).max().unwrap();
7163 assert!(
7164 new_latest > branch_latest,
7165 "restore should add a branch version"
7166 );
7167
7168 let main_after = list_versions(&namespace, "users", None)
7169 .await
7170 .unwrap()
7171 .len();
7172 assert_eq!(main_after, main_before, "main must be unaffected");
7173 }
7174
7175 #[tokio::test]
7176 async fn test_batch_delete_table_versions_on_branch() {
7177 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7178
7179 let (namespace, _temp_dir) = create_test_namespace().await;
7180 create_scalar_table(&namespace, "users").await;
7181 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7182
7183 let before = list_versions(&namespace, "users", Some("exp"))
7184 .await
7185 .unwrap();
7186 let main_before = list_versions(&namespace, "users", None).await.unwrap();
7187
7188 let req = BatchDeleteTableVersionsRequest {
7193 id: Some(vec!["users".to_string()]),
7194 branch: Some("exp".to_string()),
7195 ranges: vec![VersionRange::new(0, -1)],
7196 ..Default::default()
7197 };
7198 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7199 assert_eq!(
7200 resp.deleted_count,
7201 Some(before.len() as i64),
7202 "every branch manifest should be physically deleted"
7203 );
7204
7205 assert!(
7207 list_versions(&namespace, "users", Some("exp"))
7208 .await
7209 .is_err()
7210 );
7211 let main_after = list_versions(&namespace, "users", None).await.unwrap();
7212 assert_eq!(
7213 main_after.len(),
7214 main_before.len(),
7215 "main must be untouched"
7216 );
7217 }
7218
7219 #[tokio::test]
7220 async fn test_create_table_version_on_branch() {
7221 use futures::TryStreamExt;
7222 use lance_namespace::models::CreateTableVersionRequest;
7223
7224 let (namespace, _temp_dir) = create_test_namespace().await;
7225 create_scalar_table(&namespace, "users").await;
7226 let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 1).await;
7227
7228 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7230 let versions_dir = branch_ds.versions_dir();
7231 let store = branch_ds.object_store(None).await.unwrap();
7232 let existing = store
7233 .inner
7234 .list(Some(&versions_dir))
7235 .try_collect::<Vec<_>>()
7236 .await
7237 .unwrap()
7238 .into_iter()
7239 .find(|m| {
7240 m.location
7241 .filename()
7242 .map(|f| f.ends_with(".manifest"))
7243 .unwrap_or(false)
7244 })
7245 .expect("a branch manifest");
7246 let bytes = store
7247 .inner
7248 .get(&existing.location)
7249 .await
7250 .unwrap()
7251 .bytes()
7252 .await
7253 .unwrap();
7254 let staging = versions_dir.join("staging_manifest");
7255 store.inner.put(&staging, bytes.into()).await.unwrap();
7256
7257 let main_before = list_versions(&namespace, "users", None)
7258 .await
7259 .unwrap()
7260 .len();
7261 let new_version = list_versions(&namespace, "users", Some("exp"))
7262 .await
7263 .unwrap()
7264 .iter()
7265 .map(|v| v.version)
7266 .max()
7267 .unwrap()
7268 + 1;
7269
7270 let req = CreateTableVersionRequest {
7271 id: Some(vec!["users".to_string()]),
7272 version: new_version,
7273 manifest_path: staging.to_string(),
7274 naming_scheme: Some("V2".to_string()),
7275 branch: Some("exp".to_string()),
7276 ..Default::default()
7277 };
7278 let resp = namespace.create_table_version(req).await.unwrap();
7279 let info = resp.version.expect("version info");
7280 assert!(
7282 info.manifest_path.contains("tree/exp"),
7283 "got {}",
7284 info.manifest_path
7285 );
7286
7287 let after = list_versions(&namespace, "users", Some("exp"))
7289 .await
7290 .unwrap();
7291 assert!(after.iter().any(|v| v.version == new_version));
7292 let main_after = list_versions(&namespace, "users", None)
7293 .await
7294 .unwrap()
7295 .len();
7296 assert_eq!(main_after, main_before, "main must be unaffected");
7297 }
7298
7299 #[tokio::test]
7304 async fn test_external_manifest_store_resolves_branch_from_base_path() {
7305 use futures::TryStreamExt;
7306 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
7307 use lance_table::io::commit::external_manifest::ExternalManifestStore;
7308
7309 let (namespace, _temp_dir) = create_test_namespace().await;
7310 create_scalar_table(&namespace, "users").await; let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 2).await;
7312
7313 let namespace = Arc::new(namespace);
7314 let table_id = vec!["users".to_string()];
7315 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7316 let branch_base = branch_ds.branch_location().path;
7317 let root_base = branch_ds.branch_location().find_main().unwrap().path;
7318 let store = LanceNamespaceExternalManifestStore::new(
7319 namespace.clone(),
7320 table_id.clone(),
7321 root_base.clone(),
7322 );
7323
7324 let (branch_latest, branch_path) = store
7327 .get_latest_version(branch_base.as_ref())
7328 .await
7329 .unwrap()
7330 .expect("branch has versions");
7331 let (_main_latest, main_path) = store
7332 .get_latest_version(root_base.as_ref())
7333 .await
7334 .unwrap()
7335 .expect("main has versions");
7336 assert!(
7337 branch_path.contains("tree/exp"),
7338 "branch latest must resolve to the branch tree: {}",
7339 branch_path
7340 );
7341 assert!(
7342 !main_path.contains("tree/exp"),
7343 "main latest must not resolve to a branch tree: {}",
7344 main_path
7345 );
7346
7347 let described = store
7349 .get(branch_base.as_ref(), branch_latest)
7350 .await
7351 .unwrap();
7352 assert!(
7353 described.contains("tree/exp"),
7354 "describe on the branch must resolve to the branch tree: {}",
7355 described
7356 );
7357
7358 assert!(store.get_latest_version("somewhere/else").await.is_err());
7360
7361 let versions_dir = branch_ds.versions_dir();
7364 let obj = branch_ds.object_store(None).await.unwrap();
7365 let existing = obj
7366 .inner
7367 .list(Some(&versions_dir))
7368 .try_collect::<Vec<_>>()
7369 .await
7370 .unwrap()
7371 .into_iter()
7372 .find(|m| {
7373 m.location
7374 .filename()
7375 .map(|f| f.ends_with(".manifest"))
7376 .unwrap_or(false)
7377 })
7378 .expect("a branch manifest");
7379 let bytes = obj
7380 .inner
7381 .get(&existing.location)
7382 .await
7383 .unwrap()
7384 .bytes()
7385 .await
7386 .unwrap();
7387 let size = bytes.len() as u64;
7388 let staging = versions_dir.clone().join("staging_manifest");
7389 obj.inner.put(&staging, bytes.into()).await.unwrap();
7390
7391 let committed = store
7392 .put(
7393 &branch_base,
7394 branch_latest + 1,
7395 &staging,
7396 size,
7397 None,
7398 obj.inner.as_ref(),
7399 ManifestNamingScheme::V2,
7400 )
7401 .await
7402 .unwrap();
7403 assert!(
7404 committed.path.to_string().contains("tree/exp"),
7405 "a commit through a branch-qualified base must land on the branch tree: {}",
7406 committed.path
7407 );
7408 }
7409
7410 #[tokio::test]
7418 async fn test_write_into_namespace_on_branch_appends_to_branch() {
7419 use lance::dataset::builder::DatasetBuilder;
7420 use lance_namespace::models::CreateTableBranchRequest;
7421
7422 let temp = TempStdDir::default();
7423 let namespace = Arc::new(
7424 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
7425 .manifest_enabled(true)
7426 .table_version_tracking_enabled(true)
7427 .ops_metrics_enabled(true)
7428 .build()
7429 .await
7430 .unwrap(),
7431 );
7432 let ns: Arc<dyn LanceNamespace> = namespace.clone();
7433 let table_id = vec!["t".to_string()];
7434 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
7436 id: Some(table_id.clone()),
7437 name: "exp".to_string(),
7438 ..Default::default()
7439 })
7440 .await
7441 .unwrap();
7442
7443 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
7444 ns.list_table_versions(ListTableVersionsRequest {
7445 id: Some(table_id),
7446 ..Default::default()
7447 })
7448 .await
7449 .unwrap()
7450 .versions
7451 .len()
7452 };
7453 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
7454 let commits_before = namespace
7455 .retrieve_ops_metrics()
7456 .get("create_table_version")
7457 .copied()
7458 .unwrap_or(0);
7459
7460 let branch_ds = Dataset::write_into_namespace_on_branch(
7461 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
7462 ns.clone(),
7463 table_id.clone(),
7464 "exp",
7465 Some(WriteParams {
7466 mode: WriteMode::Append,
7467 ..Default::default()
7468 }),
7469 )
7470 .await
7471 .unwrap();
7472 assert_eq!(branch_ds.manifest.branch.as_deref(), Some("exp"));
7473 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
7474
7475 let commits_after = namespace
7478 .retrieve_ops_metrics()
7479 .get("create_table_version")
7480 .copied()
7481 .unwrap_or(0);
7482 assert_eq!(
7483 commits_after,
7484 commits_before + 1,
7485 "the branch append must register its version via create_table_version"
7486 );
7487 let exp_versions = ns
7488 .list_table_versions(ListTableVersionsRequest {
7489 id: Some(table_id.clone()),
7490 branch: Some("exp".to_string()),
7491 ..Default::default()
7492 })
7493 .await
7494 .unwrap()
7495 .versions;
7496 assert!(
7497 exp_versions
7498 .iter()
7499 .all(|v| v.manifest_path.contains("tree/exp")),
7500 "branch versions must resolve to the branch tree: {:?}",
7501 exp_versions
7502 );
7503 assert_eq!(
7504 main_chain_len(ns.clone(), table_id.clone()).await,
7505 main_before,
7506 "main's catalog must be untouched by the branch append"
7507 );
7508
7509 Dataset::write_into_namespace(
7512 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
7513 ns.clone(),
7514 table_id.clone(),
7515 Some(WriteParams {
7516 mode: WriteMode::Append,
7517 ..Default::default()
7518 }),
7519 )
7520 .await
7521 .unwrap();
7522 assert_eq!(
7523 main_chain_len(ns.clone(), table_id.clone()).await,
7524 main_before + 1,
7525 "a managed main append must register its version in the catalog"
7526 );
7527 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
7528 .await
7529 .unwrap()
7530 .load()
7531 .await
7532 .unwrap();
7533 assert_eq!(
7534 scan_id_column(&fresh).await,
7535 vec![1, 2, 100],
7536 "a fresh managed open must resolve the appended version, not a stale latest"
7537 );
7538 }
7539
7540 #[tokio::test]
7542 async fn test_write_into_namespace_on_branch_rejects_create() {
7543 use arrow::array::{Int32Array, StringArray};
7544 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7545
7546 let (namespace, _temp_dir) = create_test_namespace().await;
7547 let namespace = Arc::new(namespace);
7548
7549 let schema = Arc::new(ArrowSchema::new(vec![
7550 Field::new("id", DataType::Int32, false),
7551 Field::new("name", DataType::Utf8, true),
7552 ]));
7553 let batch = arrow::record_batch::RecordBatch::try_new(
7554 schema.clone(),
7555 vec![
7556 Arc::new(Int32Array::from(vec![1])),
7557 Arc::new(StringArray::from(vec![Some("a")])),
7558 ],
7559 )
7560 .unwrap();
7561 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
7562
7563 let result = Dataset::write_into_namespace_on_branch(
7564 reader,
7565 namespace.clone(),
7566 vec!["new_table".to_string()],
7567 "exp",
7568 Some(WriteParams {
7569 mode: WriteMode::Create,
7570 ..Default::default()
7571 }),
7572 )
7573 .await;
7574 assert!(result.is_err(), "create on a branch must be rejected");
7575 assert!(
7576 result.unwrap_err().to_string().contains("branch"),
7577 "error should mention the branch restriction"
7578 );
7579 }
7580
7581 #[tokio::test]
7582 async fn test_branch_name_validation_rejects_traversal() {
7583 let (namespace, _temp_dir) = create_test_namespace().await;
7584 create_scalar_table(&namespace, "users").await;
7585
7586 let err = list_versions(&namespace, "users", Some("../evil")).await;
7589 assert!(err.is_err());
7590 assert!(err.unwrap_err().to_string().contains("invalid branch name"));
7591 }
7592
7593 #[tokio::test]
7594 async fn test_branch_ops_reject_zombie_branch() {
7595 use futures::TryStreamExt;
7596 use lance_namespace::models::{
7597 BatchDeleteTableVersionsRequest, CreateTableVersionRequest, RestoreTableRequest,
7598 VersionRange,
7599 };
7600
7601 let (namespace, _temp_dir) = create_test_namespace().await;
7602 create_scalar_table(&namespace, "users").await;
7603
7604 let dataset = open_dataset(&namespace, "users").await;
7605 let store = dataset.object_store(None).await.unwrap();
7606 let manifest = store
7607 .inner
7608 .list(Some(&dataset.versions_dir()))
7609 .try_collect::<Vec<_>>()
7610 .await
7611 .unwrap()
7612 .into_iter()
7613 .find(|m| {
7614 m.location
7615 .filename()
7616 .map(|f| f.ends_with(".manifest"))
7617 .unwrap_or(false)
7618 })
7619 .expect("a manifest");
7620 let bytes = store
7621 .inner
7622 .get(&manifest.location)
7623 .await
7624 .unwrap()
7625 .bytes()
7626 .await
7627 .unwrap();
7628 let zombie = dataset
7629 .branch_location()
7630 .find_branch(Some("ghost"))
7631 .unwrap()
7632 .path
7633 .join(VERSIONS_DIR)
7634 .join(manifest.location.filename().unwrap());
7635 store.inner.put(&zombie, bytes.into()).await.unwrap();
7636
7637 assert!(dataset.branches().get("ghost").await.is_err());
7638
7639 fn rejected<T: std::fmt::Debug>(label: &str, r: Result<T>) {
7640 match r {
7641 Ok(v) => panic!("{label} must reject the zombie branch, got Ok({v:?})"),
7642 Err(e) => assert!(e.to_string().contains("not found"), "{label}: {e}"),
7643 }
7644 }
7645
7646 rejected(
7647 "list",
7648 list_versions(&namespace, "users", Some("ghost")).await,
7649 );
7650 rejected(
7651 "describe",
7652 namespace
7653 .describe_table_version(DescribeTableVersionRequest {
7654 id: Some(vec!["users".to_string()]),
7655 branch: Some("ghost".to_string()),
7656 ..Default::default()
7657 })
7658 .await,
7659 );
7660 rejected(
7661 "create",
7662 namespace
7663 .create_table_version(CreateTableVersionRequest {
7664 id: Some(vec!["users".to_string()]),
7665 version: 2,
7666 manifest_path: zombie.to_string(),
7667 branch: Some("ghost".to_string()),
7668 ..Default::default()
7669 })
7670 .await,
7671 );
7672 rejected(
7673 "restore",
7674 namespace
7675 .restore_table(RestoreTableRequest {
7676 id: Some(vec!["users".to_string()]),
7677 version: 1,
7678 branch: Some("ghost".to_string()),
7679 ..Default::default()
7680 })
7681 .await,
7682 );
7683 rejected(
7684 "batch_delete",
7685 namespace
7686 .batch_delete_table_versions(BatchDeleteTableVersionsRequest {
7687 id: Some(vec!["users".to_string()]),
7688 branch: Some("ghost".to_string()),
7689 ranges: vec![VersionRange::new(1, 1)],
7690 ..Default::default()
7691 })
7692 .await,
7693 );
7694 }
7695
7696 #[tokio::test]
7701 async fn test_batch_delete_table_versions_main_v2() {
7702 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7703
7704 let (namespace, _temp_dir) = create_test_namespace().await;
7705 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
7707 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
7711 assert!(before.len() >= 3);
7712 assert!(
7715 before
7716 .iter()
7717 .all(|v| v.manifest_path.rsplit('/').next().unwrap().len() == 29),
7718 "expected V2-named manifests: {:?}",
7719 before
7720 );
7721 let min_v = before.iter().map(|v| v.version).min().unwrap();
7722 let max_v = before.iter().map(|v| v.version).max().unwrap();
7723
7724 let req = BatchDeleteTableVersionsRequest {
7727 id: Some(vec!["users".to_string()]),
7728 ranges: vec![VersionRange::new(min_v, max_v)],
7729 ..Default::default()
7730 };
7731 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7732 assert_eq!(
7733 resp.deleted_count,
7734 Some((before.len() - 1) as i64),
7735 "V2 manifests must actually be deleted (was 0 before the fix)"
7736 );
7737
7738 let after = list_versions(&namespace, "users", None).await.unwrap();
7739 assert_eq!(after.len(), 1);
7740 assert_eq!(after[0].version, max_v);
7741 }
7742
7743 #[tokio::test]
7745 async fn test_batch_delete_end_is_exclusive() {
7746 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7747
7748 let (namespace, _temp_dir) = create_test_namespace().await;
7749 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
7751 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
7755 let min_v = before.iter().map(|v| v.version).min().unwrap();
7756
7757 let req = BatchDeleteTableVersionsRequest {
7758 id: Some(vec!["users".to_string()]),
7759 ranges: vec![VersionRange::new(min_v, min_v + 1)],
7760 ..Default::default()
7761 };
7762 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7763 assert_eq!(
7764 resp.deleted_count,
7765 Some(1),
7766 "only min_v is in [min_v, min_v+1)"
7767 );
7768
7769 let after = list_versions(&namespace, "users", None).await.unwrap();
7770 assert!(
7771 !after.iter().any(|v| v.version == min_v),
7772 "min_v must be deleted"
7773 );
7774 assert_eq!(after.len(), before.len() - 1, "exactly one version removed");
7775 }
7776
7777 #[tokio::test]
7778 async fn test_batch_delete_rejects_unbounded_range() {
7779 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7780
7781 let (namespace, _temp_dir) = create_test_namespace().await;
7782 create_scalar_table(&namespace, "users").await;
7783
7784 let req = BatchDeleteTableVersionsRequest {
7787 id: Some(vec!["users".to_string()]),
7788 ranges: vec![VersionRange::new(0, i64::MAX)],
7789 ..Default::default()
7790 };
7791 let err = namespace.batch_delete_table_versions(req).await;
7792 assert!(err.is_err());
7793 assert!(
7794 err.unwrap_err().to_string().contains("limit"),
7795 "expected a range-too-large error"
7796 );
7797 }
7798
7799 async fn create_managed_namespace(path: &str) -> Arc<dyn LanceNamespace> {
7801 Arc::new(
7802 DirectoryNamespaceBuilder::new(path)
7803 .manifest_enabled(true)
7804 .table_version_tracking_enabled(true)
7805 .build()
7806 .await
7807 .unwrap(),
7808 )
7809 }
7810
7811 fn single_int_schema() -> Arc<arrow::datatypes::Schema> {
7812 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7813 Arc::new(ArrowSchema::new(vec![Field::new(
7814 "id",
7815 DataType::Int32,
7816 false,
7817 )]))
7818 }
7819
7820 fn single_int_batch(seed: i32) -> arrow::record_batch::RecordBatch {
7821 use arrow::array::Int32Array;
7822 arrow::record_batch::RecordBatch::try_new(
7823 single_int_schema(),
7824 vec![Arc::new(Int32Array::from(vec![seed]))],
7825 )
7826 .unwrap()
7827 }
7828
7829 async fn create_managed_table(ns: &Arc<dyn LanceNamespace>, table_id: &[String]) -> Dataset {
7832 let mut ds = Dataset::write_into_namespace(
7833 RecordBatchIterator::new(vec![Ok(single_int_batch(1))], single_int_schema()),
7834 ns.clone(),
7835 table_id.to_vec(),
7836 Some(WriteParams {
7837 mode: WriteMode::Create,
7838 ..Default::default()
7839 }),
7840 )
7841 .await
7842 .unwrap();
7843 ds.append(
7844 RecordBatchIterator::new(vec![Ok(single_int_batch(2))], single_int_schema()),
7845 None,
7846 )
7847 .await
7848 .unwrap();
7849 ds
7850 }
7851
7852 async fn scan_id_column(ds: &Dataset) -> Vec<i32> {
7854 use arrow::array::Int32Array;
7855 use futures::TryStreamExt;
7856 let batches: Vec<arrow::record_batch::RecordBatch> = ds
7857 .scan()
7858 .try_into_stream()
7859 .await
7860 .unwrap()
7861 .try_collect()
7862 .await
7863 .unwrap();
7864 let mut ids: Vec<i32> = batches
7865 .iter()
7866 .flat_map(|b| {
7867 b.column(0)
7868 .as_any()
7869 .downcast_ref::<Int32Array>()
7870 .unwrap()
7871 .values()
7872 .to_vec()
7873 })
7874 .collect();
7875 ids.sort();
7876 ids
7877 }
7878
7879 #[tokio::test]
7884 async fn test_managed_branch_open_and_commit() {
7885 use futures::TryStreamExt;
7886 use lance::dataset::builder::DatasetBuilder;
7887 use lance_namespace::models::CreateTableBranchRequest;
7888
7889 let temp = TempStdDir::default();
7890 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
7891 let table_id = vec!["t".to_string()];
7892 create_managed_table(&ns, &table_id).await;
7893 let main_before = ns
7894 .list_table_versions(ListTableVersionsRequest {
7895 id: Some(table_id.clone()),
7896 ..Default::default()
7897 })
7898 .await
7899 .unwrap()
7900 .versions
7901 .len();
7902
7903 ns.create_table_branch(CreateTableBranchRequest {
7906 id: Some(table_id.clone()),
7907 name: "exp".to_string(),
7908 ..Default::default()
7909 })
7910 .await
7911 .unwrap();
7912
7913 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
7916 .await
7917 .unwrap()
7918 .with_branch("exp", None)
7919 .load()
7920 .await
7921 .unwrap();
7922 assert_eq!(
7923 branch_ds.manifest.branch.as_deref(),
7924 Some("exp"),
7925 "with_branch on a managed table must open the branch chain"
7926 );
7927 let branch_base = branch_ds.branch_location().path;
7928 assert!(
7929 branch_base.as_ref().ends_with("tree/exp"),
7930 "the branch dataset must be rooted at the branch chain: {}",
7931 branch_base
7932 );
7933 let branch_v_before = branch_ds.version().version;
7934
7935 branch_ds
7937 .append(
7938 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
7939 None,
7940 )
7941 .await
7942 .unwrap();
7943 assert_eq!(
7944 branch_ds.manifest.branch.as_deref(),
7945 Some("exp"),
7946 "the commit must stay on the branch"
7947 );
7948 assert!(
7949 branch_ds.version().version > branch_v_before,
7950 "the branch version must advance after the commit"
7951 );
7952 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
7953
7954 let store = branch_ds.object_store(None).await.unwrap();
7958 let branch_data = branch_base.clone().join("data");
7959 let branch_files = store
7960 .inner
7961 .list(Some(&branch_data))
7962 .try_collect::<Vec<_>>()
7963 .await
7964 .unwrap();
7965 assert!(
7966 !branch_files.is_empty(),
7967 "the branch commit must place data files under the branch chain"
7968 );
7969
7970 let table_uri = ns
7972 .describe_table(DescribeTableRequest {
7973 id: Some(table_id.clone()),
7974 ..Default::default()
7975 })
7976 .await
7977 .unwrap()
7978 .location
7979 .unwrap();
7980 let fs_branch_ds = DatasetBuilder::from_uri(&table_uri)
7981 .with_branch("exp", None)
7982 .load()
7983 .await
7984 .unwrap();
7985 assert_eq!(fs_branch_ds.manifest.branch.as_deref(), Some("exp"));
7986 assert_eq!(scan_id_column(&fs_branch_ds).await, vec![1, 2, 3]);
7987
7988 let main_after = ns
7991 .list_table_versions(ListTableVersionsRequest {
7992 id: Some(table_id.clone()),
7993 ..Default::default()
7994 })
7995 .await
7996 .unwrap()
7997 .versions
7998 .len();
7999 assert_eq!(
8000 main_after, main_before,
8001 "committing on the branch must not change main's chain"
8002 );
8003 let main_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8004 .await
8005 .unwrap()
8006 .load()
8007 .await
8008 .unwrap();
8009 assert_eq!(main_ds.manifest.branch, None);
8010 assert_eq!(scan_id_column(&main_ds).await, vec![1, 2]);
8011 }
8012
8013 #[tokio::test]
8018 async fn test_managed_branch_tags() {
8019 use lance::dataset::builder::DatasetBuilder;
8020 use lance::dataset::refs::Ref;
8021 use lance_namespace::models::CreateTableBranchRequest;
8022
8023 let temp = TempStdDir::default();
8024 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8025 let table_id = vec!["t".to_string()];
8026 let main_ds = create_managed_table(&ns, &table_id).await;
8027 ns.create_table_branch(CreateTableBranchRequest {
8028 id: Some(table_id.clone()),
8029 name: "exp".to_string(),
8030 ..Default::default()
8031 })
8032 .await
8033 .unwrap();
8034 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8035 .await
8036 .unwrap()
8037 .with_branch("exp", None)
8038 .load()
8039 .await
8040 .unwrap();
8041 branch_ds
8042 .append(
8043 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8044 None,
8045 )
8046 .await
8047 .unwrap();
8048 let branch_version = branch_ds.version().version;
8049
8050 main_ds
8053 .tags()
8054 .create("exp-tag", ("exp", Some(branch_version)))
8055 .await
8056 .unwrap();
8057 let tag = main_ds.tags().get("exp-tag").await.unwrap();
8058 assert_eq!(tag.branch.as_deref(), Some("exp"));
8059 assert_eq!(tag.version, branch_version);
8060
8061 branch_ds
8063 .tags()
8064 .create("exp-tag2", branch_version)
8065 .await
8066 .unwrap();
8067 let tag2 = branch_ds.tags().get("exp-tag2").await.unwrap();
8068 assert_eq!(tag2.branch.as_deref(), Some("exp"));
8069
8070 let tag_open = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8073 .await
8074 .unwrap()
8075 .with_tag("exp-tag")
8076 .load()
8077 .await
8078 .unwrap();
8079 assert_eq!(tag_open.manifest.branch.as_deref(), Some("exp"));
8080 assert_eq!(tag_open.version().version, branch_version);
8081 assert_eq!(scan_id_column(&tag_open).await, vec![1, 2, 3]);
8082
8083 let tag_checkout = main_ds
8085 .checkout_version(Ref::Tag("exp-tag".to_string()))
8086 .await
8087 .unwrap();
8088 assert_eq!(tag_checkout.manifest.branch.as_deref(), Some("exp"));
8089 assert_eq!(scan_id_column(&tag_checkout).await, vec![1, 2, 3]);
8090
8091 let err = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8093 .await
8094 .unwrap()
8095 .with_tag("no-such-tag")
8096 .load()
8097 .await;
8098 assert!(err.is_err(), "a missing tag must error");
8099 }
8100
8101 #[tokio::test]
8106 async fn test_managed_cross_branch_checkout() {
8107 use lance::dataset::builder::DatasetBuilder;
8108 use lance::dataset::refs::Ref;
8109 use lance_namespace::models::CreateTableBranchRequest;
8110
8111 let temp = TempStdDir::default();
8112 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8113 let table_id = vec!["t".to_string()];
8114 let mut main_ds = create_managed_table(&ns, &table_id).await;
8115 ns.create_table_branch(CreateTableBranchRequest {
8116 id: Some(table_id.clone()),
8117 name: "exp".to_string(),
8118 ..Default::default()
8119 })
8120 .await
8121 .unwrap();
8122
8123 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8125 .await
8126 .unwrap()
8127 .with_branch("exp", None)
8128 .load()
8129 .await
8130 .unwrap();
8131 branch_ds
8132 .append(
8133 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8134 None,
8135 )
8136 .await
8137 .unwrap();
8138 let overlap_version = branch_ds.version().version;
8139 while main_ds.version().version < overlap_version {
8140 main_ds
8141 .append(
8142 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
8143 None,
8144 )
8145 .await
8146 .unwrap();
8147 }
8148
8149 let on_branch = main_ds
8152 .checkout_version(Ref::Version(Some("exp".to_string()), Some(overlap_version)))
8153 .await
8154 .unwrap();
8155 assert_eq!(on_branch.manifest.branch.as_deref(), Some("exp"));
8156 assert_eq!(scan_id_column(&on_branch).await, vec![1, 2, 3]);
8157
8158 let mut on_branch_latest = main_ds.checkout_branch("exp").await.unwrap();
8160 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8161 assert_eq!(on_branch_latest.version().version, overlap_version);
8162
8163 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
8166 ns.list_table_versions(ListTableVersionsRequest {
8167 id: Some(table_id),
8168 ..Default::default()
8169 })
8170 .await
8171 .unwrap()
8172 .versions
8173 .len()
8174 };
8175 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
8176 on_branch_latest
8177 .append(
8178 RecordBatchIterator::new(vec![Ok(single_int_batch(4))], single_int_schema()),
8179 None,
8180 )
8181 .await
8182 .unwrap();
8183 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8184 assert_eq!(scan_id_column(&on_branch_latest).await, vec![1, 2, 3, 4]);
8185 assert_eq!(
8186 main_chain_len(ns.clone(), table_id.clone()).await,
8187 main_before,
8188 "a commit on the checked-out branch must not advance main's chain"
8189 );
8190
8191 let on_main = branch_ds
8193 .checkout_version(Ref::Version(None, Some(1)))
8194 .await
8195 .unwrap();
8196 assert_eq!(on_main.manifest.branch, None);
8197 assert_eq!(scan_id_column(&on_main).await, vec![1]);
8198
8199 ns.create_table_branch(CreateTableBranchRequest {
8201 id: Some(table_id.clone()),
8202 name: "exp2".to_string(),
8203 ..Default::default()
8204 })
8205 .await
8206 .unwrap();
8207 let on_branch2 = branch_ds.checkout_branch("exp2").await.unwrap();
8208 assert_eq!(on_branch2.manifest.branch.as_deref(), Some("exp2"));
8209
8210 let err = main_ds
8212 .checkout_version(Ref::Version(Some("exp".to_string()), Some(999)))
8213 .await;
8214 assert!(err.is_err(), "a version missing from the branch must error");
8215 }
8216
8217 #[tokio::test]
8223 async fn test_commit_builder_honors_explicit_handler_for_dataset_dest() {
8224 use lance::dataset::write::{CommitBuilder, InsertBuilder};
8225 use lance::dataset::{WriteDestination, builder::DatasetBuilder};
8226 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
8227 use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
8228
8229 let temp = TempStdDir::default();
8230 let namespace = Arc::new(
8231 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
8232 .manifest_enabled(true)
8233 .table_version_tracking_enabled(true)
8234 .ops_metrics_enabled(true)
8235 .build()
8236 .await
8237 .unwrap(),
8238 );
8239 let ns: Arc<dyn LanceNamespace> = namespace.clone();
8240 let table_id = vec!["t".to_string()];
8241 create_managed_table(&ns, &table_id).await; let table_uri = ns
8245 .describe_table(DescribeTableRequest {
8246 id: Some(table_id.clone()),
8247 ..Default::default()
8248 })
8249 .await
8250 .unwrap()
8251 .location
8252 .unwrap();
8253 let plain_ds = Arc::new(Dataset::open(&table_uri).await.unwrap());
8254
8255 let transaction = InsertBuilder::new(WriteDestination::Dataset(plain_ds.clone()))
8256 .with_params(&WriteParams {
8257 mode: WriteMode::Append,
8258 ..Default::default()
8259 })
8260 .execute_uncommitted(vec![single_int_batch(3)])
8261 .await
8262 .unwrap();
8263
8264 let handler = Arc::new(ExternalManifestCommitHandler {
8265 external_manifest_store: Arc::new(
8266 LanceNamespaceExternalManifestStore::for_table_uri(
8267 ns.clone(),
8268 table_id.clone(),
8269 &table_uri,
8270 )
8271 .unwrap(),
8272 ),
8273 });
8274 let commits_before = namespace
8275 .retrieve_ops_metrics()
8276 .get("create_table_version")
8277 .copied()
8278 .unwrap_or(0);
8279 let committed = CommitBuilder::new(WriteDestination::Dataset(plain_ds))
8280 .with_commit_handler(handler)
8281 .execute(transaction)
8282 .await
8283 .unwrap();
8284 assert_eq!(scan_id_column(&committed).await, vec![1, 2, 3]);
8285
8286 let commits_after = namespace
8287 .retrieve_ops_metrics()
8288 .get("create_table_version")
8289 .copied()
8290 .unwrap_or(0);
8291 assert_eq!(
8292 commits_after,
8293 commits_before + 1,
8294 "the explicit handler must route the commit through create_table_version"
8295 );
8296 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8297 .await
8298 .unwrap()
8299 .load()
8300 .await
8301 .unwrap();
8302 assert_eq!(
8303 scan_id_column(&fresh).await,
8304 vec![1, 2, 3],
8305 "a fresh managed open must resolve the committed version"
8306 );
8307 }
8308
8309 #[tokio::test]
8311 async fn test_managed_branch_from_non_latest_fork() {
8312 use lance::dataset::builder::DatasetBuilder;
8313 use lance_namespace::models::CreateTableBranchRequest;
8314
8315 let temp = TempStdDir::default();
8316 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8317 let table_id = vec!["t".to_string()];
8318 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
8321 id: Some(table_id.clone()),
8322 name: "old".to_string(),
8323 from_version: Some(1),
8324 ..Default::default()
8325 })
8326 .await
8327 .unwrap();
8328
8329 let old_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8330 .await
8331 .unwrap()
8332 .with_branch("old", None)
8333 .load()
8334 .await
8335 .unwrap();
8336 assert_eq!(old_ds.manifest.branch.as_deref(), Some("old"));
8337 assert_eq!(
8338 scan_id_column(&old_ds).await,
8339 vec![1],
8340 "the fork must contain only the fork-point data"
8341 );
8342 }
8343
8344 #[test]
8347 fn test_manifest_version_from_filename() {
8348 assert_eq!(
8350 DirectoryNamespace::manifest_version_from_filename("5.manifest"),
8351 Some(5)
8352 );
8353 assert_eq!(
8354 DirectoryNamespace::manifest_version_from_filename("0.manifest"),
8355 Some(0)
8356 );
8357 let v2_five = format!("{:020}.manifest", u64::MAX - 5);
8359 assert_eq!(
8360 DirectoryNamespace::manifest_version_from_filename(&v2_five),
8361 Some(5)
8362 );
8363 let v2_zero = format!("{:020}.manifest", u64::MAX);
8364 assert_eq!(
8365 DirectoryNamespace::manifest_version_from_filename(&v2_zero),
8366 Some(0)
8367 );
8368 assert_eq!(
8370 DirectoryNamespace::manifest_version_from_filename("data.lance"),
8371 None
8372 );
8373 assert_eq!(
8374 DirectoryNamespace::manifest_version_from_filename("d5.manifest"),
8375 None
8376 );
8377 }
8378
8379 #[tokio::test]
8380 async fn test_create_table() {
8381 let (namespace, _temp_dir) = create_test_namespace().await;
8382
8383 let schema = create_test_schema();
8385 let ipc_data = create_test_ipc_data(&schema);
8386
8387 let mut request = CreateTableRequest::new();
8388 request.id = Some(vec!["test_table".to_string()]);
8389
8390 let response = namespace
8391 .create_table(request, bytes::Bytes::from(ipc_data))
8392 .await
8393 .unwrap();
8394
8395 assert!(response.location.is_some());
8396 assert!(response.location.unwrap().ends_with("test_table.lance"));
8397 assert_eq!(response.version, Some(1));
8398 }
8399
8400 #[tokio::test]
8401 async fn test_create_table_without_data() {
8402 let (namespace, _temp_dir) = create_test_namespace().await;
8403
8404 let mut request = CreateTableRequest::new();
8405 request.id = Some(vec!["test_table".to_string()]);
8406
8407 let result = namespace.create_table(request, bytes::Bytes::new()).await;
8408 assert!(result.is_err());
8409 assert!(
8410 result
8411 .unwrap_err()
8412 .to_string()
8413 .contains("Arrow IPC stream) is required")
8414 );
8415 }
8416
8417 #[tokio::test]
8418 async fn test_create_table_with_invalid_id() {
8419 let (namespace, _temp_dir) = create_test_namespace().await;
8420
8421 let schema = create_test_schema();
8423 let ipc_data = create_test_ipc_data(&schema);
8424
8425 let mut request = CreateTableRequest::new();
8427 request.id = Some(vec![]);
8428
8429 let result = namespace
8430 .create_table(request, bytes::Bytes::from(ipc_data.clone()))
8431 .await;
8432 assert!(result.is_err());
8433
8434 let mut create_ns_req = CreateNamespaceRequest::new();
8437 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
8438 namespace.create_namespace(create_ns_req).await.unwrap();
8439
8440 let mut request = CreateTableRequest::new();
8442 request.id = Some(vec!["test_namespace".to_string(), "table".to_string()]);
8443
8444 let result = namespace
8445 .create_table(request, bytes::Bytes::from(ipc_data))
8446 .await;
8447 assert!(
8449 result.is_ok(),
8450 "Multi-level table IDs should work with manifest enabled"
8451 );
8452 }
8453
8454 #[tokio::test]
8455 async fn test_list_tables() {
8456 let (namespace, _temp_dir) = create_test_namespace().await;
8457
8458 let mut request = ListTablesRequest::new();
8460 request.id = Some(vec![]);
8461 let response = namespace.list_tables(request).await.unwrap();
8462 assert_eq!(response.tables.len(), 0);
8463
8464 let schema = create_test_schema();
8466 let ipc_data = create_test_ipc_data(&schema);
8467
8468 let mut create_request = CreateTableRequest::new();
8470 create_request.id = Some(vec!["table1".to_string()]);
8471 namespace
8472 .create_table(create_request, bytes::Bytes::from(ipc_data.clone()))
8473 .await
8474 .unwrap();
8475
8476 let mut create_request = CreateTableRequest::new();
8478 create_request.id = Some(vec!["table2".to_string()]);
8479 namespace
8480 .create_table(create_request, bytes::Bytes::from(ipc_data))
8481 .await
8482 .unwrap();
8483
8484 let mut request = ListTablesRequest::new();
8486 request.id = Some(vec![]);
8487 let response = namespace.list_tables(request).await.unwrap();
8488 let tables = response.tables;
8489 assert_eq!(tables.len(), 2);
8490 assert!(tables.contains(&"table1".to_string()));
8491 assert!(tables.contains(&"table2".to_string()));
8492 }
8493
8494 #[tokio::test]
8495 async fn test_list_tables_pagination() {
8496 let (namespace, _temp_dir) = create_test_namespace().await;
8497
8498 let schema = create_test_schema();
8499 let ipc_data = create_test_ipc_data(&schema);
8500
8501 for name in ["alpha", "bravo", "charlie"] {
8502 let mut req = CreateTableRequest::new();
8503 req.id = Some(vec![name.to_string()]);
8504 namespace
8505 .create_table(req, bytes::Bytes::from(ipc_data.clone()))
8506 .await
8507 .unwrap();
8508 }
8509
8510 let first_page = namespace
8512 .list_tables(ListTablesRequest {
8513 id: Some(vec![]),
8514 limit: Some(2),
8515 ..Default::default()
8516 })
8517 .await
8518 .unwrap();
8519
8520 assert_eq!(first_page.tables, vec!["alpha", "bravo"]);
8521 assert_eq!(first_page.page_token.as_deref(), Some("bravo"));
8522
8523 let second_page = namespace
8525 .list_tables(ListTablesRequest {
8526 id: Some(vec![]),
8527 limit: Some(2),
8528 page_token: first_page.page_token.clone(),
8529 ..Default::default()
8530 })
8531 .await
8532 .unwrap();
8533
8534 assert_eq!(second_page.tables, vec!["charlie"]);
8535 assert!(second_page.page_token.is_none());
8536 }
8537
8538 #[tokio::test]
8539 async fn test_list_tables_pagination_limit_zero() {
8540 let (namespace, _temp_dir) = create_test_namespace().await;
8541
8542 let schema = create_test_schema();
8543 let ipc_data = create_test_ipc_data(&schema);
8544
8545 let mut req = CreateTableRequest::new();
8546 req.id = Some(vec!["alpha".to_string()]);
8547 namespace
8548 .create_table(req, bytes::Bytes::from(ipc_data))
8549 .await
8550 .unwrap();
8551
8552 let response = namespace
8553 .list_tables(ListTablesRequest {
8554 id: Some(vec![]),
8555 limit: Some(0),
8556 ..Default::default()
8557 })
8558 .await
8559 .unwrap();
8560
8561 assert!(response.tables.is_empty());
8562 assert!(response.page_token.is_none());
8563 }
8564
8565 #[tokio::test]
8566 async fn test_list_tables_with_namespace_id() {
8567 let (namespace, _temp_dir) = create_test_namespace().await;
8568
8569 let mut create_ns_req = CreateNamespaceRequest::new();
8571 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
8572 namespace.create_namespace(create_ns_req).await.unwrap();
8573
8574 let mut request = ListTablesRequest::new();
8576 request.id = Some(vec!["test_namespace".to_string()]);
8577
8578 let result = namespace.list_tables(request).await;
8579 assert!(
8581 result.is_ok(),
8582 "list_tables should work with child namespace when manifest is enabled"
8583 );
8584 let response = result.unwrap();
8585 assert_eq!(
8586 response.tables.len(),
8587 0,
8588 "Namespace should have no tables yet"
8589 );
8590 }
8591
8592 #[tokio::test]
8593 async fn test_create_scalar_index() {
8594 let (namespace, _temp_dir) = create_test_namespace().await;
8595 create_scalar_table(&namespace, "users").await;
8596
8597 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8598 let dataset = open_dataset(&namespace, "users").await;
8599 let expected_transaction_id = dataset
8600 .read_transaction()
8601 .await
8602 .unwrap()
8603 .map(|transaction| transaction.uuid);
8604 assert_eq!(transaction_id, expected_transaction_id);
8605 let indices = dataset.load_indices().await.unwrap();
8606 assert!(indices.iter().any(|index| index.name == "users_id_idx"));
8607 }
8608
8609 #[tokio::test]
8610 async fn test_create_vector_index() {
8611 use lance_namespace::models::CreateTableIndexRequest;
8612
8613 let (namespace, _temp_dir) = create_test_namespace().await;
8614 create_vector_table(&namespace, "vectors").await;
8615
8616 let mut create_index_request =
8617 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
8618 create_index_request.id = Some(vec!["vectors".to_string()]);
8619 create_index_request.name = Some("vector_idx".to_string());
8620 create_index_request.distance_type = Some("l2".to_string());
8621 let transaction_id = namespace
8622 .create_table_index(create_index_request)
8623 .await
8624 .unwrap()
8625 .transaction_id;
8626
8627 let dataset = open_dataset(&namespace, "vectors").await;
8628 let expected_transaction_id = dataset
8629 .read_transaction()
8630 .await
8631 .unwrap()
8632 .map(|transaction| transaction.uuid);
8633 assert_eq!(transaction_id, expected_transaction_id);
8634 let indices = dataset.load_indices().await.unwrap();
8635 assert!(indices.iter().any(|index| index.name == "vector_idx"));
8636 }
8637
8638 #[tokio::test]
8639 async fn test_list_table_indices() {
8640 use lance_namespace::models::{CreateTableIndexRequest, ListTableIndicesRequest};
8641
8642 let (namespace, _temp_dir) = create_test_namespace().await;
8643 create_scalar_table(&namespace, "users").await;
8644 create_scalar_index(&namespace, "users", "a_idx").await;
8645 create_scalar_index(&namespace, "users", "b_idx").await;
8646 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8647
8648 let response = namespace
8649 .list_table_indices(ListTableIndicesRequest {
8650 id: Some(vec!["users".to_string()]),
8651 ..Default::default()
8652 })
8653 .await
8654 .unwrap();
8655
8656 assert_eq!(response.indexes.len(), 3);
8657 assert_eq!(response.indexes[0].index_name, "a_idx");
8658 assert_eq!(response.indexes[1].index_name, "b_idx");
8659 assert_eq!(response.indexes[2].index_name, "users_id_idx");
8660 assert!(response.page_token.is_none());
8661 let users_id_idx = response
8662 .indexes
8663 .iter()
8664 .find(|index| index.index_name == "users_id_idx")
8665 .unwrap();
8666 assert_eq!(users_id_idx.columns, vec!["id"]);
8667 assert_eq!(users_id_idx.status, "SUCCEEDED");
8668
8669 assert_eq!(users_id_idx.index_type.as_deref(), Some("BTree"));
8671 assert!(
8672 users_id_idx
8673 .type_url
8674 .as_deref()
8675 .is_some_and(|s| !s.is_empty())
8676 );
8677 assert_eq!(users_id_idx.num_indexed_rows, Some(3));
8678 assert_eq!(users_id_idx.num_unindexed_rows, Some(0));
8679 assert_eq!(users_id_idx.num_segments, Some(1));
8680 assert!(users_id_idx.size_bytes.is_some_and(|size| size > 0));
8681 assert!(users_id_idx.created_at.is_some());
8682 assert!(users_id_idx.index_version.is_some());
8683 assert!(users_id_idx.index_details.is_some());
8684
8685 let dataset = open_dataset(&namespace, "users").await;
8686 let expected_transaction_id = dataset
8687 .read_transaction()
8688 .await
8689 .unwrap()
8690 .map(|transaction| transaction.uuid);
8691 assert_eq!(transaction_id, expected_transaction_id);
8692 let indices = dataset.load_indices().await.unwrap();
8693 assert_eq!(
8694 indices
8695 .iter()
8696 .filter(|index| index.name == "users_id_idx")
8697 .count(),
8698 1
8699 );
8700
8701 let first_page = namespace
8702 .list_table_indices(ListTableIndicesRequest {
8703 id: Some(vec!["users".to_string()]),
8704 limit: Some(2),
8705 ..Default::default()
8706 })
8707 .await
8708 .unwrap();
8709
8710 assert_eq!(first_page.indexes.len(), 2);
8711 assert_eq!(first_page.indexes[0].index_name, "a_idx");
8712 assert_eq!(first_page.indexes[1].index_name, "b_idx");
8713 assert_eq!(first_page.page_token.as_deref(), Some("b_idx"));
8714
8715 let second_page = namespace
8716 .list_table_indices(ListTableIndicesRequest {
8717 id: Some(vec!["users".to_string()]),
8718 page_token: first_page.page_token.clone(),
8719 limit: Some(2),
8720 ..Default::default()
8721 })
8722 .await
8723 .unwrap();
8724
8725 assert_eq!(second_page.indexes.len(), 1);
8726 assert_eq!(second_page.indexes[0].index_name, "users_id_idx");
8727 assert!(second_page.page_token.is_none());
8728
8729 create_vector_table(&namespace, "vectors").await;
8731 let mut create_index_request =
8732 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
8733 create_index_request.id = Some(vec!["vectors".to_string()]);
8734 create_index_request.name = Some("vector_idx".to_string());
8735 create_index_request.distance_type = Some("l2".to_string());
8736 namespace
8737 .create_table_index(create_index_request)
8738 .await
8739 .unwrap();
8740
8741 let vector_response = namespace
8742 .list_table_indices(ListTableIndicesRequest {
8743 id: Some(vec!["vectors".to_string()]),
8744 ..Default::default()
8745 })
8746 .await
8747 .unwrap();
8748
8749 assert_eq!(vector_response.indexes.len(), 1);
8750 let vector_idx = &vector_response.indexes[0];
8751 assert_eq!(vector_idx.index_name, "vector_idx");
8752 assert_eq!(vector_idx.columns, vec!["vector"]);
8753 assert_eq!(vector_idx.index_type.as_deref(), Some("IVF_FLAT"));
8754 assert!(
8755 vector_idx
8756 .type_url
8757 .as_deref()
8758 .is_some_and(|s| !s.is_empty())
8759 );
8760 assert!(vector_idx.num_indexed_rows.is_some());
8761 assert!(vector_idx.num_unindexed_rows.is_some());
8762 assert_eq!(vector_idx.num_segments, Some(1));
8763 assert!(vector_idx.created_at.is_some());
8764 assert!(vector_idx.index_version.is_some());
8765 assert!(vector_idx.index_details.is_some());
8766 }
8767
8768 #[tokio::test]
8769 async fn test_describe_table_index_stats() {
8770 use lance_namespace::models::DescribeTableIndexStatsRequest;
8771
8772 let (namespace, _temp_dir) = create_test_namespace().await;
8773 create_scalar_table(&namespace, "users").await;
8774 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8775
8776 let response = namespace
8777 .describe_table_index_stats(DescribeTableIndexStatsRequest {
8778 id: Some(vec!["users".to_string()]),
8779 index_name: Some("users_id_idx".to_string()),
8780 ..Default::default()
8781 })
8782 .await
8783 .unwrap();
8784 assert_eq!(response.index_type, Some("BTree".to_string()));
8785 assert_eq!(response.num_indices, Some(1));
8786 assert_eq!(response.num_indexed_rows, Some(3));
8787 assert_eq!(response.num_unindexed_rows, Some(0));
8788
8789 let dataset = open_dataset(&namespace, "users").await;
8790 let expected_transaction_id = dataset
8791 .read_transaction()
8792 .await
8793 .unwrap()
8794 .map(|transaction| transaction.uuid);
8795 assert_eq!(transaction_id, expected_transaction_id);
8796 let stats: serde_json::Value =
8797 serde_json::from_str(&dataset.index_statistics("users_id_idx").await.unwrap()).unwrap();
8798 assert_eq!(stats["index_type"], "BTree");
8799 assert_eq!(stats["num_indices"], 1);
8800 assert_eq!(stats["num_indexed_rows"], 3);
8801 assert_eq!(stats["num_unindexed_rows"], 0);
8802 }
8803
8804 #[tokio::test]
8805 async fn test_describe_transaction() {
8806 use lance_namespace::models::DescribeTransactionRequest;
8807
8808 let (namespace, _temp_dir) = create_test_namespace().await;
8809 create_scalar_table(&namespace, "users").await;
8810 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8811 let dataset = open_dataset(&namespace, "users").await;
8812 let latest_transaction = dataset.read_transaction().await.unwrap();
8813 assert_eq!(
8814 transaction_id,
8815 latest_transaction
8816 .as_ref()
8817 .map(|transaction| transaction.uuid.clone())
8818 );
8819
8820 if let Some(transaction_id) = transaction_id {
8821 let response = namespace
8822 .describe_transaction(DescribeTransactionRequest {
8823 id: Some(vec!["users".to_string(), transaction_id.clone()]),
8824 ..Default::default()
8825 })
8826 .await
8827 .unwrap();
8828 assert_eq!(response.status, "SUCCEEDED");
8829 assert_eq!(
8830 response
8831 .properties
8832 .as_ref()
8833 .and_then(|props| props.get("operation")),
8834 Some(&"CreateIndex".to_string())
8835 );
8836 assert_eq!(
8837 response
8838 .properties
8839 .as_ref()
8840 .and_then(|props| props.get("uuid")),
8841 Some(&transaction_id)
8842 );
8843 } else {
8844 assert!(latest_transaction.is_none());
8845 }
8846 }
8847
8848 #[tokio::test]
8849 async fn test_drop_table_index() {
8850 use lance_namespace::models::{DropTableIndexRequest, ListTableIndicesRequest};
8851
8852 let (namespace, _temp_dir) = create_test_namespace().await;
8853 create_scalar_table(&namespace, "users").await;
8854 let create_transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
8855
8856 let drop_transaction_id = namespace
8857 .drop_table_index(DropTableIndexRequest {
8858 id: Some(vec!["users".to_string()]),
8859 index_name: Some("users_id_idx".to_string()),
8860 ..Default::default()
8861 })
8862 .await
8863 .unwrap()
8864 .transaction_id;
8865
8866 let dataset = open_dataset(&namespace, "users").await;
8867 let previous_dataset = dataset
8868 .checkout_version(dataset.version().version - 1)
8869 .await
8870 .unwrap();
8871 let previous_transaction_id = previous_dataset
8872 .read_transaction()
8873 .await
8874 .unwrap()
8875 .map(|transaction| transaction.uuid);
8876 assert_eq!(create_transaction_id, previous_transaction_id);
8877 let expected_drop_transaction_id = dataset
8878 .read_transaction()
8879 .await
8880 .unwrap()
8881 .map(|transaction| transaction.uuid);
8882 assert_eq!(drop_transaction_id, expected_drop_transaction_id);
8883 let indices = dataset.load_indices().await.unwrap();
8884 assert!(!indices.iter().any(|index| index.name == "users_id_idx"));
8885
8886 let list_response = namespace
8887 .list_table_indices(ListTableIndicesRequest {
8888 id: Some(vec!["users".to_string()]),
8889 ..Default::default()
8890 })
8891 .await
8892 .unwrap();
8893 assert!(list_response.indexes.is_empty());
8894 }
8895
8896 #[tokio::test]
8897 async fn test_describe_table() {
8898 let (namespace, _temp_dir) = create_test_namespace().await;
8899
8900 let schema = create_test_schema();
8902 let ipc_data = create_test_ipc_data(&schema);
8903
8904 let mut create_request = CreateTableRequest::new();
8905 create_request.id = Some(vec!["test_table".to_string()]);
8906 namespace
8907 .create_table(create_request, bytes::Bytes::from(ipc_data))
8908 .await
8909 .unwrap();
8910
8911 let mut request = DescribeTableRequest::new();
8913 request.id = Some(vec!["test_table".to_string()]);
8914 let response = namespace.describe_table(request).await.unwrap();
8915
8916 assert!(response.location.is_some());
8917 assert!(response.location.unwrap().ends_with("test_table.lance"));
8918 }
8919
8920 #[tokio::test]
8921 async fn test_describe_nonexistent_table() {
8922 let (namespace, _temp_dir) = create_test_namespace().await;
8923
8924 let mut request = DescribeTableRequest::new();
8925 request.id = Some(vec!["nonexistent".to_string()]);
8926
8927 let result = namespace.describe_table(request).await;
8928 assert!(result.is_err());
8929 assert!(result.unwrap_err().to_string().contains("Table not found"));
8930 }
8931
8932 #[tokio::test]
8933 async fn test_table_exists() {
8934 let (namespace, _temp_dir) = create_test_namespace().await;
8935
8936 let schema = create_test_schema();
8938 let ipc_data = create_test_ipc_data(&schema);
8939
8940 let mut create_request = CreateTableRequest::new();
8941 create_request.id = Some(vec!["existing_table".to_string()]);
8942 namespace
8943 .create_table(create_request, bytes::Bytes::from(ipc_data))
8944 .await
8945 .unwrap();
8946
8947 let mut request = TableExistsRequest::new();
8949 request.id = Some(vec!["existing_table".to_string()]);
8950 let result = namespace.table_exists(request).await;
8951 assert!(result.is_ok());
8952
8953 let mut request = TableExistsRequest::new();
8955 request.id = Some(vec!["nonexistent".to_string()]);
8956 let result = namespace.table_exists(request).await;
8957 assert!(result.is_err());
8958 assert!(result.unwrap_err().to_string().contains("Table not found"));
8959 }
8960
8961 #[tokio::test]
8962 async fn test_drop_table() {
8963 let (namespace, _temp_dir) = create_test_namespace().await;
8964
8965 let schema = create_test_schema();
8967 let ipc_data = create_test_ipc_data(&schema);
8968
8969 let mut create_request = CreateTableRequest::new();
8970 create_request.id = Some(vec!["table_to_drop".to_string()]);
8971 namespace
8972 .create_table(create_request, bytes::Bytes::from(ipc_data))
8973 .await
8974 .unwrap();
8975
8976 let mut exists_request = TableExistsRequest::new();
8978 exists_request.id = Some(vec!["table_to_drop".to_string()]);
8979 assert!(namespace.table_exists(exists_request.clone()).await.is_ok());
8980
8981 let mut drop_request = DropTableRequest::new();
8983 drop_request.id = Some(vec!["table_to_drop".to_string()]);
8984 let response = namespace.drop_table(drop_request).await.unwrap();
8985 assert!(response.location.is_some());
8986
8987 assert!(namespace.table_exists(exists_request).await.is_err());
8989 }
8990
8991 #[tokio::test]
8992 async fn test_drop_nonexistent_table() {
8993 let (namespace, _temp_dir) = create_test_namespace().await;
8994
8995 let mut request = DropTableRequest::new();
8996 request.id = Some(vec!["nonexistent".to_string()]);
8997
8998 let result = namespace.drop_table(request).await;
9000 let _ = result;
9003 }
9004
9005 #[tokio::test]
9006 async fn test_root_namespace_operations() {
9007 let (namespace, _temp_dir) = create_test_namespace().await;
9008
9009 let mut request = ListNamespacesRequest::new();
9011 request.id = Some(vec![]);
9012 let result = namespace.list_namespaces(request).await;
9013 assert!(result.is_ok());
9014 assert_eq!(result.unwrap().namespaces.len(), 0);
9015
9016 let mut request = DescribeNamespaceRequest::new();
9018 request.id = Some(vec![]);
9019 let result = namespace.describe_namespace(request).await;
9020 assert!(result.is_ok());
9021
9022 let mut request = NamespaceExistsRequest::new();
9024 request.id = Some(vec![]);
9025 let result = namespace.namespace_exists(request).await;
9026 assert!(result.is_ok());
9027
9028 let mut request = CreateNamespaceRequest::new();
9030 request.id = Some(vec![]);
9031 let result = namespace.create_namespace(request).await;
9032 assert!(result.is_err());
9033 assert!(result.unwrap_err().to_string().contains("already exists"));
9034
9035 let mut request = DropNamespaceRequest::new();
9037 request.id = Some(vec![]);
9038 let result = namespace.drop_namespace(request).await;
9039 assert!(result.is_err());
9040 assert!(
9041 result
9042 .unwrap_err()
9043 .to_string()
9044 .contains("cannot be dropped")
9045 );
9046 }
9047
9048 #[tokio::test]
9049 async fn test_non_root_namespace_operations() {
9050 let (namespace, _temp_dir) = create_test_namespace().await;
9051
9052 let mut request = CreateNamespaceRequest::new();
9055 request.id = Some(vec!["child".to_string()]);
9056 let result = namespace.create_namespace(request).await;
9057 assert!(
9058 result.is_ok(),
9059 "Child namespace creation should succeed with manifest enabled"
9060 );
9061
9062 let mut request = NamespaceExistsRequest::new();
9064 request.id = Some(vec!["child".to_string()]);
9065 let result = namespace.namespace_exists(request).await;
9066 assert!(
9067 result.is_ok(),
9068 "Child namespace should exist after creation"
9069 );
9070
9071 let mut request = DropNamespaceRequest::new();
9073 request.id = Some(vec!["child".to_string()]);
9074 let result = namespace.drop_namespace(request).await;
9075 assert!(
9076 result.is_ok(),
9077 "Child namespace drop should succeed with manifest enabled"
9078 );
9079
9080 let mut request = NamespaceExistsRequest::new();
9082 request.id = Some(vec!["child".to_string()]);
9083 let result = namespace.namespace_exists(request).await;
9084 assert!(
9085 result.is_err(),
9086 "Child namespace should not exist after drop"
9087 );
9088 }
9089
9090 #[tokio::test]
9091 async fn test_config_custom_root() {
9092 let temp_dir = TempStdDir::default();
9093 let custom_path = temp_dir.join("custom");
9094 std::fs::create_dir(&custom_path).unwrap();
9095
9096 let namespace = DirectoryNamespaceBuilder::new(custom_path.to_string_lossy().to_string())
9097 .build()
9098 .await
9099 .unwrap();
9100
9101 let schema = create_test_schema();
9103 let ipc_data = create_test_ipc_data(&schema);
9104
9105 let mut request = CreateTableRequest::new();
9107 request.id = Some(vec!["test_table".to_string()]);
9108
9109 let response = namespace
9110 .create_table(request, bytes::Bytes::from(ipc_data))
9111 .await
9112 .unwrap();
9113
9114 assert!(response.location.unwrap().contains("custom"));
9115 }
9116
9117 #[tokio::test]
9118 async fn test_config_storage_options() {
9119 let temp_dir = TempStdDir::default();
9120
9121 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9122 .storage_option("option1", "value1")
9123 .storage_option("option2", "value2")
9124 .build()
9125 .await
9126 .unwrap();
9127
9128 let schema = create_test_schema();
9130 let ipc_data = create_test_ipc_data(&schema);
9131
9132 let mut request = CreateTableRequest::new();
9134 request.id = Some(vec!["test_table".to_string()]);
9135
9136 let response = namespace
9137 .create_table(request, bytes::Bytes::from(ipc_data))
9138 .await
9139 .unwrap();
9140
9141 let storage_options = response.storage_options.unwrap();
9142 assert_eq!(storage_options.get("option1"), Some(&"value1".to_string()));
9143 assert_eq!(storage_options.get("option2"), Some(&"value2".to_string()));
9144 }
9145
9146 #[tokio::test]
9150 async fn test_no_storage_options_without_vendor() {
9151 use lance_namespace::models::DeclareTableRequest;
9152
9153 let temp_dir = TempStdDir::default();
9154
9155 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9157 .manifest_enabled(false)
9158 .storage_option("aws_access_key_id", "AKID")
9159 .storage_option("aws_secret_access_key", "SECRET")
9160 .storage_option("region", "us-east-1")
9161 .build()
9162 .await
9163 .unwrap();
9164
9165 let schema = create_test_schema();
9166 let ipc_data = create_test_ipc_data(&schema);
9167
9168 let mut create_req = CreateTableRequest::new();
9170 create_req.id = Some(vec!["t1".to_string()]);
9171 namespace
9172 .create_table(create_req, bytes::Bytes::from(ipc_data))
9173 .await
9174 .unwrap();
9175
9176 let mut desc_req = DescribeTableRequest::new();
9178 desc_req.id = Some(vec!["t1".to_string()]);
9179 let resp = namespace.describe_table(desc_req).await.unwrap();
9180 assert!(resp.storage_options.is_none());
9181
9182 let mut decl_req = DeclareTableRequest::new();
9184 decl_req.id = Some(vec!["t2".to_string()]);
9185 let resp = namespace.declare_table(decl_req).await.unwrap();
9186 assert!(resp.storage_options.is_none());
9187 }
9188
9189 #[tokio::test]
9191 async fn test_no_storage_options_without_vendor_manifest() {
9192 let temp_dir = TempStdDir::default();
9193
9194 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9195 .storage_option("aws_access_key_id", "AKID")
9196 .storage_option("aws_secret_access_key", "SECRET")
9197 .storage_option("region", "us-east-1")
9198 .build()
9199 .await
9200 .unwrap();
9201
9202 let schema = create_test_schema();
9203 let ipc_data = create_test_ipc_data(&schema);
9204
9205 let mut create_req = CreateTableRequest::new();
9206 create_req.id = Some(vec!["t1".to_string()]);
9207 namespace
9208 .create_table(create_req, bytes::Bytes::from(ipc_data))
9209 .await
9210 .unwrap();
9211
9212 let mut desc_req = DescribeTableRequest::new();
9214 desc_req.id = Some(vec!["t1".to_string()]);
9215 let resp = namespace.describe_table(desc_req).await.unwrap();
9216 assert!(resp.storage_options.is_none());
9217 }
9218
9219 #[tokio::test]
9220 async fn test_from_properties_manifest_enabled() {
9221 let temp_dir = TempStdDir::default();
9222
9223 let mut properties = HashMap::new();
9224 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9225 properties.insert("manifest_enabled".to_string(), "true".to_string());
9226 properties.insert("dir_listing_enabled".to_string(), "false".to_string());
9227
9228 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9229 assert!(builder.manifest_enabled);
9230 assert!(!builder.dir_listing_enabled);
9231
9232 let namespace = builder.build().await.unwrap();
9233
9234 let schema = create_test_schema();
9236 let ipc_data = create_test_ipc_data(&schema);
9237
9238 let mut request = CreateTableRequest::new();
9240 request.id = Some(vec!["test_table".to_string()]);
9241
9242 let response = namespace
9243 .create_table(request, bytes::Bytes::from(ipc_data))
9244 .await
9245 .unwrap();
9246
9247 assert!(response.location.is_some());
9248 }
9249
9250 #[tokio::test]
9251 async fn test_from_properties_dir_listing_enabled() {
9252 let temp_dir = TempStdDir::default();
9253
9254 let mut properties = HashMap::new();
9255 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9256 properties.insert("manifest_enabled".to_string(), "false".to_string());
9257 properties.insert("dir_listing_enabled".to_string(), "true".to_string());
9258
9259 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9260 assert!(!builder.manifest_enabled);
9261 assert!(builder.dir_listing_enabled);
9262
9263 let namespace = builder.build().await.unwrap();
9264
9265 let schema = create_test_schema();
9267 let ipc_data = create_test_ipc_data(&schema);
9268
9269 let mut request = CreateTableRequest::new();
9271 request.id = Some(vec!["test_table".to_string()]);
9272
9273 let response = namespace
9274 .create_table(request, bytes::Bytes::from(ipc_data))
9275 .await
9276 .unwrap();
9277
9278 assert!(response.location.is_some());
9279 }
9280
9281 #[tokio::test]
9282 async fn test_from_properties_defaults() {
9283 let temp_dir = TempStdDir::default();
9284
9285 let mut properties = HashMap::new();
9286 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9287
9288 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9289 assert!(builder.manifest_enabled);
9291 assert!(builder.dir_listing_enabled);
9292 }
9293
9294 #[tokio::test]
9295 async fn test_from_properties_with_storage_options() {
9296 let temp_dir = TempStdDir::default();
9297
9298 let mut properties = HashMap::new();
9299 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9300 properties.insert("manifest_enabled".to_string(), "true".to_string());
9301 properties.insert("storage.region".to_string(), "us-west-2".to_string());
9302 properties.insert("storage.bucket".to_string(), "my-bucket".to_string());
9303
9304 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9305 assert!(builder.manifest_enabled);
9306 assert!(builder.storage_options.is_some());
9307
9308 let storage_options = builder.storage_options.unwrap();
9309 assert_eq!(
9310 storage_options.get("region"),
9311 Some(&"us-west-2".to_string())
9312 );
9313 assert_eq!(
9314 storage_options.get("bucket"),
9315 Some(&"my-bucket".to_string())
9316 );
9317 }
9318
9319 #[tokio::test]
9320 async fn test_various_arrow_types() {
9321 let (namespace, _temp_dir) = create_test_namespace().await;
9322
9323 let fields = vec![
9325 JsonArrowField {
9326 name: "bool_col".to_string(),
9327 r#type: Box::new(JsonArrowDataType::new("bool".to_string())),
9328 nullable: true,
9329 metadata: None,
9330 },
9331 JsonArrowField {
9332 name: "int8_col".to_string(),
9333 r#type: Box::new(JsonArrowDataType::new("int8".to_string())),
9334 nullable: true,
9335 metadata: None,
9336 },
9337 JsonArrowField {
9338 name: "float64_col".to_string(),
9339 r#type: Box::new(JsonArrowDataType::new("float64".to_string())),
9340 nullable: true,
9341 metadata: None,
9342 },
9343 JsonArrowField {
9344 name: "binary_col".to_string(),
9345 r#type: Box::new(JsonArrowDataType::new("binary".to_string())),
9346 nullable: true,
9347 metadata: None,
9348 },
9349 ];
9350
9351 let schema = JsonArrowSchema {
9352 fields,
9353 metadata: None,
9354 };
9355
9356 let ipc_data = create_test_ipc_data(&schema);
9358
9359 let mut request = CreateTableRequest::new();
9360 request.id = Some(vec!["complex_table".to_string()]);
9361
9362 let response = namespace
9363 .create_table(request, bytes::Bytes::from(ipc_data))
9364 .await
9365 .unwrap();
9366
9367 assert!(response.location.is_some());
9368 }
9369
9370 #[tokio::test]
9371 async fn test_connect_dir() {
9372 let temp_dir = TempStdDir::default();
9373
9374 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9375 .build()
9376 .await
9377 .unwrap();
9378
9379 let mut request = ListTablesRequest::new();
9381 request.id = Some(vec![]);
9382 let response = namespace.list_tables(request).await.unwrap();
9383 assert_eq!(response.tables.len(), 0);
9384 }
9385
9386 #[tokio::test]
9387 async fn test_create_table_with_ipc_data() {
9388 use arrow::array::{Int32Array, StringArray};
9389 use arrow::ipc::writer::StreamWriter;
9390
9391 let (namespace, _temp_dir) = create_test_namespace().await;
9392
9393 let schema = create_test_schema();
9395
9396 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
9398 let arrow_schema = Arc::new(arrow_schema);
9399
9400 let id_array = Int32Array::from(vec![1, 2, 3]);
9402 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
9403 let batch = arrow::record_batch::RecordBatch::try_new(
9404 arrow_schema.clone(),
9405 vec![Arc::new(id_array), Arc::new(name_array)],
9406 )
9407 .unwrap();
9408
9409 let mut buffer = Vec::new();
9411 {
9412 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
9413 writer.write(&batch).unwrap();
9414 writer.finish().unwrap();
9415 }
9416
9417 let mut request = CreateTableRequest::new();
9419 request.id = Some(vec!["test_table_with_data".to_string()]);
9420
9421 let response = namespace
9422 .create_table(request, Bytes::from(buffer))
9423 .await
9424 .unwrap();
9425
9426 assert_eq!(response.version, Some(1));
9427 assert!(
9428 response
9429 .location
9430 .unwrap()
9431 .contains("test_table_with_data.lance")
9432 );
9433
9434 let mut exists_request = TableExistsRequest::new();
9436 exists_request.id = Some(vec!["test_table_with_data".to_string()]);
9437 namespace.table_exists(exists_request).await.unwrap();
9438 }
9439
9440 #[tokio::test]
9441 async fn test_child_namespace_create_and_list() {
9442 let (namespace, _temp_dir) = create_test_namespace().await;
9443
9444 for i in 1..=3 {
9446 let mut create_req = CreateNamespaceRequest::new();
9447 create_req.id = Some(vec![format!("ns{}", i)]);
9448 let result = namespace.create_namespace(create_req).await;
9449 assert!(result.is_ok(), "Failed to create child namespace ns{}", i);
9450 }
9451
9452 let list_req = ListNamespacesRequest {
9454 id: Some(vec![]),
9455 ..Default::default()
9456 };
9457 let result = namespace.list_namespaces(list_req).await;
9458 assert!(result.is_ok());
9459 let namespaces = result.unwrap().namespaces;
9460 assert_eq!(namespaces.len(), 3);
9461 assert!(namespaces.contains(&"ns1".to_string()));
9462 assert!(namespaces.contains(&"ns2".to_string()));
9463 assert!(namespaces.contains(&"ns3".to_string()));
9464 }
9465
9466 #[tokio::test]
9467 async fn test_nested_namespace_hierarchy() {
9468 let (namespace, _temp_dir) = create_test_namespace().await;
9469
9470 let mut create_req = CreateNamespaceRequest::new();
9472 create_req.id = Some(vec!["parent".to_string()]);
9473 namespace.create_namespace(create_req).await.unwrap();
9474
9475 let mut create_req = CreateNamespaceRequest::new();
9477 create_req.id = Some(vec!["parent".to_string(), "child1".to_string()]);
9478 namespace.create_namespace(create_req).await.unwrap();
9479
9480 let mut create_req = CreateNamespaceRequest::new();
9481 create_req.id = Some(vec!["parent".to_string(), "child2".to_string()]);
9482 namespace.create_namespace(create_req).await.unwrap();
9483
9484 let list_req = ListNamespacesRequest {
9486 id: Some(vec!["parent".to_string()]),
9487 ..Default::default()
9488 };
9489 let result = namespace.list_namespaces(list_req).await;
9490 assert!(result.is_ok());
9491 let children = result.unwrap().namespaces;
9492 assert_eq!(children.len(), 2);
9493 assert!(children.contains(&"child1".to_string()));
9494 assert!(children.contains(&"child2".to_string()));
9495
9496 let list_req = ListNamespacesRequest {
9498 id: Some(vec![]),
9499 ..Default::default()
9500 };
9501 let result = namespace.list_namespaces(list_req).await;
9502 assert!(result.is_ok());
9503 let root_namespaces = result.unwrap().namespaces;
9504 assert_eq!(root_namespaces.len(), 1);
9505 assert_eq!(root_namespaces[0], "parent");
9506 }
9507
9508 #[tokio::test]
9509 async fn test_table_in_child_namespace() {
9510 let (namespace, _temp_dir) = create_test_namespace().await;
9511
9512 let mut create_ns_req = CreateNamespaceRequest::new();
9514 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9515 namespace.create_namespace(create_ns_req).await.unwrap();
9516
9517 let schema = create_test_schema();
9519 let ipc_data = create_test_ipc_data(&schema);
9520 let mut create_table_req = CreateTableRequest::new();
9521 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9522 let result = namespace
9523 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9524 .await;
9525 assert!(result.is_ok(), "Failed to create table in child namespace");
9526
9527 let list_req = ListTablesRequest {
9529 id: Some(vec!["test_ns".to_string()]),
9530 ..Default::default()
9531 };
9532 let result = namespace.list_tables(list_req).await;
9533 assert!(result.is_ok());
9534 let tables = result.unwrap().tables;
9535 assert_eq!(tables.len(), 1);
9536 assert_eq!(tables[0], "table1");
9537
9538 let mut exists_req = TableExistsRequest::new();
9540 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9541 let result = namespace.table_exists(exists_req).await;
9542 assert!(result.is_ok());
9543
9544 let mut describe_req = DescribeTableRequest::new();
9546 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9547 let result = namespace.describe_table(describe_req).await;
9548 assert!(result.is_ok());
9549 let response = result.unwrap();
9550 assert!(response.location.is_some());
9551 }
9552
9553 #[tokio::test]
9554 async fn test_multiple_tables_in_child_namespace() {
9555 let (namespace, _temp_dir) = create_test_namespace().await;
9556
9557 let mut create_ns_req = CreateNamespaceRequest::new();
9559 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9560 namespace.create_namespace(create_ns_req).await.unwrap();
9561
9562 let schema = create_test_schema();
9564 let ipc_data = create_test_ipc_data(&schema);
9565 for i in 1..=3 {
9566 let mut create_table_req = CreateTableRequest::new();
9567 create_table_req.id = Some(vec!["test_ns".to_string(), format!("table{}", i)]);
9568 namespace
9569 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
9570 .await
9571 .unwrap();
9572 }
9573
9574 let list_req = ListTablesRequest {
9576 id: Some(vec!["test_ns".to_string()]),
9577 ..Default::default()
9578 };
9579 let result = namespace.list_tables(list_req).await;
9580 assert!(result.is_ok());
9581 let tables = result.unwrap().tables;
9582 assert_eq!(tables.len(), 3);
9583 assert!(tables.contains(&"table1".to_string()));
9584 assert!(tables.contains(&"table2".to_string()));
9585 assert!(tables.contains(&"table3".to_string()));
9586 }
9587
9588 #[tokio::test]
9589 async fn test_drop_table_in_child_namespace() {
9590 let (namespace, _temp_dir) = create_test_namespace().await;
9591
9592 let mut create_ns_req = CreateNamespaceRequest::new();
9594 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9595 namespace.create_namespace(create_ns_req).await.unwrap();
9596
9597 let schema = create_test_schema();
9599 let ipc_data = create_test_ipc_data(&schema);
9600 let mut create_table_req = CreateTableRequest::new();
9601 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9602 namespace
9603 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9604 .await
9605 .unwrap();
9606
9607 let mut drop_req = DropTableRequest::new();
9609 drop_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9610 let result = namespace.drop_table(drop_req).await;
9611 assert!(result.is_ok(), "Failed to drop table in child namespace");
9612
9613 let mut exists_req = TableExistsRequest::new();
9615 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9616 let result = namespace.table_exists(exists_req).await;
9617 assert!(result.is_err());
9618 }
9619
9620 #[tokio::test]
9621 async fn test_deeply_nested_namespace() {
9622 let (namespace, _temp_dir) = create_test_namespace().await;
9623
9624 let mut create_req = CreateNamespaceRequest::new();
9626 create_req.id = Some(vec!["level1".to_string()]);
9627 namespace.create_namespace(create_req).await.unwrap();
9628
9629 let mut create_req = CreateNamespaceRequest::new();
9630 create_req.id = Some(vec!["level1".to_string(), "level2".to_string()]);
9631 namespace.create_namespace(create_req).await.unwrap();
9632
9633 let mut create_req = CreateNamespaceRequest::new();
9634 create_req.id = Some(vec![
9635 "level1".to_string(),
9636 "level2".to_string(),
9637 "level3".to_string(),
9638 ]);
9639 namespace.create_namespace(create_req).await.unwrap();
9640
9641 let schema = create_test_schema();
9643 let ipc_data = create_test_ipc_data(&schema);
9644 let mut create_table_req = CreateTableRequest::new();
9645 create_table_req.id = Some(vec![
9646 "level1".to_string(),
9647 "level2".to_string(),
9648 "level3".to_string(),
9649 "table1".to_string(),
9650 ]);
9651 let result = namespace
9652 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9653 .await;
9654 assert!(
9655 result.is_ok(),
9656 "Failed to create table in deeply nested namespace"
9657 );
9658
9659 let mut exists_req = TableExistsRequest::new();
9661 exists_req.id = Some(vec![
9662 "level1".to_string(),
9663 "level2".to_string(),
9664 "level3".to_string(),
9665 "table1".to_string(),
9666 ]);
9667 let result = namespace.table_exists(exists_req).await;
9668 assert!(result.is_ok());
9669 }
9670
9671 #[tokio::test]
9672 async fn test_namespace_with_properties() {
9673 let (namespace, _temp_dir) = create_test_namespace().await;
9674
9675 let mut properties = HashMap::new();
9677 properties.insert("owner".to_string(), "test_user".to_string());
9678 properties.insert("description".to_string(), "Test namespace".to_string());
9679
9680 let mut create_req = CreateNamespaceRequest::new();
9681 create_req.id = Some(vec!["test_ns".to_string()]);
9682 create_req.properties = Some(properties.clone());
9683 namespace.create_namespace(create_req).await.unwrap();
9684
9685 let describe_req = DescribeNamespaceRequest {
9687 id: Some(vec!["test_ns".to_string()]),
9688 ..Default::default()
9689 };
9690 let result = namespace.describe_namespace(describe_req).await;
9691 assert!(result.is_ok());
9692 let response = result.unwrap();
9693 assert!(response.properties.is_some());
9694 let props = response.properties.unwrap();
9695 assert_eq!(props.get("owner"), Some(&"test_user".to_string()));
9696 assert_eq!(
9697 props.get("description"),
9698 Some(&"Test namespace".to_string())
9699 );
9700 }
9701
9702 #[tokio::test]
9703 async fn test_cannot_drop_namespace_with_tables() {
9704 let (namespace, _temp_dir) = create_test_namespace().await;
9705
9706 let mut create_ns_req = CreateNamespaceRequest::new();
9708 create_ns_req.id = Some(vec!["test_ns".to_string()]);
9709 namespace.create_namespace(create_ns_req).await.unwrap();
9710
9711 let schema = create_test_schema();
9713 let ipc_data = create_test_ipc_data(&schema);
9714 let mut create_table_req = CreateTableRequest::new();
9715 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
9716 namespace
9717 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9718 .await
9719 .unwrap();
9720
9721 let mut drop_req = DropNamespaceRequest::new();
9723 drop_req.id = Some(vec!["test_ns".to_string()]);
9724 let result = namespace.drop_namespace(drop_req).await;
9725 assert!(
9726 result.is_err(),
9727 "Should not be able to drop namespace with tables"
9728 );
9729 }
9730
9731 #[tokio::test]
9732 async fn test_isolation_between_namespaces() {
9733 let (namespace, _temp_dir) = create_test_namespace().await;
9734
9735 let mut create_req = CreateNamespaceRequest::new();
9737 create_req.id = Some(vec!["ns1".to_string()]);
9738 namespace.create_namespace(create_req).await.unwrap();
9739
9740 let mut create_req = CreateNamespaceRequest::new();
9741 create_req.id = Some(vec!["ns2".to_string()]);
9742 namespace.create_namespace(create_req).await.unwrap();
9743
9744 let schema = create_test_schema();
9746 let ipc_data = create_test_ipc_data(&schema);
9747
9748 let mut create_table_req = CreateTableRequest::new();
9749 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
9750 namespace
9751 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
9752 .await
9753 .unwrap();
9754
9755 let mut create_table_req = CreateTableRequest::new();
9756 create_table_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
9757 namespace
9758 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
9759 .await
9760 .unwrap();
9761
9762 let list_req = ListTablesRequest {
9764 id: Some(vec!["ns1".to_string()]),
9765 page_token: None,
9766 limit: None,
9767 ..Default::default()
9768 };
9769 let result = namespace.list_tables(list_req).await.unwrap();
9770 assert_eq!(result.tables.len(), 1);
9771 assert_eq!(result.tables[0], "table1");
9772
9773 let list_req = ListTablesRequest {
9774 id: Some(vec!["ns2".to_string()]),
9775 page_token: None,
9776 limit: None,
9777 ..Default::default()
9778 };
9779 let result = namespace.list_tables(list_req).await.unwrap();
9780 assert_eq!(result.tables.len(), 1);
9781 assert_eq!(result.tables[0], "table1");
9782
9783 let mut drop_req = DropTableRequest::new();
9785 drop_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
9786 namespace.drop_table(drop_req).await.unwrap();
9787
9788 let mut exists_req = TableExistsRequest::new();
9790 exists_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
9791 assert!(namespace.table_exists(exists_req).await.is_err());
9792
9793 let mut exists_req = TableExistsRequest::new();
9794 exists_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
9795 assert!(namespace.table_exists(exists_req).await.is_ok());
9796 }
9797
9798 #[tokio::test]
9799 async fn test_migrate_directory_tables() {
9800 let temp_dir = TempStdDir::default();
9801 let temp_path = temp_dir.to_str().unwrap();
9802
9803 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
9805 .manifest_enabled(false)
9806 .dir_listing_enabled(true)
9807 .build()
9808 .await
9809 .unwrap();
9810
9811 let schema = create_test_schema();
9813 let ipc_data = create_test_ipc_data(&schema);
9814
9815 for i in 1..=3 {
9816 let mut create_req = CreateTableRequest::new();
9817 create_req.id = Some(vec![format!("table{}", i)]);
9818 dir_only_ns
9819 .create_table(create_req, bytes::Bytes::from(ipc_data.clone()))
9820 .await
9821 .unwrap();
9822 }
9823
9824 drop(dir_only_ns);
9825
9826 let dual_mode_ns = DirectoryNamespaceBuilder::new(temp_path)
9828 .manifest_enabled(true)
9829 .dir_listing_enabled(true)
9830 .build()
9831 .await
9832 .unwrap();
9833
9834 let mut list_req = ListTablesRequest::new();
9836 list_req.id = Some(vec![]);
9837 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
9838 assert_eq!(tables.len(), 3);
9839
9840 let migrated_count = dual_mode_ns.migrate().await.unwrap();
9842 assert_eq!(migrated_count, 3, "Should migrate all 3 tables");
9843
9844 let mut list_req = ListTablesRequest::new();
9846 list_req.id = Some(vec![]);
9847 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
9848 assert_eq!(tables.len(), 3);
9849
9850 let migrated_count = dual_mode_ns.migrate().await.unwrap();
9852 assert_eq!(
9853 migrated_count, 0,
9854 "Should not migrate already-migrated tables"
9855 );
9856
9857 drop(dual_mode_ns);
9858
9859 let manifest_only_ns = DirectoryNamespaceBuilder::new(temp_path)
9861 .manifest_enabled(true)
9862 .dir_listing_enabled(false)
9863 .build()
9864 .await
9865 .unwrap();
9866
9867 let mut list_req = ListTablesRequest::new();
9869 list_req.id = Some(vec![]);
9870 let tables = manifest_only_ns.list_tables(list_req).await.unwrap().tables;
9871 assert_eq!(tables.len(), 3);
9872 assert!(tables.contains(&"table1".to_string()));
9873 assert!(tables.contains(&"table2".to_string()));
9874 assert!(tables.contains(&"table3".to_string()));
9875 }
9876
9877 #[tokio::test]
9878 async fn test_migrate_without_manifest() {
9879 let temp_dir = TempStdDir::default();
9880 let temp_path = temp_dir.to_str().unwrap();
9881
9882 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9884 .manifest_enabled(false)
9885 .dir_listing_enabled(true)
9886 .build()
9887 .await
9888 .unwrap();
9889
9890 let migrated_count = namespace.migrate().await.unwrap();
9892 assert_eq!(migrated_count, 0);
9893 }
9894
9895 #[tokio::test]
9896 async fn test_register_table() {
9897 use lance_namespace::models::{RegisterTableRequest, TableExistsRequest};
9898
9899 let temp_dir = TempStdDir::default();
9900 let temp_path = temp_dir.to_str().unwrap();
9901
9902 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9903 .dir_listing_to_manifest_migration_enabled(true)
9904 .build()
9905 .await
9906 .unwrap();
9907
9908 let schema = create_test_schema();
9910 let ipc_data = create_test_ipc_data(&schema);
9911
9912 let table_uri = format!("{}/external_table.lance", temp_path);
9913 let cursor = Cursor::new(ipc_data);
9914 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
9915 let batches: Vec<_> = stream_reader
9916 .collect::<std::result::Result<Vec<_>, _>>()
9917 .unwrap();
9918 let schema = batches[0].schema();
9919 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
9920 let reader = RecordBatchIterator::new(batch_results, schema);
9921 Dataset::write(Box::new(reader), &table_uri, None)
9922 .await
9923 .unwrap();
9924
9925 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
9927 register_req.id = Some(vec!["registered_table".to_string()]);
9928
9929 let response = namespace.register_table(register_req).await.unwrap();
9930 assert_eq!(response.location, Some("external_table.lance".to_string()));
9931
9932 let mut exists_req = TableExistsRequest::new();
9934 exists_req.id = Some(vec!["registered_table".to_string()]);
9935 assert!(namespace.table_exists(exists_req).await.is_ok());
9936
9937 let mut list_req = ListTablesRequest::new();
9939 list_req.id = Some(vec![]);
9940 let tables = namespace.list_tables(list_req).await.unwrap();
9941 assert!(tables.tables.contains(&"registered_table".to_string()));
9942 }
9943
9944 #[tokio::test]
9945 async fn test_register_table_duplicate_fails() {
9946 use lance_namespace::models::RegisterTableRequest;
9947
9948 let temp_dir = TempStdDir::default();
9949 let temp_path = temp_dir.to_str().unwrap();
9950
9951 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9952 .build()
9953 .await
9954 .unwrap();
9955
9956 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
9958 register_req.id = Some(vec!["test_table".to_string()]);
9959
9960 namespace
9961 .register_table(register_req.clone())
9962 .await
9963 .unwrap();
9964
9965 let result = namespace.register_table(register_req).await;
9967 assert!(result.is_err());
9968 assert!(result.unwrap_err().to_string().contains("already exists"));
9969 }
9970
9971 #[tokio::test]
9972 async fn test_deregister_table() {
9973 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
9974
9975 let temp_dir = TempStdDir::default();
9976 let temp_path = temp_dir.to_str().unwrap();
9977
9978 let namespace = DirectoryNamespaceBuilder::new(temp_path)
9981 .manifest_enabled(true)
9982 .dir_listing_enabled(false)
9983 .build()
9984 .await
9985 .unwrap();
9986
9987 let schema = create_test_schema();
9989 let ipc_data = create_test_ipc_data(&schema);
9990
9991 let mut create_req = CreateTableRequest::new();
9992 create_req.id = Some(vec!["test_table".to_string()]);
9993 namespace
9994 .create_table(create_req, bytes::Bytes::from(ipc_data))
9995 .await
9996 .unwrap();
9997
9998 let mut exists_req = TableExistsRequest::new();
10000 exists_req.id = Some(vec!["test_table".to_string()]);
10001 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
10002
10003 let mut deregister_req = DeregisterTableRequest::new();
10005 deregister_req.id = Some(vec!["test_table".to_string()]);
10006 let response = namespace.deregister_table(deregister_req).await.unwrap();
10007
10008 assert!(
10010 response.location.is_some(),
10011 "Deregister should return location"
10012 );
10013 let location = response.location.as_ref().unwrap();
10014 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10017 .expect("Failed to convert temp path to URL");
10018 let expected_prefix = expected_url.to_string();
10019 assert!(
10020 location.starts_with(&expected_prefix),
10021 "Location should start with '{}', got: {}",
10022 expected_prefix,
10023 location
10024 );
10025 assert!(
10026 location.contains("test_table"),
10027 "Location should contain table name: {}",
10028 location
10029 );
10030 assert_eq!(response.id, Some(vec!["test_table".to_string()]));
10031
10032 assert!(namespace.table_exists(exists_req).await.is_err());
10034
10035 let dataset = Dataset::open(location).await;
10037 assert!(
10038 dataset.is_ok(),
10039 "Physical table data should still exist at {}",
10040 location
10041 );
10042 }
10043
10044 #[tokio::test]
10045 async fn test_deregister_table_in_child_namespace() {
10046 use lance_namespace::models::{
10047 CreateNamespaceRequest, DeregisterTableRequest, TableExistsRequest,
10048 };
10049
10050 let temp_dir = TempStdDir::default();
10051 let temp_path = temp_dir.to_str().unwrap();
10052
10053 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10054 .build()
10055 .await
10056 .unwrap();
10057
10058 let mut create_ns_req = CreateNamespaceRequest::new();
10060 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10061 namespace.create_namespace(create_ns_req).await.unwrap();
10062
10063 let schema = create_test_schema();
10065 let ipc_data = create_test_ipc_data(&schema);
10066
10067 let mut create_req = CreateTableRequest::new();
10068 create_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10069 namespace
10070 .create_table(create_req, bytes::Bytes::from(ipc_data))
10071 .await
10072 .unwrap();
10073
10074 let mut deregister_req = DeregisterTableRequest::new();
10076 deregister_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10077 let response = namespace.deregister_table(deregister_req).await.unwrap();
10078
10079 assert!(
10081 response.location.is_some(),
10082 "Deregister should return location"
10083 );
10084 let location = response.location.as_ref().unwrap();
10085 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10088 .expect("Failed to convert temp path to URL");
10089 let expected_prefix = expected_url.to_string();
10090 assert!(
10091 location.starts_with(&expected_prefix),
10092 "Location should start with '{}', got: {}",
10093 expected_prefix,
10094 location
10095 );
10096 assert!(
10097 location.contains("test_ns") && location.contains("test_table"),
10098 "Location should contain namespace and table name: {}",
10099 location
10100 );
10101 assert_eq!(
10102 response.id,
10103 Some(vec!["test_ns".to_string(), "test_table".to_string()])
10104 );
10105
10106 let mut exists_req = TableExistsRequest::new();
10108 exists_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10109 assert!(namespace.table_exists(exists_req).await.is_err());
10110 }
10111
10112 #[tokio::test]
10113 async fn test_register_without_manifest_fails() {
10114 use lance_namespace::models::RegisterTableRequest;
10115
10116 let temp_dir = TempStdDir::default();
10117 let temp_path = temp_dir.to_str().unwrap();
10118
10119 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10121 .manifest_enabled(false)
10122 .build()
10123 .await
10124 .unwrap();
10125
10126 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10128 register_req.id = Some(vec!["test_table".to_string()]);
10129 let result = namespace.register_table(register_req).await;
10130 assert!(result.is_err());
10131 assert!(
10132 result
10133 .unwrap_err()
10134 .to_string()
10135 .contains("manifest mode is enabled")
10136 );
10137
10138 }
10141
10142 #[tokio::test]
10143 async fn test_register_table_rejects_absolute_uri() {
10144 use lance_namespace::models::RegisterTableRequest;
10145
10146 let temp_dir = TempStdDir::default();
10147 let temp_path = temp_dir.to_str().unwrap();
10148
10149 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10150 .build()
10151 .await
10152 .unwrap();
10153
10154 let mut register_req = RegisterTableRequest::new("s3://bucket/table.lance".to_string());
10156 register_req.id = Some(vec!["test_table".to_string()]);
10157 let result = namespace.register_table(register_req).await;
10158 assert!(result.is_err());
10159 let err_msg = result.unwrap_err().to_string();
10160 assert!(err_msg.contains("Absolute URIs are not allowed"));
10161 }
10162
10163 #[tokio::test]
10164 async fn test_register_table_rejects_absolute_path() {
10165 use lance_namespace::models::RegisterTableRequest;
10166
10167 let temp_dir = TempStdDir::default();
10168 let temp_path = temp_dir.to_str().unwrap();
10169
10170 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10171 .build()
10172 .await
10173 .unwrap();
10174
10175 let mut register_req = RegisterTableRequest::new("/tmp/table.lance".to_string());
10177 register_req.id = Some(vec!["test_table".to_string()]);
10178 let result = namespace.register_table(register_req).await;
10179 assert!(result.is_err());
10180 let err_msg = result.unwrap_err().to_string();
10181 assert!(err_msg.contains("Absolute paths are not allowed"));
10182 }
10183
10184 #[tokio::test]
10185 async fn test_register_table_rejects_path_traversal() {
10186 use lance_namespace::models::RegisterTableRequest;
10187
10188 let temp_dir = TempStdDir::default();
10189 let temp_path = temp_dir.to_str().unwrap();
10190
10191 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10192 .build()
10193 .await
10194 .unwrap();
10195
10196 let mut register_req = RegisterTableRequest::new("../outside/table.lance".to_string());
10198 register_req.id = Some(vec!["test_table".to_string()]);
10199 let result = namespace.register_table(register_req).await;
10200 assert!(result.is_err());
10201 let err_msg = result.unwrap_err().to_string();
10202 assert!(err_msg.contains("Path traversal is not allowed"));
10203 }
10204
10205 #[tokio::test]
10206 async fn test_namespace_write() {
10207 use arrow::array::Int32Array;
10208 use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
10209 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
10210 use lance::dataset::{Dataset, WriteMode, WriteParams};
10211 use lance_namespace::LanceNamespace;
10212
10213 let (namespace, _temp_dir) = create_test_namespace().await;
10214 let namespace = Arc::new(namespace) as Arc<dyn LanceNamespace>;
10215
10216 let table_id = vec!["test_ns".to_string(), "test_table".to_string()];
10218 let schema = Arc::new(ArrowSchema::new(vec![
10219 ArrowField::new("a", DataType::Int32, false),
10220 ArrowField::new("b", DataType::Int32, false),
10221 ]));
10222
10223 let data1 = RecordBatch::try_new(
10225 schema.clone(),
10226 vec![
10227 Arc::new(Int32Array::from(vec![1, 2, 3])),
10228 Arc::new(Int32Array::from(vec![10, 20, 30])),
10229 ],
10230 )
10231 .unwrap();
10232
10233 let reader1 = RecordBatchIterator::new(vec![data1].into_iter().map(Ok), schema.clone());
10234 let dataset =
10235 Dataset::write_into_namespace(reader1, namespace.clone(), table_id.clone(), None)
10236 .await
10237 .unwrap();
10238
10239 assert_eq!(dataset.count_rows(None).await.unwrap(), 3);
10240 assert_eq!(dataset.version().version, 1);
10241
10242 let data2 = RecordBatch::try_new(
10244 schema.clone(),
10245 vec![
10246 Arc::new(Int32Array::from(vec![4, 5])),
10247 Arc::new(Int32Array::from(vec![40, 50])),
10248 ],
10249 )
10250 .unwrap();
10251
10252 let params_append = WriteParams {
10253 mode: WriteMode::Append,
10254 ..Default::default()
10255 };
10256
10257 let reader2 = RecordBatchIterator::new(vec![data2].into_iter().map(Ok), schema.clone());
10258 let dataset = Dataset::write_into_namespace(
10259 reader2,
10260 namespace.clone(),
10261 table_id.clone(),
10262 Some(params_append),
10263 )
10264 .await
10265 .unwrap();
10266
10267 assert_eq!(dataset.count_rows(None).await.unwrap(), 5);
10268 assert_eq!(dataset.version().version, 2);
10269
10270 let data3 = RecordBatch::try_new(
10272 schema.clone(),
10273 vec![
10274 Arc::new(Int32Array::from(vec![100, 200])),
10275 Arc::new(Int32Array::from(vec![1000, 2000])),
10276 ],
10277 )
10278 .unwrap();
10279
10280 let params_overwrite = WriteParams {
10281 mode: WriteMode::Overwrite,
10282 ..Default::default()
10283 };
10284
10285 let reader3 = RecordBatchIterator::new(vec![data3].into_iter().map(Ok), schema.clone());
10286 let dataset = Dataset::write_into_namespace(
10287 reader3,
10288 namespace.clone(),
10289 table_id.clone(),
10290 Some(params_overwrite),
10291 )
10292 .await
10293 .unwrap();
10294
10295 assert_eq!(dataset.count_rows(None).await.unwrap(), 2);
10296 assert_eq!(dataset.version().version, 3);
10297
10298 let result = dataset.scan().try_into_batch().await.unwrap();
10300 let a_col = result
10301 .column_by_name("a")
10302 .unwrap()
10303 .as_any()
10304 .downcast_ref::<Int32Array>()
10305 .unwrap();
10306 assert_eq!(a_col.values(), &[100, 200]);
10307 }
10308
10309 #[tokio::test]
10314 async fn test_declare_table_v1_mode() {
10315 use lance_namespace::models::{
10316 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
10317 };
10318
10319 let temp_dir = TempStdDir::default();
10320 let temp_path = temp_dir.to_str().unwrap();
10321
10322 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10324 .manifest_enabled(false)
10325 .build()
10326 .await
10327 .unwrap();
10328
10329 let mut declare_req = DeclareTableRequest::new();
10331 declare_req.id = Some(vec!["test_table".to_string()]);
10332 let response = namespace.declare_table(declare_req).await.unwrap();
10333
10334 assert!(response.location.is_some());
10336 let location = response.location.as_ref().unwrap();
10337 assert!(location.ends_with("test_table.lance"));
10338
10339 let mut exists_req = TableExistsRequest::new();
10341 exists_req.id = Some(vec!["test_table".to_string()]);
10342 assert!(namespace.table_exists(exists_req).await.is_ok());
10343
10344 let mut describe_req = DescribeTableRequest::new();
10346 describe_req.id = Some(vec!["test_table".to_string()]);
10347 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10348 assert!(describe_response.location.is_some());
10349 assert!(describe_response.version.is_none()); assert!(describe_response.schema.is_none()); assert_eq!(describe_response.is_only_declared, None);
10352
10353 let mut describe_req = DescribeTableRequest::new();
10354 describe_req.id = Some(vec!["test_table".to_string()]);
10355 describe_req.check_declared = Some(true);
10356 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10357 assert_eq!(describe_response.is_only_declared, Some(true));
10358
10359 let mut list_req = ListTablesRequest::new();
10360 list_req.id = Some(vec![]);
10361 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
10362 assert_eq!(list_response.tables, vec!["test_table".to_string()]);
10363
10364 list_req.include_declared = Some(false);
10365 let list_response = namespace.list_tables(list_req).await.unwrap();
10366 assert!(list_response.tables.is_empty());
10367 }
10368
10369 #[tokio::test]
10370 async fn test_insert_into_declared_table_promotes_it_from_declared_state() {
10371 use lance_namespace::models::{
10372 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest,
10373 };
10374
10375 let temp_dir = TempStdDir::default();
10376 let temp_path = temp_dir.to_str().unwrap();
10377
10378 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10379 .manifest_enabled(false)
10380 .build()
10381 .await
10382 .unwrap();
10383
10384 let mut declare_req = DeclareTableRequest::new();
10385 declare_req.id = Some(vec!["test_table".to_string()]);
10386 namespace.declare_table(declare_req).await.unwrap();
10387
10388 let schema = create_test_schema();
10389 let ipc_data = create_test_ipc_data(&schema);
10390 let mut insert_req = InsertIntoTableRequest::new();
10391 insert_req.id = Some(vec!["test_table".to_string()]);
10392 namespace
10393 .insert_into_table(insert_req, bytes::Bytes::from(ipc_data))
10394 .await
10395 .unwrap();
10396
10397 let mut describe_req = DescribeTableRequest::new();
10398 describe_req.id = Some(vec!["test_table".to_string()]);
10399 describe_req.load_detailed_metadata = Some(true);
10400 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10401
10402 assert_eq!(describe_response.is_only_declared, Some(false));
10403 assert_eq!(describe_response.version, Some(1));
10404 assert!(describe_response.schema.is_some());
10405
10406 let mut list_req = ListTablesRequest::new();
10407 list_req.id = Some(vec![]);
10408 list_req.include_declared = Some(false);
10409 assert_eq!(
10410 namespace.list_tables(list_req).await.unwrap().tables,
10411 vec!["test_table".to_string()]
10412 );
10413 }
10414
10415 #[tokio::test]
10416 async fn test_create_table_after_declare_table_v1_mode_creates_table() {
10417 use lance_namespace::models::{
10418 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
10419 };
10420
10421 let temp_dir = TempStdDir::default();
10422 let temp_path = temp_dir.to_str().unwrap();
10423
10424 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10425 .manifest_enabled(false)
10426 .build()
10427 .await
10428 .unwrap();
10429
10430 let mut declare_req = DeclareTableRequest::new();
10431 declare_req.id = Some(vec!["test_table".to_string()]);
10432 namespace.declare_table(declare_req).await.unwrap();
10433
10434 let mut create_req = CreateTableRequest::new();
10435 create_req.id = Some(vec!["test_table".to_string()]);
10436 let response = namespace
10437 .create_table(
10438 create_req,
10439 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10440 )
10441 .await
10442 .unwrap();
10443
10444 assert_eq!(response.version, Some(1));
10445
10446 let mut describe_req = DescribeTableRequest::new();
10447 describe_req.id = Some(vec!["test_table".to_string()]);
10448 describe_req.load_detailed_metadata = Some(true);
10449 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10450 assert_eq!(describe_response.is_only_declared, Some(false));
10451 assert_eq!(describe_response.version, Some(1));
10452
10453 let mut list_req = ListTablesRequest::new();
10454 list_req.id = Some(vec![]);
10455 list_req.include_declared = Some(false);
10456 assert_eq!(
10457 namespace.list_tables(list_req).await.unwrap().tables,
10458 vec!["test_table".to_string()]
10459 );
10460 }
10461
10462 #[tokio::test]
10463 async fn test_insert_into_declared_table_with_manifest_promotes_it() {
10464 use lance_namespace::models::{
10465 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest, ListTablesRequest,
10466 };
10467
10468 let temp_dir = TempStdDir::default();
10469 let temp_path = temp_dir.to_str().unwrap();
10470
10471 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10472 .manifest_enabled(true)
10473 .dir_listing_enabled(false)
10474 .build()
10475 .await
10476 .unwrap();
10477
10478 let mut declare_req = DeclareTableRequest::new();
10479 declare_req.id = Some(vec!["test_table".to_string()]);
10480 namespace.declare_table(declare_req).await.unwrap();
10481
10482 let mut insert_req = InsertIntoTableRequest::new();
10483 insert_req.id = Some(vec!["test_table".to_string()]);
10484 namespace
10485 .insert_into_table(
10486 insert_req,
10487 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10488 )
10489 .await
10490 .unwrap();
10491
10492 let mut describe_req = DescribeTableRequest::new();
10493 describe_req.id = Some(vec!["test_table".to_string()]);
10494 describe_req.load_detailed_metadata = Some(true);
10495 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10496 assert_eq!(describe_response.is_only_declared, Some(false));
10497 assert_eq!(describe_response.version, Some(1));
10498
10499 let mut list_req = ListTablesRequest::new();
10500 list_req.id = Some(vec![]);
10501 list_req.include_declared = Some(false);
10502 assert_eq!(
10503 namespace.list_tables(list_req).await.unwrap().tables,
10504 vec!["test_table".to_string()]
10505 );
10506 }
10507
10508 #[tokio::test]
10509 async fn test_create_table_after_declare_table_with_manifest_creates_table() {
10510 use lance_namespace::models::{
10511 CreateTableRequest, DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
10512 };
10513
10514 let temp_dir = TempStdDir::default();
10515 let temp_path = temp_dir.to_str().unwrap();
10516
10517 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10518 .manifest_enabled(true)
10519 .dir_listing_enabled(false)
10520 .build()
10521 .await
10522 .unwrap();
10523
10524 let mut declare_req = DeclareTableRequest::new();
10525 declare_req.id = Some(vec!["test_table".to_string()]);
10526 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10527 namespace.declare_table(declare_req).await.unwrap();
10528
10529 let mut create_req = CreateTableRequest::new();
10530 create_req.id = Some(vec!["test_table".to_string()]);
10531 create_req.mode = Some("Overwrite".to_string());
10532 let response = namespace
10533 .create_table(
10534 create_req,
10535 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10536 )
10537 .await
10538 .unwrap();
10539
10540 assert_eq!(response.version, Some(1));
10541 assert_eq!(
10542 response
10543 .properties
10544 .as_ref()
10545 .and_then(|properties| properties.get("owner")),
10546 Some(&"alice".to_string())
10547 );
10548
10549 let mut describe_req = DescribeTableRequest::new();
10550 describe_req.id = Some(vec!["test_table".to_string()]);
10551 describe_req.load_detailed_metadata = Some(true);
10552 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10553 assert_eq!(describe_response.is_only_declared, Some(false));
10554 assert_eq!(describe_response.version, Some(1));
10555 assert_eq!(
10556 describe_response
10557 .properties
10558 .as_ref()
10559 .and_then(|properties| properties.get("owner")),
10560 Some(&"alice".to_string())
10561 );
10562
10563 let mut list_req = ListTablesRequest::new();
10564 list_req.id = Some(vec![]);
10565 list_req.include_declared = Some(false);
10566 assert_eq!(
10567 namespace.list_tables(list_req).await.unwrap().tables,
10568 vec!["test_table".to_string()]
10569 );
10570 }
10571
10572 #[tokio::test]
10573 async fn test_create_table_after_declare_table_with_manifest_rejects_new_properties() {
10574 use lance_namespace::models::{CreateTableRequest, DeclareTableRequest};
10575
10576 let temp_dir = TempStdDir::default();
10577 let temp_path = temp_dir.to_str().unwrap();
10578
10579 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10580 .manifest_enabled(true)
10581 .dir_listing_enabled(false)
10582 .build()
10583 .await
10584 .unwrap();
10585
10586 let mut declare_req = DeclareTableRequest::new();
10587 declare_req.id = Some(vec!["test_table".to_string()]);
10588 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10589 namespace.declare_table(declare_req).await.unwrap();
10590
10591 let mut create_req = CreateTableRequest::new();
10592 create_req.id = Some(vec!["test_table".to_string()]);
10593 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
10594
10595 let result = namespace
10596 .create_table(
10597 create_req,
10598 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10599 )
10600 .await;
10601
10602 assert!(result.is_err());
10603 assert!(
10604 result
10605 .unwrap_err()
10606 .to_string()
10607 .contains("cannot set properties for already declared table")
10608 );
10609 }
10610
10611 #[tokio::test]
10612 async fn test_create_table_with_manifest_exist_ok_keeps_existing_table() {
10613 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
10614
10615 let temp_dir = TempStdDir::default();
10616 let temp_path = temp_dir.to_str().unwrap();
10617
10618 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10619 .manifest_enabled(true)
10620 .dir_listing_enabled(false)
10621 .build()
10622 .await
10623 .unwrap();
10624
10625 let mut create_req = CreateTableRequest::new();
10626 create_req.id = Some(vec!["test_table".to_string()]);
10627 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10628 namespace
10629 .create_table(
10630 create_req,
10631 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10632 )
10633 .await
10634 .unwrap();
10635
10636 let mut create_req = CreateTableRequest::new();
10637 create_req.id = Some(vec!["test_table".to_string()]);
10638 create_req.mode = Some("ExistOk".to_string());
10639 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
10640 let response = namespace
10641 .create_table(
10642 create_req,
10643 bytes::Bytes::from(create_single_row_test_ipc_data()),
10644 )
10645 .await
10646 .unwrap();
10647
10648 assert_eq!(
10649 response
10650 .properties
10651 .as_ref()
10652 .and_then(|properties| properties.get("owner")),
10653 Some(&"alice".to_string())
10654 );
10655 assert_eq!(
10656 open_dataset(&namespace, "test_table")
10657 .await
10658 .count_rows(None)
10659 .await
10660 .unwrap(),
10661 2
10662 );
10663
10664 let mut describe_req = DescribeTableRequest::new();
10665 describe_req.id = Some(vec!["test_table".to_string()]);
10666 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10667 assert_eq!(
10668 describe_response
10669 .properties
10670 .as_ref()
10671 .and_then(|properties| properties.get("owner")),
10672 Some(&"alice".to_string())
10673 );
10674 }
10675
10676 #[tokio::test]
10677 async fn test_create_table_with_manifest_overwrite_replaces_existing_table() {
10678 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
10679
10680 let temp_dir = TempStdDir::default();
10681 let temp_path = temp_dir.to_str().unwrap();
10682
10683 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10684 .manifest_enabled(true)
10685 .dir_listing_enabled(false)
10686 .build()
10687 .await
10688 .unwrap();
10689
10690 let mut create_req = CreateTableRequest::new();
10691 create_req.id = Some(vec!["test_table".to_string()]);
10692 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10693 namespace
10694 .create_table(
10695 create_req,
10696 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10697 )
10698 .await
10699 .unwrap();
10700
10701 let mut create_req = CreateTableRequest::new();
10702 create_req.id = Some(vec!["test_table".to_string()]);
10703 create_req.mode = Some("overwrite".to_string());
10704 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
10705 let response = namespace
10706 .create_table(
10707 create_req,
10708 bytes::Bytes::from(create_single_row_test_ipc_data()),
10709 )
10710 .await
10711 .unwrap();
10712
10713 assert_eq!(response.version, Some(2));
10714 assert_eq!(
10715 response
10716 .properties
10717 .as_ref()
10718 .and_then(|properties| properties.get("owner")),
10719 Some(&"bob".to_string())
10720 );
10721 assert_eq!(
10722 open_dataset(&namespace, "test_table")
10723 .await
10724 .count_rows(None)
10725 .await
10726 .unwrap(),
10727 1
10728 );
10729
10730 let mut describe_req = DescribeTableRequest::new();
10731 describe_req.id = Some(vec!["test_table".to_string()]);
10732 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10733 assert_eq!(
10734 describe_response
10735 .properties
10736 .as_ref()
10737 .and_then(|properties| properties.get("owner")),
10738 Some(&"bob".to_string())
10739 );
10740 }
10741
10742 #[tokio::test]
10743 async fn test_create_table_with_manifest_invalid_mode_rejected() {
10744 use lance_namespace::models::CreateTableRequest;
10745
10746 let temp_dir = TempStdDir::default();
10747 let temp_path = temp_dir.to_str().unwrap();
10748
10749 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10750 .manifest_enabled(true)
10751 .dir_listing_enabled(false)
10752 .build()
10753 .await
10754 .unwrap();
10755
10756 let mut create_req = CreateTableRequest::new();
10757 create_req.id = Some(vec!["test_table".to_string()]);
10758 create_req.mode = Some("append".to_string());
10759 let result = namespace
10760 .create_table(
10761 create_req,
10762 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10763 )
10764 .await;
10765
10766 assert!(result.is_err());
10767 assert!(
10768 result
10769 .unwrap_err()
10770 .to_string()
10771 .contains("Unsupported create_table mode")
10772 );
10773 }
10774
10775 #[tokio::test]
10776 async fn test_merge_insert_into_declared_table_v1_mode_creates_table() {
10777 use lance_namespace::models::{
10778 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
10779 MergeInsertIntoTableRequest,
10780 };
10781
10782 let temp_dir = TempStdDir::default();
10783 let temp_path = temp_dir.to_str().unwrap();
10784
10785 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10786 .manifest_enabled(false)
10787 .build()
10788 .await
10789 .unwrap();
10790
10791 let mut declare_req = DeclareTableRequest::new();
10792 declare_req.id = Some(vec!["test_table".to_string()]);
10793 namespace.declare_table(declare_req).await.unwrap();
10794
10795 let mut merge_req = MergeInsertIntoTableRequest::new();
10796 merge_req.id = Some(vec!["test_table".to_string()]);
10797 merge_req.on = Some("id".to_string());
10798 let response = namespace
10799 .merge_insert_into_table(
10800 merge_req,
10801 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10802 )
10803 .await
10804 .unwrap();
10805
10806 assert_eq!(response.num_inserted_rows, Some(2));
10807 assert_eq!(response.num_updated_rows, Some(0));
10808
10809 let mut describe_req = DescribeTableRequest::new();
10810 describe_req.id = Some(vec!["test_table".to_string()]);
10811 describe_req.load_detailed_metadata = Some(true);
10812 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10813 assert_eq!(describe_response.is_only_declared, Some(false));
10814 assert_eq!(describe_response.version, Some(1));
10815
10816 let mut list_req = ListTablesRequest::new();
10817 list_req.id = Some(vec![]);
10818 list_req.include_declared = Some(false);
10819 assert_eq!(
10820 namespace.list_tables(list_req).await.unwrap().tables,
10821 vec!["test_table".to_string()]
10822 );
10823 }
10824
10825 #[tokio::test]
10826 async fn test_merge_insert_into_declared_table_with_manifest_creates_table() {
10827 use lance_namespace::models::{
10828 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
10829 MergeInsertIntoTableRequest,
10830 };
10831
10832 let temp_dir = TempStdDir::default();
10833 let temp_path = temp_dir.to_str().unwrap();
10834
10835 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10836 .manifest_enabled(true)
10837 .dir_listing_enabled(false)
10838 .build()
10839 .await
10840 .unwrap();
10841
10842 let mut declare_req = DeclareTableRequest::new();
10843 declare_req.id = Some(vec!["test_table".to_string()]);
10844 namespace.declare_table(declare_req).await.unwrap();
10845
10846 let mut merge_req = MergeInsertIntoTableRequest::new();
10847 merge_req.id = Some(vec!["test_table".to_string()]);
10848 merge_req.on = Some("id".to_string());
10849 let response = namespace
10850 .merge_insert_into_table(
10851 merge_req,
10852 bytes::Bytes::from(create_non_empty_test_ipc_data()),
10853 )
10854 .await
10855 .unwrap();
10856
10857 assert_eq!(response.num_inserted_rows, Some(2));
10858 assert_eq!(response.num_updated_rows, Some(0));
10859
10860 let mut describe_req = DescribeTableRequest::new();
10861 describe_req.id = Some(vec!["test_table".to_string()]);
10862 describe_req.load_detailed_metadata = Some(true);
10863 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10864 assert_eq!(describe_response.is_only_declared, Some(false));
10865 assert_eq!(describe_response.version, Some(1));
10866
10867 let mut list_req = ListTablesRequest::new();
10868 list_req.id = Some(vec![]);
10869 list_req.include_declared = Some(false);
10870 assert_eq!(
10871 namespace.list_tables(list_req).await.unwrap().tables,
10872 vec!["test_table".to_string()]
10873 );
10874 }
10875
10876 #[tokio::test]
10877 async fn test_declare_table_with_manifest() {
10878 use lance_namespace::models::{
10879 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
10880 };
10881
10882 let temp_dir = TempStdDir::default();
10883 let temp_path = temp_dir.to_str().unwrap();
10884
10885 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10887 .manifest_enabled(true)
10888 .dir_listing_enabled(false)
10889 .build()
10890 .await
10891 .unwrap();
10892
10893 let mut declare_req = DeclareTableRequest::new();
10895 declare_req.id = Some(vec!["test_table".to_string()]);
10896 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
10897 let response = namespace.declare_table(declare_req).await.unwrap();
10898
10899 assert!(response.location.is_some());
10901 assert_eq!(
10902 response
10903 .properties
10904 .as_ref()
10905 .and_then(|properties| properties.get("owner")),
10906 Some(&"alice".to_string())
10907 );
10908
10909 let mut exists_req = TableExistsRequest::new();
10911 exists_req.id = Some(vec!["test_table".to_string()]);
10912 assert!(namespace.table_exists(exists_req).await.is_ok());
10913
10914 let mut describe_req = DescribeTableRequest::new();
10915 describe_req.id = Some(vec!["test_table".to_string()]);
10916 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10917 assert_eq!(describe_response.is_only_declared, None);
10918
10919 let mut describe_req = DescribeTableRequest::new();
10920 describe_req.id = Some(vec!["test_table".to_string()]);
10921 describe_req.check_declared = Some(true);
10922 let describe_response = namespace.describe_table(describe_req).await.unwrap();
10923 assert_eq!(describe_response.is_only_declared, Some(true));
10924 assert_eq!(
10925 describe_response
10926 .properties
10927 .as_ref()
10928 .and_then(|properties| properties.get("owner")),
10929 Some(&"alice".to_string())
10930 );
10931
10932 let mut list_req = ListTablesRequest::new();
10933 list_req.id = Some(vec![]);
10934 assert_eq!(
10935 namespace
10936 .list_tables(list_req.clone())
10937 .await
10938 .unwrap()
10939 .tables,
10940 vec!["test_table".to_string()]
10941 );
10942 list_req.include_declared = Some(false);
10943 assert!(
10944 namespace
10945 .list_tables(list_req)
10946 .await
10947 .unwrap()
10948 .tables
10949 .is_empty()
10950 );
10951 }
10952
10953 #[tokio::test]
10954 async fn test_declare_table_with_manifest_marker_already_exists() {
10955 use lance_namespace::error::ErrorCode;
10958 use lance_namespace::models::DeclareTableRequest;
10959
10960 let temp_dir = TempStdDir::default();
10961 let temp_path = temp_dir.to_str().unwrap();
10962
10963 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10964 .manifest_enabled(true)
10965 .dir_listing_enabled(true)
10966 .build()
10967 .await
10968 .unwrap();
10969
10970 let table_dir = temp_dir.join("test_table.lance");
10971 std::fs::create_dir_all(&table_dir).unwrap();
10972 std::fs::write(table_dir.join(".lance-reserved"), b"reserved").unwrap();
10973
10974 let mut declare_req = DeclareTableRequest::new();
10975 declare_req.id = Some(vec!["test_table".to_string()]);
10976 let err = namespace
10977 .declare_table(declare_req)
10978 .await
10979 .expect_err("declare with existing marker must fail");
10980 let msg = err.to_string();
10981 assert!(
10982 msg.contains("already exists") || msg.contains("TableAlreadyExists"),
10983 "expected TableAlreadyExists, got: {msg}"
10984 );
10985 assert_eq!(
10986 mutation_error_code(err),
10987 ErrorCode::TableAlreadyExists,
10988 "expected TableAlreadyExists error code"
10989 );
10990 }
10991
10992 #[tokio::test]
10993 async fn test_declare_table_when_table_exists() {
10994 use lance_namespace::models::DeclareTableRequest;
10995
10996 let temp_dir = TempStdDir::default();
10997 let temp_path = temp_dir.to_str().unwrap();
10998
10999 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11000 .manifest_enabled(false)
11001 .build()
11002 .await
11003 .unwrap();
11004
11005 let schema = create_test_schema();
11007 let ipc_data = create_test_ipc_data(&schema);
11008 let mut create_req = CreateTableRequest::new();
11009 create_req.id = Some(vec!["test_table".to_string()]);
11010 namespace
11011 .create_table(create_req, bytes::Bytes::from(ipc_data))
11012 .await
11013 .unwrap();
11014
11015 let mut declare_req = DeclareTableRequest::new();
11017 declare_req.id = Some(vec!["test_table".to_string()]);
11018 let result = namespace.declare_table(declare_req).await;
11019 assert!(result.is_err());
11020 }
11021
11022 #[tokio::test]
11027 async fn test_deregister_table_v1_mode() {
11028 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
11029
11030 let temp_dir = TempStdDir::default();
11031 let temp_path = temp_dir.to_str().unwrap();
11032
11033 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11035 .manifest_enabled(false)
11036 .dir_listing_enabled(true)
11037 .build()
11038 .await
11039 .unwrap();
11040
11041 let schema = create_test_schema();
11043 let ipc_data = create_test_ipc_data(&schema);
11044 let mut create_req = CreateTableRequest::new();
11045 create_req.id = Some(vec!["test_table".to_string()]);
11046 namespace
11047 .create_table(create_req, bytes::Bytes::from(ipc_data))
11048 .await
11049 .unwrap();
11050
11051 let mut exists_req = TableExistsRequest::new();
11053 exists_req.id = Some(vec!["test_table".to_string()]);
11054 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
11055
11056 let mut deregister_req = DeregisterTableRequest::new();
11058 deregister_req.id = Some(vec!["test_table".to_string()]);
11059 let response = namespace.deregister_table(deregister_req).await.unwrap();
11060
11061 assert!(response.location.is_some());
11063 let location = response.location.as_ref().unwrap();
11064 assert!(location.contains("test_table"));
11065
11066 let result = namespace.table_exists(exists_req).await;
11068 assert!(result.is_err());
11069 assert!(result.unwrap_err().to_string().contains("deregistered"));
11070
11071 let dataset = Dataset::open(location).await;
11073 assert!(dataset.is_ok(), "Physical table data should still exist");
11074 }
11075
11076 #[tokio::test]
11077 async fn test_deregister_table_v1_already_deregistered() {
11078 use lance_namespace::models::DeregisterTableRequest;
11079
11080 let temp_dir = TempStdDir::default();
11081 let temp_path = temp_dir.to_str().unwrap();
11082
11083 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11084 .manifest_enabled(false)
11085 .dir_listing_enabled(true)
11086 .build()
11087 .await
11088 .unwrap();
11089
11090 let schema = create_test_schema();
11092 let ipc_data = create_test_ipc_data(&schema);
11093 let mut create_req = CreateTableRequest::new();
11094 create_req.id = Some(vec!["test_table".to_string()]);
11095 namespace
11096 .create_table(create_req, bytes::Bytes::from(ipc_data))
11097 .await
11098 .unwrap();
11099
11100 let mut deregister_req = DeregisterTableRequest::new();
11102 deregister_req.id = Some(vec!["test_table".to_string()]);
11103 namespace
11104 .deregister_table(deregister_req.clone())
11105 .await
11106 .unwrap();
11107
11108 let result = namespace.deregister_table(deregister_req).await;
11110 assert!(result.is_err());
11111 assert!(
11112 result
11113 .unwrap_err()
11114 .to_string()
11115 .contains("already deregistered")
11116 );
11117 }
11118
11119 #[tokio::test]
11124 async fn test_list_tables_skips_deregistered_v1() {
11125 use lance_namespace::models::DeregisterTableRequest;
11126
11127 let temp_dir = TempStdDir::default();
11128 let temp_path = temp_dir.to_str().unwrap();
11129
11130 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11131 .manifest_enabled(false)
11132 .dir_listing_enabled(true)
11133 .build()
11134 .await
11135 .unwrap();
11136
11137 let schema = create_test_schema();
11139 let ipc_data = create_test_ipc_data(&schema);
11140
11141 let mut create_req1 = CreateTableRequest::new();
11142 create_req1.id = Some(vec!["table1".to_string()]);
11143 namespace
11144 .create_table(create_req1, bytes::Bytes::from(ipc_data.clone()))
11145 .await
11146 .unwrap();
11147
11148 let mut create_req2 = CreateTableRequest::new();
11149 create_req2.id = Some(vec!["table2".to_string()]);
11150 namespace
11151 .create_table(create_req2, bytes::Bytes::from(ipc_data))
11152 .await
11153 .unwrap();
11154
11155 let mut list_req = ListTablesRequest::new();
11157 list_req.id = Some(vec![]);
11158 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
11159 assert_eq!(list_response.tables.len(), 2);
11160
11161 let mut deregister_req = DeregisterTableRequest::new();
11163 deregister_req.id = Some(vec!["table1".to_string()]);
11164 namespace.deregister_table(deregister_req).await.unwrap();
11165
11166 let list_response = namespace.list_tables(list_req).await.unwrap();
11168 assert_eq!(list_response.tables.len(), 1);
11169 assert!(list_response.tables.contains(&"table2".to_string()));
11170 assert!(!list_response.tables.contains(&"table1".to_string()));
11171 }
11172
11173 #[tokio::test]
11178 async fn test_describe_table_fails_for_deregistered_v1() {
11179 use lance_namespace::models::{DeregisterTableRequest, DescribeTableRequest};
11180
11181 let temp_dir = TempStdDir::default();
11182 let temp_path = temp_dir.to_str().unwrap();
11183
11184 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11185 .manifest_enabled(false)
11186 .dir_listing_enabled(true)
11187 .build()
11188 .await
11189 .unwrap();
11190
11191 let schema = create_test_schema();
11193 let ipc_data = create_test_ipc_data(&schema);
11194 let mut create_req = CreateTableRequest::new();
11195 create_req.id = Some(vec!["test_table".to_string()]);
11196 namespace
11197 .create_table(create_req, bytes::Bytes::from(ipc_data))
11198 .await
11199 .unwrap();
11200
11201 let mut describe_req = DescribeTableRequest::new();
11203 describe_req.id = Some(vec!["test_table".to_string()]);
11204 assert!(namespace.describe_table(describe_req.clone()).await.is_ok());
11205
11206 let mut deregister_req = DeregisterTableRequest::new();
11208 deregister_req.id = Some(vec!["test_table".to_string()]);
11209 namespace.deregister_table(deregister_req).await.unwrap();
11210
11211 let result = namespace.describe_table(describe_req).await;
11213 assert!(result.is_err());
11214 let err = result.unwrap_err();
11215 assert!(matches!(err, Error::Namespace { .. }));
11216 let err_msg = err.to_string();
11217 assert!(err_msg.contains("deregistered"));
11218 assert!(err_msg.contains("table id 'test_table'"));
11219 }
11220
11221 #[tokio::test]
11222 async fn test_table_exists_fails_for_deregistered_v1() {
11223 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
11224
11225 let temp_dir = TempStdDir::default();
11226 let temp_path = temp_dir.to_str().unwrap();
11227
11228 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11229 .manifest_enabled(false)
11230 .dir_listing_enabled(true)
11231 .build()
11232 .await
11233 .unwrap();
11234
11235 let schema = create_test_schema();
11237 let ipc_data = create_test_ipc_data(&schema);
11238 let mut create_req = CreateTableRequest::new();
11239 create_req.id = Some(vec!["test_table".to_string()]);
11240 namespace
11241 .create_table(create_req, bytes::Bytes::from(ipc_data))
11242 .await
11243 .unwrap();
11244
11245 let mut exists_req = TableExistsRequest::new();
11247 exists_req.id = Some(vec!["test_table".to_string()]);
11248 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
11249
11250 let mut deregister_req = DeregisterTableRequest::new();
11252 deregister_req.id = Some(vec!["test_table".to_string()]);
11253 namespace.deregister_table(deregister_req).await.unwrap();
11254
11255 let result = namespace.table_exists(exists_req).await;
11257 assert!(result.is_err());
11258 let err = result.unwrap_err();
11259 assert!(matches!(err, Error::Namespace { .. }));
11260 let err_msg = err.to_string();
11261 assert!(err_msg.contains("deregistered"));
11262 assert!(err_msg.contains("table id 'test_table'"));
11263 }
11264
11265 #[tokio::test]
11266 async fn test_atomic_table_status_check() {
11267 let temp_dir = TempStdDir::default();
11271 let temp_path = temp_dir.to_str().unwrap();
11272
11273 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11274 .manifest_enabled(false)
11275 .dir_listing_enabled(true)
11276 .build()
11277 .await
11278 .unwrap();
11279
11280 let schema = create_test_schema();
11282 let ipc_data = create_test_ipc_data(&schema);
11283 let mut create_req = CreateTableRequest::new();
11284 create_req.id = Some(vec!["test_table".to_string()]);
11285 namespace
11286 .create_table(create_req, bytes::Bytes::from(ipc_data))
11287 .await
11288 .unwrap();
11289
11290 let status = namespace.check_table_status("test_table").await.unwrap();
11292 assert!(status.exists);
11293 assert!(!status.is_deregistered);
11294 assert!(!status.has_reserved_file);
11295 }
11296
11297 #[tokio::test]
11298 async fn test_table_version_tracking_enabled_managed_versioning() {
11299 use lance_namespace::models::DescribeTableRequest;
11300
11301 let temp_dir = TempStdDir::default();
11302 let temp_path = temp_dir.to_str().unwrap();
11303
11304 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11306 .table_version_tracking_enabled(true)
11307 .build()
11308 .await
11309 .unwrap();
11310
11311 let schema = create_test_schema();
11313 let ipc_data = create_test_ipc_data(&schema);
11314 let mut create_req = CreateTableRequest::new();
11315 create_req.id = Some(vec!["test_table".to_string()]);
11316 namespace
11317 .create_table(create_req, bytes::Bytes::from(ipc_data))
11318 .await
11319 .unwrap();
11320
11321 let mut describe_req = DescribeTableRequest::new();
11323 describe_req.id = Some(vec!["test_table".to_string()]);
11324 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
11325
11326 assert_eq!(
11328 describe_resp.managed_versioning,
11329 Some(true),
11330 "managed_versioning should be true when table_version_tracking_enabled=true"
11331 );
11332 }
11333
11334 #[tokio::test]
11335 async fn test_table_version_tracking_disabled_no_managed_versioning() {
11336 use lance_namespace::models::DescribeTableRequest;
11337
11338 let temp_dir = TempStdDir::default();
11339 let temp_path = temp_dir.to_str().unwrap();
11340
11341 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11343 .table_version_tracking_enabled(false)
11344 .build()
11345 .await
11346 .unwrap();
11347
11348 let schema = create_test_schema();
11350 let ipc_data = create_test_ipc_data(&schema);
11351 let mut create_req = CreateTableRequest::new();
11352 create_req.id = Some(vec!["test_table".to_string()]);
11353 namespace
11354 .create_table(create_req, bytes::Bytes::from(ipc_data))
11355 .await
11356 .unwrap();
11357
11358 let mut describe_req = DescribeTableRequest::new();
11360 describe_req.id = Some(vec!["test_table".to_string()]);
11361 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
11362
11363 assert!(
11365 describe_resp.managed_versioning.is_none(),
11366 "managed_versioning should be None when table_version_tracking_enabled=false, got: {:?}",
11367 describe_resp.managed_versioning
11368 );
11369 }
11370
11371 #[tokio::test]
11372 async fn test_list_table_versions() {
11373 use arrow::array::{Int32Array, RecordBatchIterator};
11374 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11375 use arrow::record_batch::RecordBatch;
11376 use lance::dataset::{Dataset, WriteMode, WriteParams};
11377 use lance_namespace::models::{CreateNamespaceRequest, ListTableVersionsRequest};
11378
11379 let temp_dir = TempStrDir::default();
11380 let temp_path: &str = &temp_dir;
11381
11382 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11383 DirectoryNamespaceBuilder::new(temp_path)
11384 .table_version_tracking_enabled(true)
11385 .build()
11386 .await
11387 .unwrap(),
11388 );
11389
11390 let mut create_ns_req = CreateNamespaceRequest::new();
11392 create_ns_req.id = Some(vec!["workspace".to_string()]);
11393 namespace.create_namespace(create_ns_req).await.unwrap();
11394
11395 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11397 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
11398 "id",
11399 DataType::Int32,
11400 false,
11401 )]));
11402 let batch = RecordBatch::try_new(
11403 arrow_schema.clone(),
11404 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
11405 )
11406 .unwrap();
11407 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
11408 let write_params = WriteParams {
11409 mode: WriteMode::Create,
11410 ..Default::default()
11411 };
11412 let mut dataset = Dataset::write_into_namespace(
11413 batches,
11414 namespace.clone(),
11415 table_id.clone(),
11416 Some(write_params),
11417 )
11418 .await
11419 .unwrap();
11420
11421 let batch2 = RecordBatch::try_new(
11423 arrow_schema.clone(),
11424 vec![Arc::new(Int32Array::from(vec![100, 200]))],
11425 )
11426 .unwrap();
11427 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
11428 dataset.append(batches, None).await.unwrap();
11429
11430 let batch3 = RecordBatch::try_new(
11432 arrow_schema.clone(),
11433 vec![Arc::new(Int32Array::from(vec![300, 400]))],
11434 )
11435 .unwrap();
11436 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
11437 dataset.append(batches, None).await.unwrap();
11438
11439 let mut list_req = ListTableVersionsRequest::new();
11441 list_req.id = Some(table_id.clone());
11442 let list_resp = namespace.list_table_versions(list_req).await.unwrap();
11443
11444 assert_eq!(
11445 list_resp.versions.len(),
11446 3,
11447 "Should have 3 versions, got: {:?}",
11448 list_resp.versions
11449 );
11450
11451 for expected_version in 1..=3 {
11453 let version = list_resp
11454 .versions
11455 .iter()
11456 .find(|v| v.version == expected_version)
11457 .unwrap_or_else(|| panic!("Expected version {}", expected_version));
11458
11459 assert!(
11460 !version.manifest_path.is_empty(),
11461 "manifest_path should be set for version {}",
11462 expected_version
11463 );
11464 assert!(
11465 version.manifest_path.contains(".manifest"),
11466 "manifest_path should contain .manifest for version {}",
11467 expected_version
11468 );
11469 assert!(
11470 version.manifest_size.is_some(),
11471 "manifest_size should be set for version {}",
11472 expected_version
11473 );
11474 assert!(
11475 version.manifest_size.unwrap() > 0,
11476 "manifest_size should be > 0 for version {}",
11477 expected_version
11478 );
11479 assert!(
11480 version.timestamp_millis.is_some(),
11481 "timestamp_millis should be set for version {}",
11482 expected_version
11483 );
11484 }
11485 }
11486
11487 #[tokio::test]
11488 async fn test_describe_table_version() {
11489 use arrow::array::{Int32Array, RecordBatchIterator};
11490 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11491 use arrow::record_batch::RecordBatch;
11492 use lance::dataset::{Dataset, WriteMode, WriteParams};
11493 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
11494
11495 let temp_dir = TempStrDir::default();
11496 let temp_path: &str = &temp_dir;
11497
11498 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11499 DirectoryNamespaceBuilder::new(temp_path)
11500 .table_version_tracking_enabled(true)
11501 .build()
11502 .await
11503 .unwrap(),
11504 );
11505
11506 let mut create_ns_req = CreateNamespaceRequest::new();
11508 create_ns_req.id = Some(vec!["workspace".to_string()]);
11509 namespace.create_namespace(create_ns_req).await.unwrap();
11510
11511 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11513 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
11514 "id",
11515 DataType::Int32,
11516 false,
11517 )]));
11518 let batch = RecordBatch::try_new(
11519 arrow_schema.clone(),
11520 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
11521 )
11522 .unwrap();
11523 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
11524 let write_params = WriteParams {
11525 mode: WriteMode::Create,
11526 ..Default::default()
11527 };
11528 let mut dataset = Dataset::write_into_namespace(
11529 batches,
11530 namespace.clone(),
11531 table_id.clone(),
11532 Some(write_params),
11533 )
11534 .await
11535 .unwrap();
11536
11537 let batch2 = RecordBatch::try_new(
11539 arrow_schema.clone(),
11540 vec![Arc::new(Int32Array::from(vec![100, 200]))],
11541 )
11542 .unwrap();
11543 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
11544 dataset.append(batches, None).await.unwrap();
11545
11546 let mut describe_req = DescribeTableVersionRequest::new();
11548 describe_req.id = Some(table_id.clone());
11549 describe_req.version = Some(1);
11550 let describe_resp = namespace
11551 .describe_table_version(describe_req)
11552 .await
11553 .unwrap();
11554
11555 let version = &describe_resp.version;
11556 assert_eq!(version.version, 1);
11557 assert!(version.timestamp_millis.is_some());
11558 assert!(
11559 !version.manifest_path.is_empty(),
11560 "manifest_path should be set"
11561 );
11562 assert!(
11563 version.manifest_path.contains(".manifest"),
11564 "manifest_path should contain .manifest"
11565 );
11566 assert!(
11567 version.manifest_size.is_some(),
11568 "manifest_size should be set"
11569 );
11570 assert!(
11571 version.manifest_size.unwrap() > 0,
11572 "manifest_size should be > 0"
11573 );
11574
11575 let mut describe_req = DescribeTableVersionRequest::new();
11577 describe_req.id = Some(table_id.clone());
11578 describe_req.version = Some(2);
11579 let describe_resp = namespace
11580 .describe_table_version(describe_req)
11581 .await
11582 .unwrap();
11583
11584 let version = &describe_resp.version;
11585 assert_eq!(version.version, 2);
11586 assert!(version.timestamp_millis.is_some());
11587 assert!(
11588 !version.manifest_path.is_empty(),
11589 "manifest_path should be set"
11590 );
11591 assert!(
11592 version.manifest_size.is_some(),
11593 "manifest_size should be set"
11594 );
11595 assert!(
11596 version.manifest_size.unwrap() > 0,
11597 "manifest_size should be > 0"
11598 );
11599 }
11600
11601 #[tokio::test]
11602 async fn test_describe_table_version_latest() {
11603 use arrow::array::{Int32Array, RecordBatchIterator};
11604 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11605 use arrow::record_batch::RecordBatch;
11606 use lance::dataset::{Dataset, WriteMode, WriteParams};
11607 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
11608
11609 let temp_dir = TempStrDir::default();
11610 let temp_path: &str = &temp_dir;
11611
11612 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11613 DirectoryNamespaceBuilder::new(temp_path)
11614 .table_version_tracking_enabled(true)
11615 .build()
11616 .await
11617 .unwrap(),
11618 );
11619
11620 let mut create_ns_req = CreateNamespaceRequest::new();
11622 create_ns_req.id = Some(vec!["workspace".to_string()]);
11623 namespace.create_namespace(create_ns_req).await.unwrap();
11624
11625 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
11627 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
11628 "id",
11629 DataType::Int32,
11630 false,
11631 )]));
11632 let batch = RecordBatch::try_new(
11633 arrow_schema.clone(),
11634 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
11635 )
11636 .unwrap();
11637 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
11638 let write_params = WriteParams {
11639 mode: WriteMode::Create,
11640 ..Default::default()
11641 };
11642 let mut dataset = Dataset::write_into_namespace(
11643 batches,
11644 namespace.clone(),
11645 table_id.clone(),
11646 Some(write_params),
11647 )
11648 .await
11649 .unwrap();
11650
11651 let batch2 = RecordBatch::try_new(
11653 arrow_schema.clone(),
11654 vec![Arc::new(Int32Array::from(vec![100, 200]))],
11655 )
11656 .unwrap();
11657 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
11658 dataset.append(batches, None).await.unwrap();
11659
11660 let batch3 = RecordBatch::try_new(
11662 arrow_schema.clone(),
11663 vec![Arc::new(Int32Array::from(vec![300, 400]))],
11664 )
11665 .unwrap();
11666 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
11667 dataset.append(batches, None).await.unwrap();
11668
11669 let mut describe_req = DescribeTableVersionRequest::new();
11671 describe_req.id = Some(table_id.clone());
11672 describe_req.version = None;
11673 let describe_resp = namespace
11674 .describe_table_version(describe_req)
11675 .await
11676 .unwrap();
11677
11678 assert_eq!(describe_resp.version.version, 3);
11680 }
11681
11682 #[tokio::test]
11683 async fn test_create_table_version() {
11684 use futures::TryStreamExt;
11685 use lance::dataset::builder::DatasetBuilder;
11686 use lance_namespace::models::CreateTableVersionRequest;
11687
11688 let temp_dir = TempStrDir::default();
11689 let temp_path: &str = &temp_dir;
11690
11691 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11692 DirectoryNamespaceBuilder::new(temp_path)
11693 .table_version_tracking_enabled(true)
11694 .build()
11695 .await
11696 .unwrap(),
11697 );
11698
11699 let schema = create_test_schema();
11701 let ipc_data = create_test_ipc_data(&schema);
11702 let mut create_req = CreateTableRequest::new();
11703 create_req.id = Some(vec!["test_table".to_string()]);
11704 namespace
11705 .create_table(create_req, bytes::Bytes::from(ipc_data))
11706 .await
11707 .unwrap();
11708
11709 let table_id = vec!["test_table".to_string()];
11711 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
11712 .await
11713 .unwrap()
11714 .load()
11715 .await
11716 .unwrap();
11717
11718 let versions_path = dataset.versions_dir();
11720 let manifest_metas: Vec<_> = dataset
11721 .object_store(None)
11722 .await
11723 .unwrap()
11724 .inner
11725 .list(Some(&versions_path))
11726 .try_collect()
11727 .await
11728 .unwrap();
11729
11730 let manifest_meta = manifest_metas
11731 .iter()
11732 .find(|m| {
11733 m.location
11734 .filename()
11735 .map(|f| f.ends_with(".manifest"))
11736 .unwrap_or(false)
11737 })
11738 .expect("No manifest file found");
11739
11740 let manifest_data = dataset
11742 .object_store(None)
11743 .await
11744 .unwrap()
11745 .inner
11746 .get(&manifest_meta.location)
11747 .await
11748 .unwrap()
11749 .bytes()
11750 .await
11751 .unwrap();
11752
11753 let staging_path = dataset.versions_dir().join("staging_manifest");
11755 dataset
11756 .object_store(None)
11757 .await
11758 .unwrap()
11759 .inner
11760 .put(&staging_path, manifest_data.into())
11761 .await
11762 .unwrap();
11763
11764 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
11767 create_version_req.id = Some(table_id.clone());
11768 create_version_req.naming_scheme = Some("V2".to_string());
11769
11770 let result = namespace.create_table_version(create_version_req).await;
11771 assert!(
11772 result.is_ok(),
11773 "create_table_version should succeed: {:?}",
11774 result
11775 );
11776
11777 let response = result.unwrap();
11779 let version_info = response
11780 .version
11781 .expect("response should contain version info");
11782 let version_2_path = Path::parse(&version_info.manifest_path).unwrap();
11783 let head_result = dataset
11784 .object_store(None)
11785 .await
11786 .unwrap()
11787 .inner
11788 .head(&version_2_path)
11789 .await;
11790 assert!(
11791 head_result.is_ok(),
11792 "Version 2 manifest should exist at {}",
11793 version_2_path
11794 );
11795
11796 let staging_head_result = dataset
11798 .object_store(None)
11799 .await
11800 .unwrap()
11801 .inner
11802 .head(&staging_path)
11803 .await;
11804 assert!(
11805 staging_head_result.is_err(),
11806 "Staging manifest should have been deleted after create_table_version"
11807 );
11808 }
11809
11810 #[tokio::test]
11811 async fn test_create_table_version_idempotent() {
11812 use futures::TryStreamExt;
11815 use lance::dataset::builder::DatasetBuilder;
11816 use lance_namespace::models::CreateTableVersionRequest;
11817
11818 let temp_dir = TempStrDir::default();
11819 let temp_path: &str = &temp_dir;
11820
11821 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11822 DirectoryNamespaceBuilder::new(temp_path)
11823 .table_version_tracking_enabled(true)
11824 .build()
11825 .await
11826 .unwrap(),
11827 );
11828
11829 let schema = create_test_schema();
11830 let ipc_data = create_test_ipc_data(&schema);
11831 let mut create_req = CreateTableRequest::new();
11832 create_req.id = Some(vec!["test_table".to_string()]);
11833 namespace
11834 .create_table(create_req, bytes::Bytes::from(ipc_data))
11835 .await
11836 .unwrap();
11837
11838 let table_id = vec!["test_table".to_string()];
11839 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
11840 .await
11841 .unwrap()
11842 .load()
11843 .await
11844 .unwrap();
11845
11846 let versions_path = dataset.versions_dir();
11847 let manifest_metas: Vec<_> = dataset
11848 .object_store(None)
11849 .await
11850 .unwrap()
11851 .inner
11852 .list(Some(&versions_path))
11853 .try_collect()
11854 .await
11855 .unwrap();
11856
11857 let manifest_meta = manifest_metas
11858 .iter()
11859 .find(|m| {
11860 m.location
11861 .filename()
11862 .map(|f| f.ends_with(".manifest"))
11863 .unwrap_or(false)
11864 })
11865 .expect("No manifest file found");
11866
11867 let manifest_data = dataset
11868 .object_store(None)
11869 .await
11870 .unwrap()
11871 .inner
11872 .get(&manifest_meta.location)
11873 .await
11874 .unwrap()
11875 .bytes()
11876 .await
11877 .unwrap();
11878
11879 let staging_path = dataset.versions_dir().join("staging_manifest");
11880 dataset
11881 .object_store(None)
11882 .await
11883 .unwrap()
11884 .inner
11885 .put(&staging_path, manifest_data.clone().into())
11886 .await
11887 .unwrap();
11888
11889 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
11890 create_version_req.id = Some(table_id.clone());
11891 create_version_req.naming_scheme = Some("V2".to_string());
11892 let first = namespace
11893 .create_table_version(create_version_req)
11894 .await
11895 .expect("first create_table_version should succeed");
11896
11897 let retry_staging = dataset.versions_dir().join("staging_manifest_retry");
11899 dataset
11900 .object_store(None)
11901 .await
11902 .unwrap()
11903 .inner
11904 .put(&retry_staging, manifest_data.into())
11905 .await
11906 .unwrap();
11907
11908 let mut retry_req = CreateTableVersionRequest::new(2, retry_staging.to_string());
11909 retry_req.id = Some(table_id.clone());
11910 retry_req.naming_scheme = Some("V2".to_string());
11911 let second = namespace
11912 .create_table_version(retry_req)
11913 .await
11914 .expect("idempotent retry must succeed");
11915
11916 assert_eq!(
11917 first.version.as_ref().map(|v| v.version),
11918 second.version.as_ref().map(|v| v.version)
11919 );
11920 assert_eq!(
11921 first.version.as_ref().map(|v| &v.manifest_path),
11922 second.version.as_ref().map(|v| &v.manifest_path)
11923 );
11924 }
11925
11926 #[tokio::test]
11927 async fn test_create_table_version_conflict() {
11928 use futures::TryStreamExt;
11930 use lance::dataset::builder::DatasetBuilder;
11931 use lance_namespace::models::CreateTableVersionRequest;
11932
11933 let temp_dir = TempStrDir::default();
11934 let temp_path: &str = &temp_dir;
11935
11936 let namespace: Arc<dyn LanceNamespace> = Arc::new(
11937 DirectoryNamespaceBuilder::new(temp_path)
11938 .table_version_tracking_enabled(true)
11939 .build()
11940 .await
11941 .unwrap(),
11942 );
11943
11944 let schema = create_test_schema();
11946 let ipc_data = create_test_ipc_data(&schema);
11947 let mut create_req = CreateTableRequest::new();
11948 create_req.id = Some(vec!["test_table".to_string()]);
11949 namespace
11950 .create_table(create_req, bytes::Bytes::from(ipc_data))
11951 .await
11952 .unwrap();
11953
11954 let table_id = vec!["test_table".to_string()];
11956 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
11957 .await
11958 .unwrap()
11959 .load()
11960 .await
11961 .unwrap();
11962
11963 let versions_path = dataset.versions_dir();
11965 let manifest_metas: Vec<_> = dataset
11966 .object_store(None)
11967 .await
11968 .unwrap()
11969 .inner
11970 .list(Some(&versions_path))
11971 .try_collect()
11972 .await
11973 .unwrap();
11974
11975 let manifest_meta = manifest_metas
11976 .iter()
11977 .find(|m| {
11978 m.location
11979 .filename()
11980 .map(|f| f.ends_with(".manifest"))
11981 .unwrap_or(false)
11982 })
11983 .expect("No manifest file found");
11984
11985 let manifest_data = dataset
11987 .object_store(None)
11988 .await
11989 .unwrap()
11990 .inner
11991 .get(&manifest_meta.location)
11992 .await
11993 .unwrap()
11994 .bytes()
11995 .await
11996 .unwrap();
11997
11998 let staging_path = dataset.versions_dir().join("staging_manifest");
12000 dataset
12001 .object_store(None)
12002 .await
12003 .unwrap()
12004 .inner
12005 .put(&staging_path, manifest_data.into())
12006 .await
12007 .unwrap();
12008
12009 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12011 create_version_req.id = Some(table_id.clone());
12012 create_version_req.naming_scheme = Some("V2".to_string());
12013 let first_result = namespace.create_table_version(create_version_req).await;
12014 assert!(
12015 first_result.is_ok(),
12016 "First create_table_version for version 2 should succeed: {:?}",
12017 first_result
12018 );
12019
12020 let version_2_path = Path::parse(
12022 &first_result
12023 .unwrap()
12024 .version
12025 .expect("response should contain version info")
12026 .manifest_path,
12027 )
12028 .unwrap();
12029
12030 let conflict_staging = dataset.versions_dir().join("staging_manifest_conflict");
12032 dataset
12033 .object_store(None)
12034 .await
12035 .unwrap()
12036 .inner
12037 .put(
12038 &conflict_staging,
12039 bytes::Bytes::from_static(b"not-a-real-manifest").into(),
12040 )
12041 .await
12042 .unwrap();
12043
12044 let mut create_version_req =
12045 CreateTableVersionRequest::new(2, conflict_staging.to_string());
12046 create_version_req.id = Some(table_id.clone());
12047 create_version_req.naming_scheme = Some("V2".to_string());
12048
12049 let result = namespace.create_table_version(create_version_req).await;
12050 assert!(
12051 result.is_err(),
12052 "create_table_version should fail for existing version with different content"
12053 );
12054 let err = result.unwrap_err().to_string();
12055 assert!(
12056 err.contains("already exists") || err.contains("ConcurrentModification"),
12057 "expected ConcurrentModification, got: {err}"
12058 );
12059
12060 let head_result = dataset
12062 .object_store(None)
12063 .await
12064 .unwrap()
12065 .inner
12066 .head(&version_2_path)
12067 .await;
12068 assert!(
12069 head_result.is_ok(),
12070 "Version 2 manifest should still exist at {}",
12071 version_2_path
12072 );
12073 }
12074
12075 #[tokio::test]
12076 async fn test_create_table_version_cas_rejects_gap() {
12077 use futures::TryStreamExt;
12079 use lance::dataset::builder::DatasetBuilder;
12080 use lance_namespace::models::CreateTableVersionRequest;
12081
12082 let temp_dir = TempStrDir::default();
12083 let temp_path: &str = &temp_dir;
12084
12085 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12086 DirectoryNamespaceBuilder::new(temp_path)
12087 .table_version_tracking_enabled(true)
12088 .build()
12089 .await
12090 .unwrap(),
12091 );
12092
12093 let schema = create_test_schema();
12094 let ipc_data = create_test_ipc_data(&schema);
12095 let mut create_req = CreateTableRequest::new();
12096 create_req.id = Some(vec!["test_table".to_string()]);
12097 namespace
12098 .create_table(create_req, bytes::Bytes::from(ipc_data))
12099 .await
12100 .unwrap();
12101
12102 let table_id = vec!["test_table".to_string()];
12103 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12104 .await
12105 .unwrap()
12106 .load()
12107 .await
12108 .unwrap();
12109
12110 let versions_path = dataset.versions_dir();
12111 let manifest_metas: Vec<_> = dataset
12112 .object_store(None)
12113 .await
12114 .unwrap()
12115 .inner
12116 .list(Some(&versions_path))
12117 .try_collect()
12118 .await
12119 .unwrap();
12120 let manifest_meta = manifest_metas
12121 .iter()
12122 .find(|m| {
12123 m.location
12124 .filename()
12125 .map(|f| f.ends_with(".manifest"))
12126 .unwrap_or(false)
12127 })
12128 .expect("No manifest file found");
12129 let manifest_data = dataset
12130 .object_store(None)
12131 .await
12132 .unwrap()
12133 .inner
12134 .get(&manifest_meta.location)
12135 .await
12136 .unwrap()
12137 .bytes()
12138 .await
12139 .unwrap();
12140
12141 let staging_path = dataset.versions_dir().join("staging_gap");
12142 dataset
12143 .object_store(None)
12144 .await
12145 .unwrap()
12146 .inner
12147 .put(&staging_path, manifest_data.into())
12148 .await
12149 .unwrap();
12150
12151 let mut req = CreateTableVersionRequest::new(5, staging_path.to_string());
12153 req.id = Some(table_id);
12154 req.naming_scheme = Some("V2".to_string());
12155 let err = namespace
12156 .create_table_version(req)
12157 .await
12158 .expect_err("gap create must fail CAS");
12159 let msg = err.to_string();
12160 assert!(
12161 msg.contains("CAS") || msg.contains("ConcurrentModification"),
12162 "expected CAS ConcurrentModification, got: {msg}"
12163 );
12164 }
12165
12166 #[tokio::test]
12167 async fn test_create_table_version_branch_cas_requires_parent_version() {
12168 use futures::TryStreamExt;
12171 use lance_namespace::models::CreateTableVersionRequest;
12172
12173 let (namespace, _temp_dir) = create_test_namespace().await;
12174 create_scalar_table(&namespace, "users").await;
12175 let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
12176 append_scalar_version(&main_uri, 10).await; let mut main = open_dataset(&namespace, "users").await;
12179 let fork_version = main.version().version;
12180 assert_eq!(fork_version, 2);
12181 let branch_uri = main
12182 .create_branch("exp", fork_version, None)
12183 .await
12184 .unwrap()
12185 .uri()
12186 .to_string();
12187
12188 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
12189 let versions_dir = branch_ds.versions_dir();
12190 let store = branch_ds.object_store(None).await.unwrap();
12191 let manifests: Vec<_> = store
12192 .inner
12193 .list(Some(&versions_dir))
12194 .try_collect()
12195 .await
12196 .unwrap();
12197 for meta in &manifests {
12198 if meta
12199 .location
12200 .filename()
12201 .is_some_and(|f| f.ends_with(".manifest"))
12202 {
12203 store.inner.delete(&meta.location).await.unwrap();
12204 }
12205 }
12206 let remaining_manifests = store
12209 .inner
12210 .list(Some(&versions_dir))
12211 .try_collect::<Vec<_>>()
12212 .await
12213 .unwrap()
12214 .into_iter()
12215 .filter(|m| {
12216 m.location
12217 .filename()
12218 .is_some_and(|f| f.ends_with(".manifest"))
12219 })
12220 .count();
12221 assert_eq!(
12222 remaining_manifests, 0,
12223 "branch version chain should be empty after deleting manifests"
12224 );
12225
12226 let main_ds = open_dataset(&namespace, "users").await;
12228 let main_versions = main_ds.versions_dir();
12229 let main_store = main_ds.object_store(None).await.unwrap();
12230 let source_meta = main_store
12231 .inner
12232 .list(Some(&main_versions))
12233 .try_collect::<Vec<_>>()
12234 .await
12235 .unwrap()
12236 .into_iter()
12237 .find(|m| {
12238 m.location
12239 .filename()
12240 .is_some_and(|f| f.ends_with(".manifest"))
12241 })
12242 .expect("main should have a manifest");
12243 let source_bytes = main_store
12244 .inner
12245 .get(&source_meta.location)
12246 .await
12247 .unwrap()
12248 .bytes()
12249 .await
12250 .unwrap();
12251
12252 let staging_wrong = versions_dir.clone().join("staging_wrong");
12253 store
12254 .inner
12255 .put(&staging_wrong, source_bytes.clone().into())
12256 .await
12257 .unwrap();
12258 let err = namespace
12259 .create_table_version(CreateTableVersionRequest {
12260 id: Some(vec!["users".to_string()]),
12261 version: 1,
12262 manifest_path: staging_wrong.to_string(),
12263 naming_scheme: Some("V2".to_string()),
12264 branch: Some("exp".to_string()),
12265 ..Default::default()
12266 })
12267 .await
12268 .expect_err("bootstrap at v1 must fail when parent_version is 2");
12269 let msg = err.to_string();
12270 assert!(
12271 msg.contains("CAS") || msg.contains("ConcurrentModification"),
12272 "expected CAS ConcurrentModification, got: {msg}"
12273 );
12274
12275 let staging_ok = versions_dir.join("staging_ok");
12276 store
12277 .inner
12278 .put(&staging_ok, source_bytes.into())
12279 .await
12280 .unwrap();
12281 let resp = namespace
12282 .create_table_version(CreateTableVersionRequest {
12283 id: Some(vec!["users".to_string()]),
12284 version: 2,
12285 manifest_path: staging_ok.to_string(),
12286 naming_scheme: Some("V2".to_string()),
12287 branch: Some("exp".to_string()),
12288 ..Default::default()
12289 })
12290 .await
12291 .expect("bootstrap at parent_version must succeed");
12292 assert_eq!(resp.version.as_ref().map(|v| v.version), Some(2));
12293 }
12294
12295 #[tokio::test]
12296 async fn test_create_table_version_table_not_found() {
12297 use lance_namespace::models::CreateTableVersionRequest;
12298
12299 let temp_dir = TempStdDir::default();
12300 let temp_path = temp_dir.to_str().unwrap();
12301
12302 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12303 .table_version_tracking_enabled(true)
12304 .build()
12305 .await
12306 .unwrap();
12307
12308 let mut create_version_req =
12310 CreateTableVersionRequest::new(1, "/some/staging/path".to_string());
12311 create_version_req.id = Some(vec!["non_existent_table".to_string()]);
12312
12313 let result = namespace.create_table_version(create_version_req).await;
12314 assert!(
12315 result.is_err(),
12316 "create_table_version should fail for non-existent table"
12317 );
12318 let err_msg = result.unwrap_err().to_string();
12319 assert!(
12320 err_msg.contains("Table not found"),
12321 "Error should mention table not found, got: {}",
12322 err_msg
12323 );
12324 }
12325
12326 mod e2e_table_version_tracking {
12328 use super::*;
12329 use std::sync::atomic::{AtomicUsize, Ordering};
12330
12331 struct TrackingNamespace {
12333 inner: DirectoryNamespace,
12334 create_table_version_count: AtomicUsize,
12335 describe_table_version_count: AtomicUsize,
12336 list_table_versions_count: AtomicUsize,
12337 }
12338
12339 impl TrackingNamespace {
12340 fn new(inner: DirectoryNamespace) -> Self {
12341 Self {
12342 inner,
12343 create_table_version_count: AtomicUsize::new(0),
12344 describe_table_version_count: AtomicUsize::new(0),
12345 list_table_versions_count: AtomicUsize::new(0),
12346 }
12347 }
12348
12349 fn create_table_version_calls(&self) -> usize {
12350 self.create_table_version_count.load(Ordering::SeqCst)
12351 }
12352
12353 fn describe_table_version_calls(&self) -> usize {
12354 self.describe_table_version_count.load(Ordering::SeqCst)
12355 }
12356
12357 fn list_table_versions_calls(&self) -> usize {
12358 self.list_table_versions_count.load(Ordering::SeqCst)
12359 }
12360 }
12361
12362 impl std::fmt::Debug for TrackingNamespace {
12363 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
12364 f.debug_struct("TrackingNamespace")
12365 .field(
12366 "create_table_version_calls",
12367 &self.create_table_version_calls(),
12368 )
12369 .finish()
12370 }
12371 }
12372
12373 #[async_trait]
12374 impl LanceNamespace for TrackingNamespace {
12375 async fn create_namespace(
12376 &self,
12377 request: CreateNamespaceRequest,
12378 ) -> Result<CreateNamespaceResponse> {
12379 self.inner.create_namespace(request).await
12380 }
12381
12382 async fn describe_namespace(
12383 &self,
12384 request: DescribeNamespaceRequest,
12385 ) -> Result<DescribeNamespaceResponse> {
12386 self.inner.describe_namespace(request).await
12387 }
12388
12389 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
12390 self.inner.namespace_exists(request).await
12391 }
12392
12393 async fn list_namespaces(
12394 &self,
12395 request: ListNamespacesRequest,
12396 ) -> Result<ListNamespacesResponse> {
12397 self.inner.list_namespaces(request).await
12398 }
12399
12400 async fn drop_namespace(
12401 &self,
12402 request: DropNamespaceRequest,
12403 ) -> Result<DropNamespaceResponse> {
12404 self.inner.drop_namespace(request).await
12405 }
12406
12407 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
12408 self.inner.list_tables(request).await
12409 }
12410
12411 async fn describe_table(
12412 &self,
12413 request: DescribeTableRequest,
12414 ) -> Result<DescribeTableResponse> {
12415 self.inner.describe_table(request).await
12416 }
12417
12418 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
12419 self.inner.table_exists(request).await
12420 }
12421
12422 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
12423 self.inner.drop_table(request).await
12424 }
12425
12426 async fn create_table(
12427 &self,
12428 request: CreateTableRequest,
12429 request_data: Bytes,
12430 ) -> Result<CreateTableResponse> {
12431 self.inner.create_table(request, request_data).await
12432 }
12433
12434 async fn declare_table(
12435 &self,
12436 request: DeclareTableRequest,
12437 ) -> Result<DeclareTableResponse> {
12438 self.inner.declare_table(request).await
12439 }
12440
12441 async fn list_table_versions(
12442 &self,
12443 request: ListTableVersionsRequest,
12444 ) -> Result<ListTableVersionsResponse> {
12445 self.list_table_versions_count
12446 .fetch_add(1, Ordering::SeqCst);
12447 self.inner.list_table_versions(request).await
12448 }
12449
12450 async fn create_table_version(
12451 &self,
12452 request: CreateTableVersionRequest,
12453 ) -> Result<CreateTableVersionResponse> {
12454 self.create_table_version_count
12455 .fetch_add(1, Ordering::SeqCst);
12456 self.inner.create_table_version(request).await
12457 }
12458
12459 async fn describe_table_version(
12460 &self,
12461 request: DescribeTableVersionRequest,
12462 ) -> Result<DescribeTableVersionResponse> {
12463 self.describe_table_version_count
12464 .fetch_add(1, Ordering::SeqCst);
12465 self.inner.describe_table_version(request).await
12466 }
12467
12468 async fn batch_delete_table_versions(
12469 &self,
12470 request: BatchDeleteTableVersionsRequest,
12471 ) -> Result<BatchDeleteTableVersionsResponse> {
12472 self.inner.batch_delete_table_versions(request).await
12473 }
12474
12475 fn namespace_id(&self) -> String {
12476 self.inner.namespace_id()
12477 }
12478 }
12479
12480 #[tokio::test]
12481 async fn test_describe_table_returns_managed_versioning() {
12482 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableRequest};
12483
12484 let temp_dir = TempStdDir::default();
12485 let temp_path = temp_dir.to_str().unwrap();
12486
12487 let ns = DirectoryNamespaceBuilder::new(temp_path)
12489 .table_version_tracking_enabled(true)
12490 .manifest_enabled(true)
12491 .build()
12492 .await
12493 .unwrap();
12494
12495 let mut create_ns_req = CreateNamespaceRequest::new();
12497 create_ns_req.id = Some(vec!["workspace".to_string()]);
12498 ns.create_namespace(create_ns_req).await.unwrap();
12499
12500 let schema = create_test_schema();
12502 let ipc_data = create_test_ipc_data(&schema);
12503 let mut create_req = CreateTableRequest::new();
12504 create_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
12505 ns.create_table(create_req, bytes::Bytes::from(ipc_data))
12506 .await
12507 .unwrap();
12508
12509 let mut describe_req = DescribeTableRequest::new();
12511 describe_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
12512 let describe_resp = ns.describe_table(describe_req).await.unwrap();
12513
12514 assert_eq!(
12516 describe_resp.managed_versioning,
12517 Some(true),
12518 "managed_versioning should be true when table_version_tracking_enabled=true"
12519 );
12520 }
12521
12522 #[tokio::test]
12523 async fn test_external_manifest_store_invokes_namespace_apis() {
12524 use arrow::array::{Int32Array, StringArray};
12525 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12526 use arrow::record_batch::RecordBatch;
12527 use lance::Dataset;
12528 use lance::dataset::builder::DatasetBuilder;
12529 use lance::dataset::{WriteMode, WriteParams};
12530 use lance_namespace::models::CreateNamespaceRequest;
12531
12532 let temp_dir = TempStdDir::default();
12533 let temp_path = temp_dir.to_str().unwrap();
12534
12535 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
12537 .table_version_tracking_enabled(true)
12538 .manifest_enabled(true)
12539 .build()
12540 .await
12541 .unwrap();
12542
12543 let tracking_ns = Arc::new(TrackingNamespace::new(inner_ns));
12544 let ns: Arc<dyn LanceNamespace> = tracking_ns.clone();
12545
12546 let mut create_ns_req = CreateNamespaceRequest::new();
12548 create_ns_req.id = Some(vec!["workspace".to_string()]);
12549 ns.create_namespace(create_ns_req).await.unwrap();
12550
12551 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12553
12554 let arrow_schema = Arc::new(ArrowSchema::new(vec![
12556 Field::new("id", DataType::Int32, false),
12557 Field::new("name", DataType::Utf8, true),
12558 ]));
12559 let batch = RecordBatch::try_new(
12560 arrow_schema.clone(),
12561 vec![
12562 Arc::new(Int32Array::from(vec![1, 2, 3])),
12563 Arc::new(StringArray::from(vec!["a", "b", "c"])),
12564 ],
12565 )
12566 .unwrap();
12567
12568 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
12570 let write_params = WriteParams {
12571 mode: WriteMode::Create,
12572 ..Default::default()
12573 };
12574 let mut dataset = Dataset::write_into_namespace(
12575 batches,
12576 ns.clone(),
12577 table_id.clone(),
12578 Some(write_params),
12579 )
12580 .await
12581 .unwrap();
12582 assert_eq!(dataset.version().version, 1);
12583
12584 assert_eq!(
12586 tracking_ns.create_table_version_calls(),
12587 1,
12588 "create_table_version should have been called once during initial write_into_namespace"
12589 );
12590
12591 let append_batch = RecordBatch::try_new(
12593 arrow_schema.clone(),
12594 vec![
12595 Arc::new(Int32Array::from(vec![4, 5, 6])),
12596 Arc::new(StringArray::from(vec!["d", "e", "f"])),
12597 ],
12598 )
12599 .unwrap();
12600 let append_batches = RecordBatchIterator::new(vec![Ok(append_batch)], arrow_schema);
12601 dataset.append(append_batches, None).await.unwrap();
12602
12603 assert_eq!(
12604 tracking_ns.create_table_version_calls(),
12605 2,
12606 "create_table_version should have been called twice (once for create, once for append)"
12607 );
12608
12609 let initial_list_calls = tracking_ns.list_table_versions_calls();
12611 let latest_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
12612 .await
12613 .unwrap()
12614 .load()
12615 .await
12616 .unwrap();
12617 assert_eq!(latest_dataset.version().version, 2);
12618 assert_eq!(
12619 tracking_ns.list_table_versions_calls(),
12620 initial_list_calls + 1,
12621 "list_table_versions should have been called exactly once during checkout_latest"
12622 );
12623
12624 let initial_describe_calls = tracking_ns.describe_table_version_calls();
12626 let v1_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
12627 .await
12628 .unwrap()
12629 .with_version(1)
12630 .load()
12631 .await
12632 .unwrap();
12633 assert_eq!(v1_dataset.version().version, 1);
12634 assert_eq!(
12635 tracking_ns.describe_table_version_calls(),
12636 initial_describe_calls + 1,
12637 "describe_table_version should have been called exactly once during checkout to version 1"
12638 );
12639 }
12640
12641 #[tokio::test]
12642 async fn test_dataset_commit_with_external_manifest_store() {
12643 use arrow::array::{Int32Array, StringArray};
12644 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12645 use arrow::record_batch::RecordBatch;
12646 use futures::TryStreamExt;
12647 use lance::dataset::{Dataset, WriteMode, WriteParams};
12648 use lance_namespace::models::CreateNamespaceRequest;
12649 use lance_table::io::commit::ManifestNamingScheme;
12650
12651 let temp_dir = TempStdDir::default();
12652 let temp_path = temp_dir.to_str().unwrap();
12653
12654 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
12656 .table_version_tracking_enabled(true)
12657 .manifest_enabled(true)
12658 .build()
12659 .await
12660 .unwrap();
12661
12662 let tracking_ns: Arc<dyn LanceNamespace> = Arc::new(TrackingNamespace::new(inner_ns));
12663
12664 let mut create_ns_req = CreateNamespaceRequest::new();
12666 create_ns_req.id = Some(vec!["workspace".to_string()]);
12667 tracking_ns.create_namespace(create_ns_req).await.unwrap();
12668
12669 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12671 let arrow_schema = Arc::new(ArrowSchema::new(vec![
12672 Field::new("id", DataType::Int32, false),
12673 Field::new("name", DataType::Utf8, true),
12674 ]));
12675 let batch = RecordBatch::try_new(
12676 arrow_schema.clone(),
12677 vec![
12678 Arc::new(Int32Array::from(vec![1, 2, 3])),
12679 Arc::new(StringArray::from(vec!["a", "b", "c"])),
12680 ],
12681 )
12682 .unwrap();
12683 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12684 let write_params = WriteParams {
12685 mode: WriteMode::Create,
12686 ..Default::default()
12687 };
12688 let dataset = Dataset::write_into_namespace(
12689 batches,
12690 tracking_ns.clone(),
12691 table_id.clone(),
12692 Some(write_params),
12693 )
12694 .await
12695 .unwrap();
12696 assert_eq!(dataset.version().version, 1);
12697
12698 let batch2 = RecordBatch::try_new(
12700 arrow_schema.clone(),
12701 vec![
12702 Arc::new(Int32Array::from(vec![4, 5, 6])),
12703 Arc::new(StringArray::from(vec!["d", "e", "f"])),
12704 ],
12705 )
12706 .unwrap();
12707 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
12708 let write_params = WriteParams {
12709 mode: WriteMode::Append,
12710 ..Default::default()
12711 };
12712 Dataset::write_into_namespace(
12713 batches,
12714 tracking_ns.clone(),
12715 table_id.clone(),
12716 Some(write_params),
12717 )
12718 .await
12719 .unwrap();
12720
12721 let manifest_metas: Vec<_> = dataset
12724 .object_store(None)
12725 .await
12726 .unwrap()
12727 .inner
12728 .list(Some(&dataset.versions_dir()))
12729 .try_collect()
12730 .await
12731 .unwrap();
12732 let version_2_found = manifest_metas.iter().any(|m| {
12733 m.location
12734 .filename()
12735 .map(|f| {
12736 f.ends_with(".manifest")
12737 && ManifestNamingScheme::V2.parse_version(f) == Some(2)
12738 })
12739 .unwrap_or(false)
12740 });
12741 assert!(
12742 version_2_found,
12743 "Version 2 manifest should exist in versions directory"
12744 );
12745 }
12746
12747 async fn create_ns_with_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
12749 use arrow::array::{Int32Array, StringArray};
12750 use arrow::ipc::writer::StreamWriter;
12751
12752 let (namespace, temp_dir) = create_test_namespace().await;
12753
12754 let schema = create_test_schema();
12755 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
12756 let arrow_schema = Arc::new(arrow_schema);
12757
12758 let id_array = Int32Array::from(vec![1, 2, 3]);
12759 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
12760 let batch = arrow::record_batch::RecordBatch::try_new(
12761 arrow_schema.clone(),
12762 vec![Arc::new(id_array), Arc::new(name_array)],
12763 )
12764 .unwrap();
12765
12766 let mut buffer = Vec::new();
12767 {
12768 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
12769 writer.write(&batch).unwrap();
12770 writer.finish().unwrap();
12771 }
12772
12773 let mut request = CreateTableRequest::new();
12774 let table_id = vec!["test_ops_table".to_string()];
12775 request.id = Some(table_id.clone());
12776
12777 namespace
12778 .create_table(request, Bytes::from(buffer))
12779 .await
12780 .unwrap();
12781
12782 (namespace, temp_dir, table_id)
12783 }
12784
12785 #[tokio::test]
12786 async fn test_count_table_rows_basic() {
12787 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12788
12789 let request = CountTableRowsRequest {
12790 id: Some(table_id),
12791 version: None,
12792 predicate: None,
12793 ..Default::default()
12794 };
12795
12796 let count = namespace.count_table_rows(request).await.unwrap();
12797 assert_eq!(count, 3);
12798 }
12799
12800 #[tokio::test]
12801 async fn test_count_table_rows_with_predicate() {
12802 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12803
12804 let request = CountTableRowsRequest {
12805 id: Some(table_id),
12806 version: None,
12807 predicate: Some("id > 1".to_string()),
12808 ..Default::default()
12809 };
12810
12811 let count = namespace.count_table_rows(request).await.unwrap();
12812 assert_eq!(count, 2);
12813 }
12814
12815 #[tokio::test]
12816 async fn test_query_table_invalid_distance_type() {
12817 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
12818
12819 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
12820 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
12821 multi_vector: None,
12822 });
12823
12824 let request = QueryTableRequest {
12825 id: Some(table_id),
12826 k: 2,
12827 vector,
12828 vector_column: Some("vector".to_string()),
12829 distance_type: Some("invalid_metric".to_string()),
12830 filter: None,
12831 offset: None,
12832 version: None,
12833 ..Default::default()
12834 };
12835
12836 let result = namespace.query_table(request).await;
12837 assert!(result.is_err());
12838 let err_msg = result.unwrap_err().to_string();
12839 assert!(
12840 err_msg.contains("Unknown distance type"),
12841 "Expected error about unknown distance type, got: {}",
12842 err_msg
12843 );
12844 }
12845
12846 #[tokio::test]
12847 async fn test_insert_into_table_append() {
12848 use arrow::array::{Int32Array, StringArray};
12849 use arrow::ipc::writer::StreamWriter;
12850
12851 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12852
12853 let schema = create_test_schema();
12855 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
12856 let arrow_schema = Arc::new(arrow_schema);
12857
12858 let id_array = Int32Array::from(vec![4, 5]);
12859 let name_array = StringArray::from(vec!["Dave", "Eve"]);
12860 let batch = arrow::record_batch::RecordBatch::try_new(
12861 arrow_schema.clone(),
12862 vec![Arc::new(id_array), Arc::new(name_array)],
12863 )
12864 .unwrap();
12865
12866 let mut buffer = Vec::new();
12867 {
12868 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
12869 writer.write(&batch).unwrap();
12870 writer.finish().unwrap();
12871 }
12872
12873 let request = InsertIntoTableRequest {
12874 id: Some(table_id.clone()),
12875 mode: Some("append".to_string()),
12876 ..Default::default()
12877 };
12878
12879 let response = namespace
12880 .insert_into_table(request, Bytes::from(buffer))
12881 .await
12882 .unwrap();
12883 assert!(response.transaction_id.is_none());
12884
12885 let count_req = CountTableRowsRequest {
12887 id: Some(table_id),
12888 version: None,
12889 predicate: None,
12890 ..Default::default()
12891 };
12892 let count = namespace.count_table_rows(count_req).await.unwrap();
12893 assert_eq!(count, 5);
12894 }
12895
12896 #[tokio::test]
12897 async fn test_insert_into_table_overwrite() {
12898 use arrow::array::{Int32Array, StringArray};
12899 use arrow::ipc::writer::StreamWriter;
12900
12901 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12902
12903 let schema = create_test_schema();
12904 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
12905 let arrow_schema = Arc::new(arrow_schema);
12906
12907 let id_array = Int32Array::from(vec![10, 20]);
12908 let name_array = StringArray::from(vec!["X", "Y"]);
12909 let batch = arrow::record_batch::RecordBatch::try_new(
12910 arrow_schema.clone(),
12911 vec![Arc::new(id_array), Arc::new(name_array)],
12912 )
12913 .unwrap();
12914
12915 let mut buffer = Vec::new();
12916 {
12917 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
12918 writer.write(&batch).unwrap();
12919 writer.finish().unwrap();
12920 }
12921
12922 let request = InsertIntoTableRequest {
12923 id: Some(table_id.clone()),
12924 mode: Some("overwrite".to_string()),
12925 ..Default::default()
12926 };
12927
12928 namespace
12929 .insert_into_table(request, Bytes::from(buffer))
12930 .await
12931 .unwrap();
12932
12933 let count_req = CountTableRowsRequest {
12935 id: Some(table_id),
12936 version: None,
12937 predicate: None,
12938 ..Default::default()
12939 };
12940 let count = namespace.count_table_rows(count_req).await.unwrap();
12941 assert_eq!(count, 2);
12942 }
12943
12944 #[tokio::test]
12945 async fn test_insert_into_table_empty_data() {
12946 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
12947
12948 let request = InsertIntoTableRequest {
12949 id: Some(table_id),
12950 mode: None,
12951 ..Default::default()
12952 };
12953
12954 let result = namespace.insert_into_table(request, Bytes::new()).await;
12955 assert!(result.is_err());
12956 assert!(
12957 result
12958 .unwrap_err()
12959 .to_string()
12960 .contains("Arrow IPC stream) is required")
12961 );
12962 }
12963
12964 #[tokio::test]
12965 async fn test_insert_into_table_with_storage_options() {
12966 use arrow::array::{Int32Array, StringArray};
12967 use arrow::ipc::writer::StreamWriter;
12968
12969 let temp_dir = TempStdDir::default();
12970
12971 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
12973 .storage_option("allow_http", "true")
12974 .build()
12975 .await
12976 .unwrap();
12977
12978 let schema = create_test_schema();
12980 let ipc_data = create_test_ipc_data(&schema);
12981 let mut create_req = CreateTableRequest::new();
12982 let table_id = vec!["so_table".to_string()];
12983 create_req.id = Some(table_id.clone());
12984 namespace
12985 .create_table(create_req, Bytes::from(ipc_data))
12986 .await
12987 .unwrap();
12988
12989 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
12991 let arrow_schema = Arc::new(arrow_schema);
12992
12993 let id_array = Int32Array::from(vec![10, 20]);
12994 let name_array = StringArray::from(vec!["X", "Y"]);
12995 let batch = arrow::record_batch::RecordBatch::try_new(
12996 arrow_schema.clone(),
12997 vec![Arc::new(id_array), Arc::new(name_array)],
12998 )
12999 .unwrap();
13000
13001 let mut buffer = Vec::new();
13002 {
13003 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13004 writer.write(&batch).unwrap();
13005 writer.finish().unwrap();
13006 }
13007
13008 let request = InsertIntoTableRequest {
13009 id: Some(table_id.clone()),
13010 mode: Some("append".to_string()),
13011 ..Default::default()
13012 };
13013
13014 let response = namespace
13015 .insert_into_table(request, Bytes::from(buffer))
13016 .await
13017 .unwrap();
13018 assert!(response.transaction_id.is_none());
13019
13020 let count_req = CountTableRowsRequest {
13022 id: Some(table_id),
13023 version: None,
13024 predicate: None,
13025 ..Default::default()
13026 };
13027 let count = namespace.count_table_rows(count_req).await.unwrap();
13028 assert_eq!(count, 2);
13029 }
13030
13031 #[tokio::test]
13032 async fn test_query_table_basic() {
13033 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13034
13035 let request = QueryTableRequest {
13036 id: Some(table_id),
13037 k: 10,
13038 filter: None,
13039 offset: None,
13040 version: None,
13041 ..Default::default()
13042 };
13043
13044 let bytes = namespace.query_table(request).await.unwrap();
13045
13046 let cursor = Cursor::new(bytes.to_vec());
13048 let reader = FileReader::try_new(cursor, None).unwrap();
13049 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13050 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13051 assert_eq!(total_rows, 3);
13052 }
13053
13054 #[tokio::test]
13055 async fn test_query_table_with_filter() {
13056 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13057
13058 let request = QueryTableRequest {
13059 id: Some(table_id),
13060 k: 10,
13061 filter: Some("id <= 2".to_string()),
13062 offset: None,
13063 version: None,
13064 ..Default::default()
13065 };
13066
13067 let bytes = namespace.query_table(request).await.unwrap();
13068
13069 let cursor = Cursor::new(bytes.to_vec());
13070 let reader = FileReader::try_new(cursor, None).unwrap();
13071 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13072 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13073 assert_eq!(total_rows, 2);
13074 }
13075
13076 #[tokio::test]
13077 async fn test_query_table_with_limit_and_offset() {
13078 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13079
13080 let request = QueryTableRequest {
13081 id: Some(table_id),
13082 k: 2,
13083 filter: None,
13084 offset: Some(1),
13085 version: None,
13086 ..Default::default()
13087 };
13088
13089 let bytes = namespace.query_table(request).await.unwrap();
13090
13091 let cursor = Cursor::new(bytes.to_vec());
13092 let reader = FileReader::try_new(cursor, None).unwrap();
13093 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13094 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13095 assert_eq!(total_rows, 2);
13096 }
13097
13098 #[tokio::test]
13099 async fn test_query_table_no_limit() {
13100 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13101
13102 let request = QueryTableRequest {
13104 id: Some(table_id),
13105 k: 0,
13106 filter: None,
13107 offset: None,
13108 version: None,
13109 ..Default::default()
13110 };
13111
13112 let bytes = namespace.query_table(request).await.unwrap();
13113
13114 let cursor = Cursor::new(bytes.to_vec());
13115 let reader = FileReader::try_new(cursor, None).unwrap();
13116 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13117 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13118 assert_eq!(total_rows, 3);
13119 }
13120
13121 #[tokio::test]
13122 async fn test_query_table_with_columns() {
13123 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13124
13125 let columns = Box::new(lance_namespace::models::QueryTableRequestColumns {
13126 column_names: Some(vec!["id".to_string()]),
13127 column_aliases: None,
13128 });
13129
13130 let request = QueryTableRequest {
13131 id: Some(table_id),
13132 k: 10,
13133 filter: None,
13134 offset: None,
13135 version: None,
13136 columns: Some(columns),
13137 ..Default::default()
13138 };
13139
13140 let bytes = namespace.query_table(request).await.unwrap();
13141
13142 let cursor = Cursor::new(bytes.to_vec());
13143 let reader = FileReader::try_new(cursor, None).unwrap();
13144 let schema = reader.schema();
13145 assert_eq!(schema.fields().len(), 1);
13146 assert_eq!(schema.field(0).name(), "id");
13147 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13148 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13149 assert_eq!(total_rows, 3);
13150 }
13151
13152 #[tokio::test]
13153 async fn test_count_table_rows_with_version() {
13154 use arrow::array::{Int32Array, StringArray};
13155 use arrow::ipc::writer::StreamWriter;
13156
13157 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13158
13159 let schema = create_test_schema();
13161 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13162 let arrow_schema = Arc::new(arrow_schema);
13163
13164 let id_array = Int32Array::from(vec![4, 5]);
13165 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13166 let batch = arrow::record_batch::RecordBatch::try_new(
13167 arrow_schema.clone(),
13168 vec![Arc::new(id_array), Arc::new(name_array)],
13169 )
13170 .unwrap();
13171
13172 let mut buffer = Vec::new();
13173 {
13174 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13175 writer.write(&batch).unwrap();
13176 writer.finish().unwrap();
13177 }
13178
13179 let request = InsertIntoTableRequest {
13180 id: Some(table_id.clone()),
13181 mode: None,
13182 ..Default::default()
13183 };
13184 namespace
13185 .insert_into_table(request, Bytes::from(buffer))
13186 .await
13187 .unwrap();
13188
13189 let count_req = CountTableRowsRequest {
13191 id: Some(table_id.clone()),
13192 version: Some(1),
13193 predicate: None,
13194 ..Default::default()
13195 };
13196 let count = namespace.count_table_rows(count_req).await.unwrap();
13197 assert_eq!(count, 3);
13198
13199 let count_req = CountTableRowsRequest {
13201 id: Some(table_id),
13202 version: None,
13203 predicate: None,
13204 ..Default::default()
13205 };
13206 let count = namespace.count_table_rows(count_req).await.unwrap();
13207 assert_eq!(count, 5);
13208 }
13209
13210 #[tokio::test]
13211 async fn test_query_table_with_version() {
13212 use arrow::array::{Int32Array, StringArray};
13213 use arrow::ipc::writer::StreamWriter;
13214
13215 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13216
13217 let schema = create_test_schema();
13219 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13220 let arrow_schema = Arc::new(arrow_schema);
13221
13222 let id_array = Int32Array::from(vec![4, 5]);
13223 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13224 let batch = arrow::record_batch::RecordBatch::try_new(
13225 arrow_schema.clone(),
13226 vec![Arc::new(id_array), Arc::new(name_array)],
13227 )
13228 .unwrap();
13229
13230 let mut buffer = Vec::new();
13231 {
13232 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13233 writer.write(&batch).unwrap();
13234 writer.finish().unwrap();
13235 }
13236
13237 let request = InsertIntoTableRequest {
13238 id: Some(table_id.clone()),
13239 mode: None,
13240 ..Default::default()
13241 };
13242 namespace
13243 .insert_into_table(request, Bytes::from(buffer))
13244 .await
13245 .unwrap();
13246
13247 let request = QueryTableRequest {
13249 id: Some(table_id.clone()),
13250 k: 100,
13251 filter: None,
13252 offset: None,
13253 version: Some(1),
13254 ..Default::default()
13255 };
13256
13257 let bytes = namespace.query_table(request).await.unwrap();
13258 let cursor = Cursor::new(bytes.to_vec());
13259 let reader = FileReader::try_new(cursor, None).unwrap();
13260 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13261 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13262 assert_eq!(total_rows, 3);
13263
13264 let request = QueryTableRequest {
13266 id: Some(table_id),
13267 k: 100,
13268 filter: None,
13269 offset: None,
13270 version: None,
13271 ..Default::default()
13272 };
13273
13274 let bytes = namespace.query_table(request).await.unwrap();
13275 let cursor = Cursor::new(bytes.to_vec());
13276 let reader = FileReader::try_new(cursor, None).unwrap();
13277 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13278 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13279 assert_eq!(total_rows, 5);
13280 }
13281
13282 async fn create_ns_with_vector_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
13285 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
13286 use arrow::ipc::writer::StreamWriter;
13287
13288 let (namespace, temp_dir) = create_test_namespace().await;
13289
13290 let arrow_schema = Arc::new(arrow::datatypes::Schema::new(vec![
13292 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
13293 arrow::datatypes::Field::new(
13294 "vector",
13295 arrow::datatypes::DataType::FixedSizeList(
13296 Arc::new(arrow::datatypes::Field::new(
13297 "item",
13298 arrow::datatypes::DataType::Float32,
13299 true,
13300 )),
13301 4,
13302 ),
13303 true,
13304 ),
13305 ]));
13306
13307 let id_array = Int32Array::from(vec![1, 2, 3]);
13308 let values = Float32Array::from(vec![
13309 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ]);
13313 let vector_array = FixedSizeListArray::try_new(
13314 Arc::new(arrow::datatypes::Field::new(
13315 "item",
13316 arrow::datatypes::DataType::Float32,
13317 true,
13318 )),
13319 4,
13320 Arc::new(values),
13321 None,
13322 )
13323 .unwrap();
13324
13325 let batch = arrow::record_batch::RecordBatch::try_new(
13326 arrow_schema.clone(),
13327 vec![Arc::new(id_array), Arc::new(vector_array)],
13328 )
13329 .unwrap();
13330
13331 let mut buffer = Vec::new();
13332 {
13333 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13334 writer.write(&batch).unwrap();
13335 writer.finish().unwrap();
13336 }
13337
13338 let table_name = "vector_table";
13340 let table_uri = format!("{}/{}.lance", temp_dir.to_str().unwrap(), table_name);
13341 let reader = arrow::record_batch::RecordBatchIterator::new(
13342 vec![Ok(batch)],
13343 arrow_schema.clone(),
13344 );
13345 Dataset::write(reader, &table_uri, None).await.unwrap();
13346
13347 let table_id = vec![table_name.to_string()];
13348 (namespace, temp_dir, table_id)
13349 }
13350
13351 #[tokio::test]
13352 async fn test_query_table_vector_search() {
13353 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13354
13355 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13356 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13357 multi_vector: None,
13358 });
13359
13360 let request = QueryTableRequest {
13361 id: Some(table_id),
13362 k: 2,
13363 vector,
13364 filter: None,
13365 offset: None,
13366 version: None,
13367 ..Default::default()
13368 };
13369
13370 let bytes = namespace.query_table(request).await.unwrap();
13371
13372 let cursor = Cursor::new(bytes.to_vec());
13373 let reader = FileReader::try_new(cursor, None).unwrap();
13374 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13375 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13376 assert_eq!(total_rows, 2);
13377 }
13378
13379 #[tokio::test]
13380 async fn test_query_table_vector_search_with_distance_type() {
13381 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13382
13383 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13384 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13385 multi_vector: None,
13386 });
13387
13388 let request = QueryTableRequest {
13389 id: Some(table_id),
13390 k: 3,
13391 vector,
13392 filter: None,
13393 offset: None,
13394 version: None,
13395 distance_type: Some("cosine".to_string()),
13396 ..Default::default()
13397 };
13398
13399 let bytes = namespace.query_table(request).await.unwrap();
13400
13401 let cursor = Cursor::new(bytes.to_vec());
13402 let reader = FileReader::try_new(cursor, None).unwrap();
13403 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13404 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13405 assert_eq!(total_rows, 3);
13406 }
13407
13408 #[tokio::test]
13409 async fn test_query_table_vector_search_with_filter() {
13410 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13411
13412 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13413 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13414 multi_vector: None,
13415 });
13416
13417 let request = QueryTableRequest {
13418 id: Some(table_id),
13419 k: 10,
13420 vector,
13421 filter: Some("id <= 2".to_string()),
13422 offset: None,
13423 version: None,
13424 ..Default::default()
13425 };
13426
13427 let bytes = namespace.query_table(request).await.unwrap();
13428
13429 let cursor = Cursor::new(bytes.to_vec());
13430 let reader = FileReader::try_new(cursor, None).unwrap();
13431 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13432 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13433 assert!(total_rows <= 2);
13434 }
13435
13436 #[tokio::test]
13437 async fn test_query_table_vector_search_with_nprobes_and_refine() {
13438 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13439
13440 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13441 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
13442 multi_vector: None,
13443 });
13444
13445 let request = QueryTableRequest {
13446 id: Some(table_id),
13447 k: 2,
13448 vector,
13449 filter: None,
13450 offset: None,
13451 version: None,
13452 nprobes: Some(1),
13453 refine_factor: Some(1),
13454 prefilter: Some(true),
13455 ..Default::default()
13456 };
13457
13458 let bytes = namespace.query_table(request).await.unwrap();
13459
13460 let cursor = Cursor::new(bytes.to_vec());
13461 let reader = FileReader::try_new(cursor, None).unwrap();
13462 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13463 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13464 assert_eq!(total_rows, 2);
13465 }
13466
13467 #[tokio::test]
13468 async fn test_namespace_id() {
13469 let (namespace, _temp_dir) = create_test_namespace().await;
13470 let id = namespace.namespace_id();
13471 assert!(id.contains("DirectoryNamespace"));
13472 assert!(id.contains("root"));
13473 }
13474
13475 #[tokio::test]
13476 async fn test_query_table_empty_table() {
13477 let (namespace, _temp_dir) = create_test_namespace().await;
13478
13479 let schema = create_test_schema();
13481 let ipc_data = create_test_ipc_data(&schema);
13482 let mut create_request = CreateTableRequest::new();
13483 create_request.id = Some(vec!["empty_table".to_string()]);
13484 namespace
13485 .create_table(create_request, bytes::Bytes::from(ipc_data))
13486 .await
13487 .unwrap();
13488
13489 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13491 single_vector: None,
13492 multi_vector: None,
13493 });
13494 let request = QueryTableRequest {
13495 id: Some(vec!["empty_table".to_string()]),
13496 k: 10,
13497 vector,
13498 ..Default::default()
13499 };
13500 let bytes = namespace.query_table(request).await.unwrap();
13501
13502 let cursor = Cursor::new(bytes.to_vec());
13503 let reader = FileReader::try_new(cursor, None).unwrap();
13504 let batches: Vec<_> = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
13505 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13506 assert_eq!(total_rows, 0, "empty table should yield no rows");
13507 }
13508
13509 #[tokio::test]
13510 async fn test_query_table_with_plain_filter_no_vector() {
13511 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13512
13513 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13515 single_vector: None,
13516 multi_vector: None,
13517 });
13518 let request = QueryTableRequest {
13519 id: Some(table_id),
13520 k: 0,
13521 vector,
13522 filter: Some("id > 1".to_string()),
13523 ..Default::default()
13524 };
13525 let bytes = namespace.query_table(request).await.unwrap();
13526
13527 let cursor = Cursor::new(bytes.to_vec());
13528 let reader = FileReader::try_new(cursor, None).unwrap();
13529 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13530 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13531 assert!(total_rows > 0);
13532 assert!(total_rows < 3);
13533 }
13534
13535 #[tokio::test]
13538 async fn test_update_full_table() {
13539 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13540
13541 let base_version = open_dataset(&namespace, &table_id[0])
13543 .await
13544 .version()
13545 .version;
13546
13547 let request = UpdateTableRequest {
13548 id: Some(table_id.clone()),
13549 updates: vec![vec!["name".to_string(), "'updated'".to_string()]],
13550 predicate: None,
13551 ..Default::default()
13552 };
13553
13554 let response = namespace.update_table(request).await.unwrap();
13555 assert_eq!(response.updated_rows, 3);
13556 assert!(response.version as u64 > base_version);
13557
13558 let count_req = CountTableRowsRequest {
13560 id: Some(table_id),
13561 version: None,
13562 predicate: Some("name = 'updated'".to_string()),
13563 ..Default::default()
13564 };
13565 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 3);
13566 }
13567
13568 #[tokio::test]
13569 async fn test_update_with_predicate() {
13570 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13571
13572 let request = UpdateTableRequest {
13573 id: Some(table_id.clone()),
13574 updates: vec![vec!["name".to_string(), "'matched'".to_string()]],
13575 predicate: Some("id > 1".to_string()),
13576 ..Default::default()
13577 };
13578
13579 let response = namespace.update_table(request).await.unwrap();
13580 assert_eq!(response.updated_rows, 2);
13581
13582 let untouched = CountTableRowsRequest {
13584 id: Some(table_id.clone()),
13585 version: None,
13586 predicate: Some("name = 'Alice'".to_string()),
13587 ..Default::default()
13588 };
13589 assert_eq!(namespace.count_table_rows(untouched).await.unwrap(), 1);
13590
13591 let touched = CountTableRowsRequest {
13592 id: Some(table_id),
13593 version: None,
13594 predicate: Some("name = 'matched'".to_string()),
13595 ..Default::default()
13596 };
13597 assert_eq!(namespace.count_table_rows(touched).await.unwrap(), 2);
13598 }
13599
13600 #[tokio::test]
13601 async fn test_update_invalid_expression_returns_invalid_input() {
13602 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13603
13604 let request = UpdateTableRequest {
13605 id: Some(table_id),
13606 updates: vec![vec!["name".to_string(), "no_such_column + 1".to_string()]],
13608 predicate: None,
13609 ..Default::default()
13610 };
13611
13612 let err = namespace.update_table(request).await.unwrap_err();
13613 let msg = err.to_string();
13614 assert!(
13615 msg.contains("Invalid input"),
13616 "expected Invalid input error, got: {}",
13617 msg
13618 );
13619 }
13620
13621 #[tokio::test]
13622 async fn test_update_rejects_duplicate_columns() {
13623 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13624
13625 let request = UpdateTableRequest {
13626 id: Some(table_id),
13627 updates: vec![
13628 vec!["name".to_string(), "'a'".to_string()],
13629 vec!["name".to_string(), "'b'".to_string()],
13630 ],
13631 predicate: None,
13632 ..Default::default()
13633 };
13634
13635 let err = namespace.update_table(request).await.unwrap_err();
13636 let msg = err.to_string();
13637 assert!(
13638 msg.contains("Invalid input") && msg.contains("more than once"),
13639 "expected duplicate column InvalidInput error, got: {}",
13640 msg
13641 );
13642 }
13643
13644 #[tokio::test]
13645 async fn test_delete_with_predicate() {
13646 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13647
13648 let request = DeleteFromTableRequest {
13649 id: Some(table_id.clone()),
13650 predicate: "id > 1".to_string(),
13651 ..Default::default()
13652 };
13653
13654 let response = namespace.delete_from_table(request).await.unwrap();
13655 assert!(response.version.is_some());
13656
13657 let count_req = CountTableRowsRequest {
13658 id: Some(table_id),
13659 version: None,
13660 predicate: None,
13661 ..Default::default()
13662 };
13663 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 1);
13665 }
13666
13667 #[tokio::test]
13668 async fn test_delete_empty_predicate_returns_invalid_input() {
13669 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13670
13671 let request = DeleteFromTableRequest {
13672 id: Some(table_id),
13673 predicate: " ".to_string(),
13674 ..Default::default()
13675 };
13676
13677 let err = namespace.delete_from_table(request).await.unwrap_err();
13678 let msg = err.to_string();
13679 assert!(
13680 msg.contains("Invalid input") && msg.contains("non-empty predicate"),
13681 "expected non-empty predicate InvalidInput error, got: {}",
13682 msg
13683 );
13684 }
13685
13686 #[tokio::test]
13687 async fn test_delete_table_not_found() {
13688 let (namespace, _temp_dir) = create_test_namespace().await;
13689
13690 let request = DeleteFromTableRequest {
13691 id: Some(vec!["does_not_exist".to_string()]),
13692 predicate: "id = 1".to_string(),
13693 ..Default::default()
13694 };
13695
13696 let err = namespace.delete_from_table(request).await.unwrap_err();
13697 let msg = err.to_string();
13698 assert!(
13699 msg.contains("Table not found"),
13700 "expected TableNotFound for missing table, got: {}",
13701 msg
13702 );
13703 }
13704
13705 #[tokio::test]
13706 async fn test_delete_invalid_predicate_returns_invalid_input() {
13707 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13708
13709 let request = DeleteFromTableRequest {
13713 id: Some(table_id),
13714 predicate: "no_such_column = 1".to_string(),
13715 ..Default::default()
13716 };
13717
13718 let err = namespace.delete_from_table(request).await.unwrap_err();
13719 let lance_core::Error::Namespace { source, .. } = &err else {
13720 panic!("expected a Namespace error, got: {}", err);
13721 };
13722 let ns_err = source
13723 .downcast_ref::<NamespaceError>()
13724 .expect("expected a NamespaceError source");
13725 assert_eq!(
13726 ns_err.code(),
13727 lance_namespace::ErrorCode::InvalidInput,
13728 "expected InvalidInput for an invalid delete predicate, got: {}",
13729 err
13730 );
13731 }
13732 }
13733
13734 #[tokio::test]
13735 async fn test_list_all_tables() {
13736 use lance_namespace::models::ListTablesRequest;
13737
13738 let (namespace, _temp_dir) = create_test_namespace().await;
13739 create_scalar_table(&namespace, "alpha").await;
13740 create_scalar_table(&namespace, "beta").await;
13741
13742 let request = ListTablesRequest {
13743 id: Some(vec![]),
13744 page_token: None,
13745 limit: None,
13746 ..Default::default()
13747 };
13748 let response = namespace.list_all_tables(request).await.unwrap();
13749 let mut tables = response.tables;
13750 tables.sort();
13751 assert_eq!(tables, vec!["alpha", "beta"]);
13752 }
13753
13754 #[tokio::test]
13755 async fn test_restore_table() {
13756 use lance_namespace::models::RestoreTableRequest;
13757
13758 let (namespace, _temp_dir) = create_test_namespace().await;
13759 create_scalar_table(&namespace, "users").await;
13760
13761 create_scalar_index(&namespace, "users", "users_id_idx").await;
13763
13764 let dataset = open_dataset(&namespace, "users").await;
13765 let current_version = dataset.version().version;
13766 assert!(current_version >= 2, "Should have at least 2 versions");
13767
13768 let mut restore_req = RestoreTableRequest::new(1);
13770 restore_req.id = Some(vec!["users".to_string()]);
13771 let response = namespace.restore_table(restore_req).await.unwrap();
13772
13773 assert!(
13775 response.transaction_id.is_some(),
13776 "restore_table should return a transaction_id"
13777 );
13778
13779 let dataset_after = open_dataset(&namespace, "users").await;
13781 assert!(
13782 dataset_after.version().version > current_version,
13783 "Restore should create a new version"
13784 );
13785 }
13786
13787 #[tokio::test]
13788 async fn test_alter_table_add_columns() {
13789 use lance_namespace::models::{
13790 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
13791 };
13792
13793 let (namespace, _temp_dir) = create_test_namespace().await;
13794
13795 let schema = create_test_schema();
13797 let ipc_data = create_test_ipc_data(&schema);
13798 let mut create_request = CreateTableRequest::new();
13799 create_request.id = Some(vec!["test_table".to_string()]);
13800 namespace
13801 .create_table(create_request, bytes::Bytes::from(ipc_data))
13802 .await
13803 .unwrap();
13804
13805 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
13807 new_col.expression = Some(Some("id * 2".to_string()));
13808 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
13809 add_request.id = Some(vec!["test_table".to_string()]);
13810
13811 let response = namespace
13812 .alter_table_add_columns(add_request)
13813 .await
13814 .unwrap();
13815 assert!(
13816 response.version > 1,
13817 "Version should increment after adding columns"
13818 );
13819
13820 let mut describe_request = DescribeTableRequest::new();
13822 describe_request.id = Some(vec!["test_table".to_string()]);
13823 describe_request.load_detailed_metadata = Some(true);
13824 let describe_response = namespace.describe_table(describe_request).await.unwrap();
13825 assert!(describe_response.schema.is_some());
13826
13827 let resp_schema = describe_response.schema.unwrap();
13828 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
13829 assert!(
13830 field_names.contains(&"doubled_id"),
13831 "Column 'doubled_id' should exist, got: {:?}",
13832 field_names
13833 );
13834 }
13835
13836 #[tokio::test]
13837 async fn test_update_table_schema_metadata() {
13838 use lance_namespace::models::UpdateTableSchemaMetadataRequest;
13839
13840 let (namespace, _temp_dir) = create_test_namespace().await;
13841 create_scalar_table(&namespace, "products").await;
13842
13843 let mut metadata = HashMap::new();
13844 metadata.insert("owner".to_string(), "team_a".to_string());
13845 metadata.insert("version".to_string(), "1.0".to_string());
13846
13847 let mut req = UpdateTableSchemaMetadataRequest::new();
13848 req.id = Some(vec!["products".to_string()]);
13849 req.metadata = Some(metadata.clone());
13850
13851 let response = namespace.update_table_schema_metadata(req).await.unwrap();
13852
13853 assert!(response.metadata.is_some());
13854 let returned = response.metadata.unwrap();
13855 assert_eq!(returned.get("owner"), Some(&"team_a".to_string()));
13856 assert_eq!(returned.get("version"), Some(&"1.0".to_string()));
13857 assert!(
13858 response.transaction_id.is_some(),
13859 "update_table_schema_metadata should return a transaction_id"
13860 );
13861 }
13862
13863 #[tokio::test]
13864 async fn test_alter_table_add_columns_missing_id() {
13865 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
13866
13867 let (namespace, _temp_dir) = create_test_namespace().await;
13868
13869 let new_col = AddColumnsEntry::new("col".to_string());
13870 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
13871 let result = namespace.alter_table_add_columns(request).await;
13872 assert!(result.is_err(), "Should fail when table ID is missing");
13873 }
13874
13875 #[tokio::test]
13876 async fn test_alter_table_alter_columns_rename() {
13877 use lance_namespace::models::{
13878 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
13879 };
13880
13881 let (namespace, _temp_dir) = create_test_namespace().await;
13882
13883 let schema = create_test_schema();
13885 let ipc_data = create_test_ipc_data(&schema);
13886 let mut create_request = CreateTableRequest::new();
13887 create_request.id = Some(vec!["test_table".to_string()]);
13888 namespace
13889 .create_table(create_request, bytes::Bytes::from(ipc_data))
13890 .await
13891 .unwrap();
13892
13893 let mut entry = AlterColumnsEntry::new("name".to_string());
13895 entry.rename = Some(Some("full_name".to_string()));
13896 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
13897 alter_request.id = Some(vec!["test_table".to_string()]);
13898
13899 let response = namespace
13900 .alter_table_alter_columns(alter_request)
13901 .await
13902 .unwrap();
13903 assert!(
13904 response.version > 1,
13905 "Version should increment after altering columns"
13906 );
13907
13908 let mut describe_request = DescribeTableRequest::new();
13910 describe_request.id = Some(vec!["test_table".to_string()]);
13911 describe_request.load_detailed_metadata = Some(true);
13912 let describe_response = namespace.describe_table(describe_request).await.unwrap();
13913 assert!(describe_response.schema.is_some());
13914
13915 let resp_schema = describe_response.schema.unwrap();
13916 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
13917 assert!(
13918 field_names.contains(&"full_name"),
13919 "Column should be renamed to 'full_name', got: {:?}",
13920 field_names
13921 );
13922 assert!(
13923 !field_names.contains(&"name"),
13924 "Old column 'name' should not exist, got: {:?}",
13925 field_names
13926 );
13927 }
13928
13929 #[tokio::test]
13930 async fn test_get_table_stats() {
13931 use lance_namespace::models::GetTableStatsRequest;
13932
13933 let (namespace, _temp_dir) = create_test_namespace().await;
13934 create_scalar_table(&namespace, "items").await;
13935 create_scalar_index(&namespace, "items", "items_id_idx").await;
13936
13937 let mut req = GetTableStatsRequest::new();
13938 req.id = Some(vec!["items".to_string()]);
13939
13940 let response = namespace.get_table_stats(req).await.unwrap();
13941 assert_eq!(response.num_rows, 3);
13942 assert_eq!(response.num_indices, 1);
13943 }
13944
13945 #[tokio::test]
13946 async fn test_explain_table_query_plan() {
13947 use lance_namespace::models::QueryTableRequestVector;
13948 use lance_namespace::models::{ExplainTableQueryPlanRequest, QueryTableRequest};
13949
13950 let (namespace, _temp_dir) = create_test_namespace().await;
13951 create_scalar_table(&namespace, "catalog").await;
13952
13953 let mut query = QueryTableRequest::new(1, QueryTableRequestVector::new());
13954 query.filter = Some("id > 1".to_string());
13955 query.columns = Some(Box::new(QueryTableRequestColumns {
13956 column_names: Some(vec!["id".to_string(), "name".to_string()]),
13957 column_aliases: None,
13958 }));
13959 query.with_row_id = Some(true);
13960
13961 let mut req = ExplainTableQueryPlanRequest::new(query);
13962 req.id = Some(vec!["catalog".to_string()]);
13963
13964 let plan_str = namespace.explain_table_query_plan(req).await.unwrap();
13965 assert_plan_contains_all(
13966 &plan_str,
13967 &[
13968 "ProjectionExec: expr=[id@0 as id, name@2 as name",
13969 "projection=[name], source=stream(_rowid)",
13970 "LanceRead: uri=",
13971 "projection=[id]",
13972 "row_id=true, row_addr=false",
13973 "full_filter=id > Int32(1)",
13974 "refine_filter=id > Int32(1)",
13975 ],
13976 "Filtered explain plan should preserve late materialization and filter pushdown",
13977 );
13978 }
13979
13980 #[tokio::test]
13981 async fn test_alter_table_alter_columns_missing_id() {
13982 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
13983
13984 let (namespace, _temp_dir) = create_test_namespace().await;
13985
13986 let entry = AlterColumnsEntry::new("name".to_string());
13987 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
13988 let result = namespace.alter_table_alter_columns(request).await;
13989 assert!(result.is_err(), "Should fail when table ID is missing");
13990 }
13991
13992 #[tokio::test]
13993 async fn test_alter_table_drop_columns() {
13994 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
13995
13996 let (namespace, _temp_dir) = create_test_namespace().await;
13997
13998 let schema = create_test_schema();
14000 let ipc_data = create_test_ipc_data(&schema);
14001 let mut create_request = CreateTableRequest::new();
14002 create_request.id = Some(vec!["test_table".to_string()]);
14003 namespace
14004 .create_table(create_request, bytes::Bytes::from(ipc_data))
14005 .await
14006 .unwrap();
14007
14008 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
14010 drop_request.id = Some(vec!["test_table".to_string()]);
14011
14012 let response = namespace
14013 .alter_table_drop_columns(drop_request)
14014 .await
14015 .unwrap();
14016 assert!(
14017 response.version > 1,
14018 "Version should increment after dropping columns"
14019 );
14020
14021 let mut describe_request = DescribeTableRequest::new();
14023 describe_request.id = Some(vec!["test_table".to_string()]);
14024 describe_request.load_detailed_metadata = Some(true);
14025 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14026 assert!(describe_response.schema.is_some());
14027
14028 let resp_schema = describe_response.schema.unwrap();
14029 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14030 assert!(
14031 !field_names.contains(&"name"),
14032 "Column 'name' should be dropped, got: {:?}",
14033 field_names
14034 );
14035 assert!(
14036 field_names.contains(&"id"),
14037 "Column 'id' should still exist, got: {:?}",
14038 field_names
14039 );
14040 }
14041
14042 #[tokio::test]
14043 async fn test_analyze_table_query_plan() {
14044 use lance_namespace::models::AnalyzeTableQueryPlanRequest;
14045 use lance_namespace::models::QueryTableRequestVector;
14046
14047 let (namespace, _temp_dir) = create_test_namespace().await;
14048 create_scalar_table(&namespace, "catalog").await;
14049
14050 let mut req = AnalyzeTableQueryPlanRequest::new(1, QueryTableRequestVector::new());
14051 req.id = Some(vec!["catalog".to_string()]);
14052 req.filter = Some("id > 0".to_string());
14053 req.columns = Some(Box::new(QueryTableRequestColumns {
14054 column_names: Some(vec!["id".to_string(), "name".to_string()]),
14055 column_aliases: None,
14056 }));
14057 req.with_row_id = Some(true);
14058
14059 let analysis_str = namespace.analyze_table_query_plan(req).await.unwrap();
14060 assert_plan_contains_all(
14061 &analysis_str,
14062 &[
14063 "AnalyzeExec verbose=true",
14064 "ProjectionExec: elapsed=",
14065 "expr=[id@0 as id, name@2 as name",
14066 "projection=[name], source=stream(_rowid)",
14067 "LanceRead: elapsed=",
14068 "projection=[id]",
14069 "row_id=true, row_addr=false",
14070 "full_filter=id > Int32(0)",
14071 "refine_filter=id > Int32(0)",
14072 "metrics=[output_rows=",
14073 ],
14074 "Filtered analyze plan should preserve late materialization and filter pushdown",
14075 );
14076 }
14077
14078 #[tokio::test]
14079 async fn test_dir_listing_no_extra_calls_without_migration() {
14080 let temp_dir = TempStdDir::default();
14081 let temp_path = temp_dir.to_str().unwrap();
14082 let root_uri = file_object_store_uri(temp_path);
14083 let listing_count = Arc::new(AtomicUsize::new(0));
14084 let session = build_listing_counting_session(listing_count.clone());
14085
14086 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
14088 .session(session.clone())
14089 .manifest_enabled(false)
14090 .dir_listing_enabled(true)
14091 .build()
14092 .await
14093 .unwrap();
14094
14095 let schema = create_test_schema();
14096 let ipc_data = create_test_ipc_data(&schema);
14097 let mut create_req = CreateTableRequest::new();
14098 create_req.id = Some(vec!["test_table".to_string()]);
14099 dir_only_ns
14100 .create_table(create_req, Bytes::from(ipc_data))
14101 .await
14102 .unwrap();
14103
14104 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
14106 .session(session)
14107 .manifest_enabled(true)
14108 .dir_listing_enabled(true)
14109 .dir_listing_to_manifest_migration_enabled(false)
14110 .build()
14111 .await
14112 .unwrap();
14113
14114 listing_count.store(0, Ordering::SeqCst);
14116
14117 let mut exists_req = TableExistsRequest::new();
14119 exists_req.id = Some(vec!["test_table".to_string()]);
14120 hybrid_ns.table_exists(exists_req).await.unwrap();
14121
14122 let count = listing_count.load(Ordering::SeqCst);
14123 assert_eq!(
14124 count, 1,
14125 "Expected exactly 1 listing call for table_exists \
14126 without migration mode, but got {}",
14127 count
14128 );
14129
14130 listing_count.store(0, Ordering::SeqCst);
14132
14133 let mut describe_req = DescribeTableRequest::new();
14134 describe_req.id = Some(vec!["test_table".to_string()]);
14135 hybrid_ns.describe_table(describe_req).await.unwrap();
14136
14137 let count = listing_count.load(Ordering::SeqCst);
14138 assert_eq!(
14139 count, 1,
14140 "Expected exactly 1 listing call for describe_table \
14141 without migration mode, but got {}",
14142 count
14143 );
14144 }
14145
14146 #[tokio::test]
14147 async fn test_build_and_root_reads_do_not_create_manifest() {
14148 let temp_dir = TempStdDir::default();
14149 let temp_path = temp_dir.to_str().unwrap();
14150 let manifest_path = std::path::Path::new(temp_path).join("__manifest");
14151
14152 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
14153 .manifest_enabled(false)
14154 .dir_listing_enabled(true)
14155 .build()
14156 .await
14157 .unwrap();
14158 create_scalar_table(&dir_only_ns, "catalog").await;
14159 assert!(!manifest_path.exists());
14160
14161 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14162 .manifest_enabled(true)
14163 .dir_listing_enabled(true)
14164 .build()
14165 .await
14166 .unwrap();
14167 assert!(!manifest_path.exists());
14168
14169 let mut exists_req = TableExistsRequest::new();
14170 exists_req.id = Some(vec!["catalog".to_string()]);
14171 namespace.table_exists(exists_req).await.unwrap();
14172 assert!(!manifest_path.exists());
14173
14174 let mut describe_req = DescribeTableRequest::new();
14175 describe_req.id = Some(vec!["catalog".to_string()]);
14176 namespace.describe_table(describe_req).await.unwrap();
14177 assert!(!manifest_path.exists());
14178
14179 let list_response = namespace
14180 .list_tables(ListTablesRequest {
14181 id: Some(vec![]),
14182 ..Default::default()
14183 })
14184 .await
14185 .unwrap();
14186 assert_eq!(list_response.tables, vec!["catalog".to_string()]);
14187 assert!(!manifest_path.exists());
14188
14189 let mut list_namespaces_req = ListNamespacesRequest::new();
14190 list_namespaces_req.id = Some(vec!["workspace".to_string()]);
14191 let err = namespace
14192 .list_namespaces(list_namespaces_req)
14193 .await
14194 .unwrap_err();
14195 assert!(err.to_string().contains("__manifest"));
14196 assert!(!manifest_path.exists());
14197
14198 let err = namespace
14199 .list_tables(ListTablesRequest {
14200 id: Some(vec!["workspace".to_string()]),
14201 ..Default::default()
14202 })
14203 .await
14204 .unwrap_err();
14205 assert!(err.to_string().contains("__manifest"));
14206 assert!(!manifest_path.exists());
14207
14208 let mut child_describe_req = DescribeTableRequest::new();
14209 child_describe_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
14210 let err = namespace
14211 .describe_table(child_describe_req)
14212 .await
14213 .unwrap_err();
14214 assert!(err.to_string().contains("__manifest"));
14215 assert!(!manifest_path.exists());
14216
14217 let mut child_exists_req = TableExistsRequest::new();
14218 child_exists_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
14219 let err = namespace.table_exists(child_exists_req).await.unwrap_err();
14220 assert!(err.to_string().contains("__manifest"));
14221 assert!(!manifest_path.exists());
14222
14223 let mut create_ns_req = CreateNamespaceRequest::new();
14224 create_ns_req.id = Some(vec!["workspace".to_string()]);
14225 namespace.create_namespace(create_ns_req).await.unwrap();
14226 assert!(manifest_path.exists());
14227 }
14228
14229 #[tokio::test]
14230 async fn test_migrate_updates_read_opened_legacy_manifest() {
14231 let temp_dir = TempStdDir::default();
14232 let temp_path = temp_dir.to_str().unwrap();
14233 create_legacy_manifest_without_primary_key_metadata(temp_path).await;
14234 assert!(!manifest_has_primary_key_metadata(temp_path).await);
14235
14236 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14237 .manifest_enabled(true)
14238 .dir_listing_enabled(true)
14239 .build()
14240 .await
14241 .unwrap();
14242 assert!(!manifest_has_primary_key_metadata(temp_path).await);
14243
14244 let migrated = namespace.migrate().await.unwrap();
14245 assert_eq!(migrated, 0);
14246 assert!(manifest_has_primary_key_metadata(temp_path).await);
14247 }
14248
14249 #[tokio::test]
14250 async fn test_describe_declared_table_checks_versions_only_when_requested() {
14251 let temp_dir = TempStdDir::default();
14252 let temp_path = temp_dir.to_str().unwrap();
14253 let root_uri = file_object_store_uri(temp_path);
14254 let listing_count = Arc::new(AtomicUsize::new(0));
14255 let session = build_listing_counting_session(listing_count.clone());
14256
14257 let namespace = DirectoryNamespaceBuilder::new(root_uri)
14258 .session(session)
14259 .manifest_enabled(false)
14260 .dir_listing_enabled(true)
14261 .build()
14262 .await
14263 .unwrap();
14264
14265 let mut declare_req = DeclareTableRequest::new();
14266 declare_req.id = Some(vec!["test_table".to_string()]);
14267 namespace.declare_table(declare_req).await.unwrap();
14268
14269 listing_count.store(0, Ordering::SeqCst);
14270
14271 let mut describe_req = DescribeTableRequest::new();
14272 describe_req.id = Some(vec!["test_table".to_string()]);
14273 let describe_response = namespace.describe_table(describe_req).await.unwrap();
14274
14275 assert_eq!(describe_response.is_only_declared, None);
14276 assert_eq!(
14277 listing_count.load(Ordering::SeqCst),
14278 1,
14279 "Default describe_table should only list the table directory"
14280 );
14281
14282 listing_count.store(0, Ordering::SeqCst);
14283
14284 let mut describe_req = DescribeTableRequest::new();
14285 describe_req.id = Some(vec!["test_table".to_string()]);
14286 describe_req.check_declared = Some(true);
14287 let describe_response = namespace.describe_table(describe_req).await.unwrap();
14288
14289 assert_eq!(describe_response.is_only_declared, Some(true));
14290 assert_eq!(
14291 listing_count.load(Ordering::SeqCst),
14292 2,
14293 "check_declared describe_table should list the table directory and _versions"
14294 );
14295 }
14296
14297 #[tokio::test]
14298 async fn test_dir_listing_extra_calls_with_migration() {
14299 let temp_dir = TempStdDir::default();
14300 let temp_path = temp_dir.to_str().unwrap();
14301 let root_uri = file_object_store_uri(temp_path);
14302 let listing_count = Arc::new(AtomicUsize::new(0));
14303 let session = build_listing_counting_session(listing_count.clone());
14304
14305 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
14307 .session(session.clone())
14308 .manifest_enabled(false)
14309 .dir_listing_enabled(true)
14310 .build()
14311 .await
14312 .unwrap();
14313
14314 let schema = create_test_schema();
14315 let ipc_data = create_test_ipc_data(&schema);
14316 let mut create_req = CreateTableRequest::new();
14317 create_req.id = Some(vec!["test_table".to_string()]);
14318 dir_only_ns
14319 .create_table(create_req, Bytes::from(ipc_data))
14320 .await
14321 .unwrap();
14322
14323 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
14324 .session(session)
14325 .manifest_enabled(true)
14326 .dir_listing_enabled(true)
14327 .dir_listing_to_manifest_migration_enabled(true)
14328 .build()
14329 .await
14330 .unwrap();
14331
14332 listing_count.store(0, Ordering::SeqCst);
14336
14337 let mut exists_req = TableExistsRequest::new();
14338 exists_req.id = Some(vec!["test_table".to_string()]);
14339 hybrid_ns.table_exists(exists_req).await.unwrap();
14340
14341 let count = listing_count.load(Ordering::SeqCst);
14342 assert_eq!(
14343 count, 1,
14344 "Expected exactly 1 listing call for table_exists with migration mode \
14345 (table directory fallback; manifest reload uses the version hint), but got {}",
14346 count
14347 );
14348
14349 listing_count.store(0, Ordering::SeqCst);
14351
14352 let mut describe_req = DescribeTableRequest::new();
14353 describe_req.id = Some(vec!["test_table".to_string()]);
14354 hybrid_ns.describe_table(describe_req).await.unwrap();
14355
14356 let count = listing_count.load(Ordering::SeqCst);
14357 assert_eq!(
14358 count, 1,
14359 "Expected exactly 1 listing call for describe_table with migration mode \
14360 (table directory fallback; manifest reload uses the version hint), but got {}",
14361 count
14362 );
14363 }
14364
14365 #[tokio::test]
14366 async fn test_manifest_reload_observes_new_version_from_other_namespace() {
14367 let temp_dir = TempStdDir::default();
14368 let temp_path = temp_dir.to_str().unwrap();
14369
14370 let namespace_a = DirectoryNamespaceBuilder::new(temp_path)
14371 .manifest_enabled(true)
14372 .dir_listing_enabled(false)
14373 .build()
14374 .await
14375 .unwrap();
14376 create_scalar_table(&namespace_a, "alpha").await;
14377
14378 let namespace_b = DirectoryNamespaceBuilder::new(temp_path)
14379 .manifest_enabled(true)
14380 .dir_listing_enabled(false)
14381 .build()
14382 .await
14383 .unwrap();
14384 create_scalar_table(&namespace_b, "beta").await;
14385
14386 let response = namespace_a
14387 .list_tables(ListTablesRequest {
14388 id: Some(vec![]),
14389 ..Default::default()
14390 })
14391 .await
14392 .unwrap();
14393
14394 let mut tables = response.tables;
14395 tables.sort();
14396 assert_eq!(tables, vec!["alpha", "beta"]);
14397 }
14398
14399 #[tokio::test]
14400 async fn test_migration_not_found_errors_include_table_id() {
14401 let temp_dir = TempStdDir::default();
14402 let temp_path = temp_dir.to_str().unwrap();
14403
14404 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14405 .manifest_enabled(true)
14406 .dir_listing_enabled(true)
14407 .dir_listing_to_manifest_migration_enabled(true)
14408 .build()
14409 .await
14410 .unwrap();
14411
14412 let mut exists_req = TableExistsRequest::new();
14413 exists_req.id = Some(vec!["missing_table".to_string()]);
14414 let err = namespace.table_exists(exists_req).await.unwrap_err();
14415 assert!(matches!(err, Error::Namespace { .. }));
14416 let err_msg = err.to_string();
14417 assert!(err_msg.contains("Table not found"));
14418 assert!(err_msg.contains("table id 'missing_table'"));
14419
14420 let mut describe_req = DescribeTableRequest::new();
14421 describe_req.id = Some(vec!["missing_table".to_string()]);
14422 let err = namespace.describe_table(describe_req).await.unwrap_err();
14423 assert!(matches!(err, Error::Namespace { .. }));
14424 let err_msg = err.to_string();
14425 assert!(err_msg.contains("Table not found"));
14426 assert!(err_msg.contains("table id 'missing_table'"));
14427 }
14428
14429 #[tokio::test]
14430 async fn test_manifest_not_found_errors_include_full_table_id() {
14431 use lance_namespace::models::CreateNamespaceRequest;
14432
14433 let temp_dir = TempStdDir::default();
14434 let temp_path = temp_dir.to_str().unwrap();
14435
14436 let namespace = DirectoryNamespaceBuilder::new(temp_path)
14437 .manifest_enabled(true)
14438 .dir_listing_enabled(true)
14439 .build()
14440 .await
14441 .unwrap();
14442
14443 let mut create_ns_req = CreateNamespaceRequest::new();
14444 create_ns_req.id = Some(vec!["workspace".to_string()]);
14445 namespace.create_namespace(create_ns_req).await.unwrap();
14446
14447 let missing_table_id = vec!["workspace".to_string(), "missing_table".to_string()];
14448
14449 let mut exists_req = TableExistsRequest::new();
14450 exists_req.id = Some(missing_table_id.clone());
14451 let err = namespace.table_exists(exists_req).await.unwrap_err();
14452 assert!(matches!(err, Error::Namespace { .. }));
14453 let err_msg = err.to_string();
14454 assert!(err_msg.contains("Table not found"));
14455 assert!(err_msg.contains("table id 'workspace$missing_table'"));
14456
14457 let mut describe_req = DescribeTableRequest::new();
14458 describe_req.id = Some(missing_table_id);
14459 let err = namespace.describe_table(describe_req).await.unwrap_err();
14460 assert!(matches!(err, Error::Namespace { .. }));
14461 let err_msg = err.to_string();
14462 assert!(err_msg.contains("Table not found"));
14463 assert!(err_msg.contains("table id 'workspace$missing_table'"));
14464 }
14465
14466 async fn create_tagged_test_table(
14469 versions: u32,
14470 ) -> (Arc<DirectoryNamespace>, TempStdDir, Vec<String>) {
14471 use arrow::array::{Int32Array, RecordBatchIterator};
14472 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
14473 use arrow::record_batch::RecordBatch;
14474 use lance::dataset::{Dataset, WriteMode, WriteParams};
14475
14476 assert!(versions >= 1, "versions must be at least 1");
14477
14478 let temp_dir = TempStdDir::default();
14479 let temp_path = temp_dir.to_str().unwrap();
14480
14481 let namespace = Arc::new(
14482 DirectoryNamespaceBuilder::new(temp_path)
14483 .build()
14484 .await
14485 .unwrap(),
14486 );
14487 let table_id = vec!["tag_table".to_string()];
14488 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
14489 "id",
14490 DataType::Int32,
14491 false,
14492 )]));
14493 let initial_batch = RecordBatch::try_new(
14494 arrow_schema.clone(),
14495 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
14496 )
14497 .unwrap();
14498 let batches = RecordBatchIterator::new(vec![Ok(initial_batch)], arrow_schema.clone());
14499 let write_params = WriteParams {
14500 mode: WriteMode::Create,
14501 ..Default::default()
14502 };
14503
14504 let mut dataset = Dataset::write_into_namespace(
14505 batches,
14506 namespace.clone() as Arc<dyn LanceNamespace>,
14507 table_id.clone(),
14508 Some(write_params),
14509 )
14510 .await
14511 .unwrap();
14512
14513 for i in 1..versions {
14514 let value_start = (i as i32) * 10;
14515 let batch = RecordBatch::try_new(
14516 arrow_schema.clone(),
14517 vec![Arc::new(Int32Array::from(vec![
14518 value_start,
14519 value_start + 1,
14520 ]))],
14521 )
14522 .unwrap();
14523 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
14524 dataset.append(batches, None).await.unwrap();
14525 }
14526
14527 (namespace, temp_dir, table_id)
14528 }
14529
14530 fn namespace_code(err: &Error) -> Option<ErrorCode> {
14532 match err {
14533 Error::Namespace { source, .. } => {
14534 source.downcast_ref::<NamespaceError>().map(|e| e.code())
14535 }
14536 _ => None,
14537 }
14538 }
14539
14540 #[tokio::test]
14541 async fn test_create_and_list_branches() {
14542 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14543
14544 namespace
14545 .create_table_branch(CreateTableBranchRequest {
14546 id: Some(table_id.clone()),
14547 name: "dev".to_string(),
14548 ..Default::default()
14549 })
14550 .await
14551 .unwrap();
14552 namespace
14553 .create_table_branch(CreateTableBranchRequest {
14554 id: Some(table_id.clone()),
14555 name: "staging".to_string(),
14556 ..Default::default()
14557 })
14558 .await
14559 .unwrap();
14560
14561 let resp = namespace
14562 .list_table_branches(ListTableBranchesRequest {
14563 id: Some(table_id.clone()),
14564 ..Default::default()
14565 })
14566 .await
14567 .unwrap();
14568 assert_eq!(
14569 resp.branches.len(),
14570 2,
14571 "expected 2 branches, got: {:?}",
14572 resp.branches
14573 );
14574 assert!(resp.branches.contains_key("dev"));
14575 assert!(resp.branches.contains_key("staging"));
14576 assert!(resp.page_token.is_none());
14577
14578 namespace
14580 .delete_table_branch(DeleteTableBranchRequest {
14581 id: Some(table_id.clone()),
14582 name: "dev".to_string(),
14583 ..Default::default()
14584 })
14585 .await
14586 .unwrap();
14587
14588 let resp = namespace
14589 .list_table_branches(ListTableBranchesRequest {
14590 id: Some(table_id),
14591 ..Default::default()
14592 })
14593 .await
14594 .unwrap();
14595 assert_eq!(resp.branches.len(), 1, "expected 1 branch after delete");
14596 assert!(!resp.branches.contains_key("dev"));
14597 assert!(resp.branches.contains_key("staging"));
14598 }
14599
14600 #[tokio::test]
14601 async fn test_create_branch_from_version() {
14602 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14603
14604 namespace
14606 .create_table_branch(CreateTableBranchRequest {
14607 id: Some(table_id.clone()),
14608 name: "from-v1".to_string(),
14609 from_version: Some(1),
14610 ..Default::default()
14611 })
14612 .await
14613 .unwrap();
14614
14615 let resp = namespace
14616 .list_table_branches(ListTableBranchesRequest {
14617 id: Some(table_id),
14618 ..Default::default()
14619 })
14620 .await
14621 .unwrap();
14622 let branch = resp
14623 .branches
14624 .get("from-v1")
14625 .expect("forked branch should be listed");
14626 assert_eq!(
14627 branch.parent_version, 1,
14628 "branch should fork from version 1"
14629 );
14630 assert!(
14631 branch.parent_branch.is_none(),
14632 "a branch forked from main has no parent branch"
14633 );
14634 }
14635
14636 #[tokio::test]
14641 async fn test_create_branch_from_other_branch() {
14642 use lance::dataset::builder::DatasetBuilder;
14643
14644 let (namespace, _temp_dir) = create_test_namespace().await;
14645 create_scalar_table(&namespace, "users").await; create_branch_with_commits(&namespace, "users", "dev", 1).await;
14648 let main_ds = open_dataset(&namespace, "users").await;
14650 append_scalar_version(main_ds.uri(), 500).await; namespace
14653 .create_table_branch(CreateTableBranchRequest {
14654 id: Some(vec!["users".to_string()]),
14655 name: "child".to_string(),
14656 from_branch: Some("dev".to_string()),
14657 from_version: Some(2),
14658 ..Default::default()
14659 })
14660 .await
14661 .unwrap();
14662
14663 let child_ds = DatasetBuilder::from_uri(main_ds.uri())
14664 .with_branch("child", None)
14665 .load()
14666 .await
14667 .unwrap();
14668 let ids = scan_id_column(&child_ds).await;
14669 assert!(
14670 ids.contains(&100) && ids.contains(&101),
14671 "child must contain dev's appended rows, got: {:?}",
14672 ids
14673 );
14674 assert!(
14675 !ids.contains(&500),
14676 "child must not contain main's diverged rows, got: {:?}",
14677 ids
14678 );
14679
14680 let listed = namespace
14682 .list_table_branches(ListTableBranchesRequest {
14683 id: Some(vec!["users".to_string()]),
14684 ..Default::default()
14685 })
14686 .await
14687 .unwrap();
14688 assert_eq!(
14689 listed
14690 .branches
14691 .get("child")
14692 .unwrap()
14693 .parent_branch
14694 .as_deref(),
14695 Some("dev")
14696 );
14697 }
14698
14699 #[tokio::test]
14700 async fn test_create_existing_branch_conflict() {
14701 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14702
14703 namespace
14704 .create_table_branch(CreateTableBranchRequest {
14705 id: Some(table_id.clone()),
14706 name: "dev".to_string(),
14707 ..Default::default()
14708 })
14709 .await
14710 .unwrap();
14711
14712 let err = namespace
14713 .create_table_branch(CreateTableBranchRequest {
14714 id: Some(table_id),
14715 name: "dev".to_string(),
14716 ..Default::default()
14717 })
14718 .await
14719 .unwrap_err();
14720 assert_eq!(
14721 namespace_code(&err),
14722 Some(ErrorCode::TableBranchAlreadyExists),
14723 "expected TableBranchAlreadyExists, got: {}",
14724 err
14725 );
14726 assert!(
14727 err.to_string().to_lowercase().contains("already exists"),
14728 "expected already-exists message, got: {}",
14729 err
14730 );
14731 }
14732
14733 #[tokio::test]
14734 async fn test_delete_unknown_branch() {
14735 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14736
14737 let err = namespace
14738 .delete_table_branch(DeleteTableBranchRequest {
14739 id: Some(table_id),
14740 name: "does-not-exist".to_string(),
14741 ..Default::default()
14742 })
14743 .await
14744 .unwrap_err();
14745 assert_eq!(
14746 namespace_code(&err),
14747 Some(ErrorCode::TableBranchNotFound),
14748 "expected TableBranchNotFound, got: {}",
14749 err
14750 );
14751 assert!(
14752 err.to_string().to_lowercase().contains("not found"),
14753 "expected not-found message, got: {}",
14754 err
14755 );
14756 }
14757
14758 #[tokio::test]
14759 async fn test_delete_referenced_branch_conflict() {
14760 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14761
14762 namespace
14764 .create_table_branch(CreateTableBranchRequest {
14765 id: Some(table_id.clone()),
14766 name: "parent".to_string(),
14767 ..Default::default()
14768 })
14769 .await
14770 .unwrap();
14771 namespace
14772 .create_table_branch(CreateTableBranchRequest {
14773 id: Some(table_id.clone()),
14774 name: "child".to_string(),
14775 from_branch: Some("parent".to_string()),
14776 ..Default::default()
14777 })
14778 .await
14779 .unwrap();
14780
14781 let listed = namespace
14783 .list_table_branches(ListTableBranchesRequest {
14784 id: Some(table_id.clone()),
14785 ..Default::default()
14786 })
14787 .await
14788 .unwrap();
14789 let child = listed
14790 .branches
14791 .get("child")
14792 .expect("child branch should be listed");
14793 assert_eq!(
14794 child.parent_branch.as_deref(),
14795 Some("parent"),
14796 "child should record parent branch as its fork point"
14797 );
14798 assert!(
14799 child.parent_version >= 1,
14800 "child should record the parent version it forked from, got {}",
14801 child.parent_version
14802 );
14803
14804 let err = namespace
14807 .delete_table_branch(DeleteTableBranchRequest {
14808 id: Some(table_id),
14809 name: "parent".to_string(),
14810 ..Default::default()
14811 })
14812 .await
14813 .unwrap_err();
14814 assert_eq!(
14815 namespace_code(&err),
14816 Some(ErrorCode::InvalidInput),
14817 "expected InvalidInput for deleting a referenced branch, got: {}",
14818 err
14819 );
14820 assert!(
14821 err.to_string().to_lowercase().contains("referenced"),
14822 "error should explain the branch is still referenced, got: {}",
14823 err
14824 );
14825 }
14826
14827 #[tokio::test]
14828 async fn test_branch_name_required() {
14829 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14830
14831 let create_err = namespace
14832 .create_table_branch(CreateTableBranchRequest {
14833 id: Some(table_id.clone()),
14834 name: String::new(),
14835 ..Default::default()
14836 })
14837 .await
14838 .unwrap_err();
14839 assert_eq!(
14840 namespace_code(&create_err),
14841 Some(ErrorCode::InvalidInput),
14842 "empty name on create should be InvalidInput, got: {}",
14843 create_err
14844 );
14845 assert!(
14846 create_err
14847 .to_string()
14848 .to_lowercase()
14849 .contains("must not be empty")
14850 );
14851
14852 let delete_err = namespace
14853 .delete_table_branch(DeleteTableBranchRequest {
14854 id: Some(table_id),
14855 name: String::new(),
14856 ..Default::default()
14857 })
14858 .await
14859 .unwrap_err();
14860 assert_eq!(
14861 namespace_code(&delete_err),
14862 Some(ErrorCode::InvalidInput),
14863 "empty name on delete should be InvalidInput, got: {}",
14864 delete_err
14865 );
14866 assert!(
14867 delete_err
14868 .to_string()
14869 .to_lowercase()
14870 .contains("must not be empty")
14871 );
14872 }
14873
14874 #[tokio::test]
14875 async fn test_create_branch_rejects_negative_from_version() {
14876 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14877
14878 let err = namespace
14879 .create_table_branch(CreateTableBranchRequest {
14880 id: Some(table_id),
14881 name: "dev".to_string(),
14882 from_version: Some(-1),
14883 ..Default::default()
14884 })
14885 .await
14886 .unwrap_err();
14887 assert_eq!(
14888 namespace_code(&err),
14889 Some(ErrorCode::InvalidInput),
14890 "negative from_version should be InvalidInput, got: {}",
14891 err
14892 );
14893 assert!(err.to_string().to_lowercase().contains("from_version"));
14894 }
14895
14896 #[tokio::test]
14897 async fn test_create_branch_nonexistent_from_version() {
14898 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14899
14900 let err = namespace
14903 .create_table_branch(CreateTableBranchRequest {
14904 id: Some(table_id),
14905 name: "dev".to_string(),
14906 from_version: Some(999),
14907 ..Default::default()
14908 })
14909 .await
14910 .unwrap_err();
14911 assert_eq!(
14912 namespace_code(&err),
14913 Some(ErrorCode::InvalidInput),
14914 "non-existent from_version should map to InvalidInput, got: {}",
14915 err
14916 );
14917 assert!(
14918 err.to_string().to_lowercase().contains("does not exist"),
14919 "error should name the missing source, got: {}",
14920 err
14921 );
14922 }
14923
14924 #[tokio::test]
14925 async fn test_create_and_list_tags() {
14926 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14927
14928 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
14929 req.id = Some(table_id.clone());
14930 namespace.create_table_tag(req).await.unwrap();
14931
14932 let mut req = CreateTableTagRequest::new("v2".to_string(), 2);
14933 req.id = Some(table_id.clone());
14934 namespace.create_table_tag(req).await.unwrap();
14935
14936 let mut list_req = ListTableTagsRequest::new();
14937 list_req.id = Some(table_id);
14938 let resp = namespace.list_table_tags(list_req).await.unwrap();
14939
14940 assert_eq!(resp.tags.len(), 2, "expected 2 tags, got: {:?}", resp.tags);
14941 assert_eq!(resp.tags.get("v1").unwrap().version, 1);
14942 assert_eq!(resp.tags.get("v2").unwrap().version, 2);
14943 assert!(resp.page_token.is_none());
14944 }
14945
14946 #[tokio::test]
14947 async fn test_create_existing_tag_conflict() {
14948 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14949
14950 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
14951 req.id = Some(table_id.clone());
14952 namespace.create_table_tag(req).await.unwrap();
14953
14954 let mut req = CreateTableTagRequest::new("v1".to_string(), 2);
14955 req.id = Some(table_id);
14956 let err = namespace.create_table_tag(req).await.unwrap_err();
14957 assert!(
14958 err.to_string().to_lowercase().contains("already exists"),
14959 "expected already-exists error, got: {}",
14960 err
14961 );
14962 }
14963
14964 #[tokio::test]
14965 async fn test_get_tag_version() {
14966 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14967
14968 let mut req = CreateTableTagRequest::new("release".to_string(), 2);
14969 req.id = Some(table_id.clone());
14970 namespace.create_table_tag(req).await.unwrap();
14971
14972 let mut get_req = GetTableTagVersionRequest::new("release".to_string());
14973 get_req.id = Some(table_id);
14974 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
14975 assert_eq!(resp.version, 2);
14976 assert_eq!(resp.branch, None);
14977 }
14978
14979 #[tokio::test]
14980 async fn test_get_unknown_tag() {
14981 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
14982
14983 let mut get_req = GetTableTagVersionRequest::new("does-not-exist".to_string());
14984 get_req.id = Some(table_id);
14985 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
14986 assert!(
14987 err.to_string().to_lowercase().contains("not found"),
14988 "expected not-found error, got: {}",
14989 err
14990 );
14991 }
14992
14993 #[tokio::test]
14994 async fn test_update_tag_to_new_version() {
14995 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
14996
14997 let mut req = CreateTableTagRequest::new("rolling".to_string(), 1);
14998 req.id = Some(table_id.clone());
14999 namespace.create_table_tag(req).await.unwrap();
15000
15001 let mut update_req = UpdateTableTagRequest::new("rolling".to_string(), 3);
15002 update_req.id = Some(table_id.clone());
15003 namespace.update_table_tag(update_req).await.unwrap();
15004
15005 let mut get_req = GetTableTagVersionRequest::new("rolling".to_string());
15006 get_req.id = Some(table_id);
15007 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15008 assert_eq!(resp.version, 3);
15009 }
15010
15011 #[tokio::test]
15012 async fn test_update_unknown_tag() {
15013 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15014
15015 let mut update_req = UpdateTableTagRequest::new("ghost".to_string(), 1);
15016 update_req.id = Some(table_id);
15017 let err = namespace.update_table_tag(update_req).await.unwrap_err();
15018 assert!(
15019 err.to_string().to_lowercase().contains("not found"),
15020 "expected not-found error, got: {}",
15021 err
15022 );
15023 }
15024
15025 #[tokio::test]
15026 async fn test_delete_tag() {
15027 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15028
15029 let mut req = CreateTableTagRequest::new("doomed".to_string(), 1);
15030 req.id = Some(table_id.clone());
15031 namespace.create_table_tag(req).await.unwrap();
15032
15033 let mut delete_req = DeleteTableTagRequest::new("doomed".to_string());
15034 delete_req.id = Some(table_id.clone());
15035 namespace.delete_table_tag(delete_req).await.unwrap();
15036
15037 let mut list_req = ListTableTagsRequest::new();
15038 list_req.id = Some(table_id.clone());
15039 let resp = namespace.list_table_tags(list_req).await.unwrap();
15040 assert!(resp.tags.is_empty(), "tag should be removed after delete");
15041
15042 let mut get_req = GetTableTagVersionRequest::new("doomed".to_string());
15044 get_req.id = Some(table_id);
15045 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
15046 assert!(err.to_string().to_lowercase().contains("not found"));
15047 }
15048
15049 #[tokio::test]
15050 async fn test_delete_unknown_tag() {
15051 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15052
15053 let mut delete_req = DeleteTableTagRequest::new("nope".to_string());
15054 delete_req.id = Some(table_id);
15055 let err = namespace.delete_table_tag(delete_req).await.unwrap_err();
15056 assert!(
15057 err.to_string().to_lowercase().contains("not found"),
15058 "expected not-found error, got: {}",
15059 err
15060 );
15061 }
15062
15063 #[tokio::test]
15064 async fn test_create_tag_invalid_version() {
15065 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15066
15067 let mut req = CreateTableTagRequest::new("v0".to_string(), 0);
15069 req.id = Some(table_id.clone());
15070 let err = namespace.create_table_tag(req).await.unwrap_err();
15071 assert!(
15072 err.to_string().to_lowercase().contains("positive"),
15073 "expected positive-version error, got: {}",
15074 err
15075 );
15076
15077 let mut req = CreateTableTagRequest::new(String::new(), 1);
15079 req.id = Some(table_id);
15080 let err = namespace.create_table_tag(req).await.unwrap_err();
15081 assert!(
15082 err.to_string().to_lowercase().contains("must not be empty"),
15083 "expected empty-tag-name error, got: {}",
15084 err
15085 );
15086 }
15087
15088 #[tokio::test]
15089 async fn test_create_tag_table_not_found() {
15090 let (namespace, _temp_dir) = create_test_namespace().await;
15091
15092 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15093 req.id = Some(vec!["does_not_exist".to_string()]);
15094 let err = namespace.create_table_tag(req).await.unwrap_err();
15095 let msg = err.to_string();
15096 assert!(
15097 msg.contains("Table") && msg.to_lowercase().contains("not found"),
15098 "expected TableNotFound error, got: {}",
15099 err
15100 );
15101 }
15102 #[tokio::test]
15103 async fn test_alter_table_drop_columns_missing_id() {
15104 use lance_namespace::models::AlterTableDropColumnsRequest;
15105
15106 let (namespace, _temp_dir) = create_test_namespace().await;
15107
15108 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
15109 let result = namespace.alter_table_drop_columns(request).await;
15110 assert!(result.is_err(), "Should fail when table ID is missing");
15111 }
15112
15113 #[tokio::test]
15114 async fn test_alter_table_drop_columns_nonexistent_table() {
15115 use lance_namespace::models::AlterTableDropColumnsRequest;
15116
15117 let (namespace, _temp_dir) = create_test_namespace().await;
15118
15119 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
15120 request.id = Some(vec!["nonexistent".to_string()]);
15121 let result = namespace.alter_table_drop_columns(request).await;
15122 assert!(result.is_err(), "Should fail when table does not exist");
15123 }
15124
15125 #[tokio::test]
15126 async fn test_create_branch_on_managed_dataset_succeeds() {
15127 use lance::dataset::builder::DatasetBuilder;
15128
15129 let temp = TempStdDir::default();
15130 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
15131 let table_id = vec!["t".to_string()];
15132 let mut main = create_managed_table(&ns, &table_id).await;
15133
15134 let fork_version = main.version().version;
15135 let branch = main
15136 .create_branch("exp", fork_version, None)
15137 .await
15138 .expect("create_branch failed");
15139 assert_eq!(branch.manifest.branch.as_deref(), Some("exp"));
15140 assert_eq!(scan_id_column(&branch).await, vec![1, 2]);
15141
15142 let reopened = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
15143 .await
15144 .unwrap()
15145 .with_branch("exp", None)
15146 .load()
15147 .await
15148 .expect("reopen branch failed");
15149 assert_eq!(scan_id_column(&reopened).await, vec![1, 2]);
15150 }
15151
15152 #[tokio::test]
15153 async fn test_alter_transaction_set_status() {
15154 use lance_namespace::models::{
15155 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15156 DescribeTransactionRequest,
15157 };
15158
15159 let (namespace, _temp_dir) = create_test_namespace().await;
15160 create_scalar_table(&namespace, "users").await;
15161 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15162 .await
15163 .expect("create_scalar_index should return a transaction id");
15164
15165 let describe_resp = namespace
15167 .describe_transaction(DescribeTransactionRequest {
15168 id: Some(vec!["users".to_string(), txn_id.clone()]),
15169 ..Default::default()
15170 })
15171 .await
15172 .unwrap();
15173 assert_eq!(describe_resp.status, "SUCCEEDED");
15174
15175 let response = namespace
15177 .alter_transaction(AlterTransactionRequest {
15178 id: Some(vec!["users".to_string(), txn_id.clone()]),
15179 actions: vec![AlterTransactionAction {
15180 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15181 status: Some("Canceled".to_string()),
15182 })),
15183 set_property_action: None,
15184 unset_property_action: None,
15185 }],
15186 ..Default::default()
15187 })
15188 .await
15189 .unwrap();
15190 assert_eq!(response.status, "Canceled");
15191 assert!(response.properties.is_some());
15192 let props = response.properties.unwrap();
15193 assert_eq!(props.get("uuid"), Some(&txn_id));
15194 assert_eq!(props.get("operation"), Some(&"CreateIndex".to_string()));
15195 }
15196
15197 #[tokio::test]
15198 async fn test_alter_transaction_set_property() {
15199 use lance_namespace::models::{
15200 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15201 };
15202
15203 let (namespace, _temp_dir) = create_test_namespace().await;
15204 create_scalar_table(&namespace, "users").await;
15205 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15206 .await
15207 .expect("create_scalar_index should return a transaction id");
15208
15209 let response = namespace
15210 .alter_transaction(AlterTransactionRequest {
15211 id: Some(vec!["users".to_string(), txn_id.clone()]),
15212 actions: vec![AlterTransactionAction {
15213 set_status_action: None,
15214 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15215 key: Some("custom_key".to_string()),
15216 value: Some("custom_value".to_string()),
15217 mode: None,
15218 })),
15219 unset_property_action: None,
15220 }],
15221 ..Default::default()
15222 })
15223 .await
15224 .unwrap();
15225 assert_eq!(response.status, "SUCCEEDED");
15226 let props = response.properties.unwrap();
15227 assert_eq!(props.get("custom_key"), Some(&"custom_value".to_string()));
15228 }
15229
15230 #[tokio::test]
15231 async fn test_alter_transaction_set_property_fail_mode() {
15232 use lance_namespace::models::{
15233 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15234 };
15235
15236 let (namespace, _temp_dir) = create_test_namespace().await;
15237 create_scalar_table(&namespace, "users").await;
15238 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15239 .await
15240 .expect("create_scalar_index should return a transaction id");
15241
15242 namespace
15244 .alter_transaction(AlterTransactionRequest {
15245 id: Some(vec!["users".to_string(), txn_id.clone()]),
15246 actions: vec![AlterTransactionAction {
15247 set_status_action: None,
15248 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15249 key: Some("custom_key".to_string()),
15250 value: Some("initial_value".to_string()),
15251 mode: None,
15252 })),
15253 unset_property_action: None,
15254 }],
15255 ..Default::default()
15256 })
15257 .await
15258 .unwrap();
15259
15260 let result = namespace
15263 .alter_transaction(AlterTransactionRequest {
15264 id: Some(vec!["users".to_string(), txn_id.clone()]),
15265 actions: vec![AlterTransactionAction {
15266 set_status_action: None,
15267 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15268 key: Some("custom_key".to_string()),
15269 value: Some("new_value".to_string()),
15270 mode: Some("Fail".to_string()),
15271 })),
15272 unset_property_action: None,
15273 }],
15274 ..Default::default()
15275 })
15276 .await;
15277 assert!(result.is_err());
15278 }
15279
15280 #[tokio::test]
15281 async fn test_alter_transaction_unset_property() {
15282 use lance_namespace::models::{
15283 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15284 AlterTransactionUnsetProperty,
15285 };
15286
15287 let (namespace, _temp_dir) = create_test_namespace().await;
15288 create_scalar_table(&namespace, "users").await;
15289 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15290 .await
15291 .expect("create_scalar_index should return a transaction id");
15292
15293 let response = namespace
15295 .alter_transaction(AlterTransactionRequest {
15296 id: Some(vec!["users".to_string(), txn_id.clone()]),
15297 actions: vec![
15298 AlterTransactionAction {
15299 set_status_action: None,
15300 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15301 key: Some("temp_key".to_string()),
15302 value: Some("temp_value".to_string()),
15303 mode: None,
15304 })),
15305 unset_property_action: None,
15306 },
15307 AlterTransactionAction {
15308 set_status_action: None,
15309 set_property_action: None,
15310 unset_property_action: Some(Box::new(AlterTransactionUnsetProperty {
15311 key: Some("temp_key".to_string()),
15312 mode: None,
15313 })),
15314 },
15315 ],
15316 ..Default::default()
15317 })
15318 .await
15319 .unwrap();
15320 assert_eq!(response.status, "SUCCEEDED");
15321 let props = response.properties.unwrap();
15322 assert!(!props.contains_key("temp_key"));
15323 }
15324
15325 #[tokio::test]
15326 async fn test_alter_transaction_invalid_status() {
15327 use lance_namespace::models::{
15328 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15329 };
15330
15331 let (namespace, _temp_dir) = create_test_namespace().await;
15332 create_scalar_table(&namespace, "users").await;
15333 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
15334 .await
15335 .expect("create_scalar_index should return a transaction id");
15336
15337 let result = namespace
15338 .alter_transaction(AlterTransactionRequest {
15339 id: Some(vec!["users".to_string(), txn_id.clone()]),
15340 actions: vec![AlterTransactionAction {
15341 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15342 status: Some("InvalidStatus".to_string()),
15343 })),
15344 set_property_action: None,
15345 unset_property_action: None,
15346 }],
15347 ..Default::default()
15348 })
15349 .await;
15350 assert!(result.is_err());
15351 }
15352
15353 #[tokio::test]
15354 async fn test_alter_transaction_not_found() {
15355 use lance_namespace::models::{
15356 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15357 };
15358
15359 let (namespace, _temp_dir) = create_test_namespace().await;
15360 create_scalar_table(&namespace, "users").await;
15361
15362 let result = namespace
15364 .alter_transaction(AlterTransactionRequest {
15365 id: Some(vec!["users".to_string(), "non_existent_txn".to_string()]),
15366 actions: vec![AlterTransactionAction {
15367 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15368 status: Some("Canceled".to_string()),
15369 })),
15370 set_property_action: None,
15371 unset_property_action: None,
15372 }],
15373 ..Default::default()
15374 })
15375 .await;
15376 assert!(result.is_err());
15377 }
15378
15379 #[tokio::test]
15380 async fn test_alter_transaction_missing_id() {
15381 use lance_namespace::models::{
15382 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
15383 };
15384
15385 let (namespace, _temp_dir) = create_test_namespace().await;
15386
15387 let result = namespace
15389 .alter_transaction(AlterTransactionRequest {
15390 id: None,
15391 actions: vec![AlterTransactionAction {
15392 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15393 status: Some("Canceled".to_string()),
15394 })),
15395 set_property_action: None,
15396 unset_property_action: None,
15397 }],
15398 ..Default::default()
15399 })
15400 .await;
15401 assert!(result.is_err());
15402
15403 let result = namespace
15405 .alter_transaction(AlterTransactionRequest {
15406 id: Some(vec!["users".to_string()]),
15407 actions: vec![AlterTransactionAction {
15408 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15409 status: Some("Canceled".to_string()),
15410 })),
15411 set_property_action: None,
15412 unset_property_action: None,
15413 }],
15414 ..Default::default()
15415 })
15416 .await;
15417 assert!(result.is_err());
15418 }
15419
15420 #[tokio::test]
15421 async fn test_alter_transaction_persists_changes() {
15422 use lance_namespace::models::{
15423 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
15424 AlterTransactionSetStatus, DescribeTransactionRequest,
15425 };
15426
15427 let (namespace, _temp_dir) = create_test_namespace().await;
15428 create_scalar_table(&namespace, "users").await;
15429 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
15430
15431 let txn_id = transaction_id.expect("scalar index should produce a transaction id");
15432
15433 namespace
15435 .alter_transaction(AlterTransactionRequest {
15436 id: Some(vec!["users".to_string(), txn_id.clone()]),
15437 actions: vec![
15438 AlterTransactionAction {
15439 set_status_action: Some(Box::new(AlterTransactionSetStatus {
15440 status: Some("Canceled".to_string()),
15441 })),
15442 set_property_action: None,
15443 unset_property_action: None,
15444 },
15445 AlterTransactionAction {
15446 set_status_action: None,
15447 set_property_action: Some(Box::new(AlterTransactionSetProperty {
15448 key: Some("owner".to_string()),
15449 value: Some("alice".to_string()),
15450 mode: None,
15451 })),
15452 unset_property_action: None,
15453 },
15454 ],
15455 ..Default::default()
15456 })
15457 .await
15458 .unwrap();
15459
15460 let describe_resp = namespace
15463 .describe_transaction(DescribeTransactionRequest {
15464 id: Some(vec!["users".to_string(), txn_id.clone()]),
15465 ..Default::default()
15466 })
15467 .await
15468 .unwrap();
15469 let props = describe_resp.properties.expect("properties should be set");
15470 assert_eq!(props.get("owner"), Some(&"alice".to_string()));
15471 assert!(!props.contains_key("_status"));
15475
15476 let follow_up = namespace
15477 .alter_transaction(AlterTransactionRequest {
15478 id: Some(vec!["users".to_string(), txn_id.clone()]),
15479 actions: vec![],
15480 ..Default::default()
15481 })
15482 .await
15483 .unwrap();
15484 assert_eq!(follow_up.status, "Canceled");
15485 let follow_up_props = follow_up.properties.unwrap();
15486 assert_eq!(follow_up_props.get("owner"), Some(&"alice".to_string()));
15487 }
15488}