1pub mod manifest;
10pub mod manifest_feature_flags;
11
12use arrow::array::Float32Array;
13use arrow::record_batch::RecordBatchIterator;
14use arrow_ipc::reader::StreamReader;
15use async_trait::async_trait;
16use bytes::Bytes;
17use futures::{StreamExt, TryStreamExt};
18use lance::dataset::builder::DatasetBuilder;
19use lance::dataset::refs::check_valid_branch;
20use lance::dataset::scanner::Scanner;
21use lance::dataset::statistics::DatasetStatisticsExt;
22use lance::dataset::transaction::{Operation, Transaction};
23use lance::dataset::{
24 Dataset, MergeInsertBuilder, UpdateBuilder, WhenMatched, WhenNotMatched,
25 WhenNotMatchedBySource, WriteMode, WriteParams,
26};
27use lance::index::{DatasetIndexExt, IndexParams, vector::VectorIndexParams};
28use lance::session::Session;
29use lance_index::scalar::{
30 BuiltinIndexType, FullTextSearchQuery, InvertedIndexParams, ScalarIndexParams,
31};
32use lance_index::vector::{
33 bq::{RABIT_MAX_NUM_BITS, RABIT_MIN_NUM_BITS, RQBuildParams, validate_supported_rq_num_bits},
34 hnsw::builder::HnswBuildParams,
35 ivf::IvfBuildParams,
36 pq::PQBuildParams,
37 sq::builder::SQBuildParams,
38};
39use lance_index::{IndexType, is_system_index};
40use lance_io::object_store::throttle::is_throttle_error;
41use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry, ReadDirOptions};
42use lance_linalg::distance::MetricType;
43use lance_table::io::commit::{ManifestNamingScheme, VERSIONS_DIR};
44use object_store::ObjectStoreExt;
45use object_store::path::Path;
46use object_store::{
47 Error as ObjectStoreError, ObjectMeta, ObjectStore as OSObjectStore, PutMode, PutOptions,
48};
49use std::collections::HashMap;
50use std::io::Cursor;
51use std::sync::{Arc, Mutex};
52use tokio::sync::OnceCell;
53
54use crate::context::DynamicContextProvider;
55use crate::merge_insert_on_columns;
56use lance_namespace::models::{
57 AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest,
58 AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse,
59 AlterTransactionRequest, AlterTransactionResponse, AnalyzeTableQueryPlanRequest,
60 BatchDeleteTableVersionsRequest, BatchDeleteTableVersionsResponse,
61 BranchContents as ModelBranchContents, CountTableRowsRequest, CreateNamespaceRequest,
62 CreateNamespaceResponse, CreateTableBranchRequest, CreateTableBranchResponse,
63 CreateTableIndexRequest, CreateTableIndexResponse, CreateTableRequest, CreateTableResponse,
64 CreateTableScalarIndexResponse, CreateTableTagRequest, CreateTableTagResponse,
65 CreateTableVersionRequest, CreateTableVersionResponse, DeclareTableRequest,
66 DeclareTableResponse, DeleteFromTableRequest, DeleteFromTableResponse,
67 DeleteTableBranchRequest, DeleteTableBranchResponse, DeleteTableTagRequest,
68 DeleteTableTagResponse, DescribeNamespaceRequest, DescribeNamespaceResponse,
69 DescribeTableIndexStatsRequest, DescribeTableIndexStatsResponse, DescribeTableRequest,
70 DescribeTableResponse, DescribeTableVersionRequest, DescribeTableVersionResponse,
71 DescribeTransactionRequest, DescribeTransactionResponse, DropNamespaceRequest,
72 DropNamespaceResponse, DropTableIndexRequest, DropTableIndexResponse, DropTableRequest,
73 DropTableResponse, ExplainTableQueryPlanRequest, FragmentStats, FragmentSummary,
74 GetTableStatsRequest, GetTableStatsResponse, GetTableTagVersionRequest,
75 GetTableTagVersionResponse, Identity, IndexContent, InsertIntoTableRequest,
76 InsertIntoTableResponse, ListNamespacesRequest, ListNamespacesResponse,
77 ListTableBranchesRequest, ListTableBranchesResponse, ListTableIndicesRequest,
78 ListTableIndicesResponse, ListTableTagsRequest, ListTableTagsResponse,
79 ListTableVersionsRequest, ListTableVersionsResponse, ListTablesRequest, ListTablesResponse,
80 MergeInsertIntoTableRequest, MergeInsertIntoTableResponse, NamespaceExistsRequest,
81 QueryTableRequest, QueryTableRequestColumns, QueryTableRequestVector, RestoreTableRequest,
82 RestoreTableResponse, TableExistsRequest, TableVersion, TagContents as ModelTagContents,
83 UpdateTableRequest, UpdateTableResponse, UpdateTableSchemaMetadataRequest,
84 UpdateTableSchemaMetadataResponse, UpdateTableTagRequest, UpdateTableTagResponse,
85};
86
87use lance_core::utils::parse::str_to_bool;
88use lance_core::{Error, Result, box_error};
89use lance_index::scalar::inverted::query::{
90 BooleanQuery, BoostQuery, FtsQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery,
91};
92use lance_namespace::LanceNamespace;
93use lance_namespace::error::NamespaceError;
94use lance_namespace::schema::arrow_schema_to_json;
95
96use crate::credentials::{
97 CredentialVendor, create_credential_vendor_for_location, has_credential_vendor_config,
98};
99
100#[derive(Debug, Default)]
105pub struct OpsMetrics {
106 counters: Mutex<HashMap<String, u64>>,
107}
108
109impl OpsMetrics {
110 pub fn increment(&self, operation: &str) {
112 if let Ok(mut counters) = self.counters.lock() {
113 *counters.entry(operation.to_string()).or_insert(0) += 1;
114 }
115 }
116
117 pub fn retrieve(&self) -> HashMap<String, u64> {
119 self.counters.lock().map(|c| c.clone()).unwrap_or_default()
120 }
121
122 pub fn reset(&self) {
124 if let Ok(mut counters) = self.counters.lock() {
125 counters.clear();
126 }
127 }
128}
129
130pub(crate) fn build_sql_expressions(
133 new_columns: &[lance_namespace::models::AddColumnsEntry],
134) -> Result<Vec<(String, String)>> {
135 new_columns
136 .iter()
137 .map(|col| {
138 let expression = col.expression.clone().and_then(|opt| opt).ok_or_else(|| {
141 Error::invalid_input(format!(
142 "Expression is required for new column '{}'",
143 col.name
144 ))
145 })?;
146 Ok((col.name.clone(), expression))
147 })
148 .collect()
149}
150
151pub(crate) fn build_column_alterations(
154 alterations: &[lance_namespace::models::AlterColumnsEntry],
155) -> Result<Vec<lance::dataset::ColumnAlteration>> {
156 alterations
157 .iter()
158 .map(|entry| {
159 let mut alteration = lance::dataset::ColumnAlteration::new(entry.path.clone());
160 if let Some(Some(rename)) = &entry.rename {
162 alteration = alteration.rename(rename.clone());
163 }
164 if let Some(Some(nullable)) = entry.nullable {
166 alteration = alteration.set_nullable(nullable);
167 }
168 if let Some(data_type) = &entry.data_type
170 && !data_type.is_null()
171 {
172 let type_str = data_type.as_str().ok_or_else(|| {
173 Error::invalid_input(format!(
174 "data_type for column '{}' must be a JSON string, got: {}",
175 entry.path, data_type
176 ))
177 })?;
178 let json_type =
179 lance_namespace::models::JsonArrowDataType::new(type_str.to_string());
180 let dt =
181 lance_namespace::schema::convert_json_arrow_type(&json_type).map_err(|e| {
182 Error::invalid_input(format!(
183 "Failed to parse data_type '{}' for column '{}': {}",
184 type_str, entry.path, e
185 ))
186 })?;
187 alteration = alteration.cast_to(dt);
188 }
189 Ok(alteration)
190 })
191 .collect()
192}
193
194pub(crate) struct TableStatus {
199 pub(crate) exists: bool,
201 pub(crate) is_deregistered: bool,
203 pub(crate) has_reserved_file: bool,
205}
206
207enum DirectoryIndexParams {
208 Scalar {
209 index_type: IndexType,
210 params: ScalarIndexParams,
211 },
212 Inverted(InvertedIndexParams),
213 Vector {
214 index_type: IndexType,
215 params: VectorIndexParams,
216 },
217}
218
219impl DirectoryIndexParams {
220 fn index_type(&self) -> IndexType {
221 match self {
222 Self::Scalar { index_type, .. } | Self::Vector { index_type, .. } => *index_type,
223 Self::Inverted(_) => IndexType::Inverted,
224 }
225 }
226
227 fn params(&self) -> &dyn IndexParams {
228 match self {
229 Self::Scalar { params, .. } => params,
230 Self::Inverted(params) => params,
231 Self::Vector { params, .. } => params,
232 }
233 }
234}
235
236#[derive(Clone)]
271pub struct DirectoryNamespaceBuilder {
272 root: String,
273 storage_options: Option<HashMap<String, String>>,
274 session: Option<Arc<Session>>,
275 manifest_enabled: bool,
276 dir_listing_enabled: bool,
277 inline_optimization_enabled: bool,
278 table_version_tracking_enabled: bool,
279 dir_listing_to_manifest_migration_enabled: bool,
284 credential_vendor_properties: HashMap<String, String>,
285 context_provider: Option<Arc<dyn DynamicContextProvider>>,
286 commit_retries: Option<u32>,
287 vend_input_storage_options: bool,
290 vend_input_storage_options_refresh_interval_millis: Option<u64>,
295 ops_metrics_enabled: bool,
297}
298
299impl std::fmt::Debug for DirectoryNamespaceBuilder {
300 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
301 f.debug_struct("DirectoryNamespaceBuilder")
302 .field("root", &self.root)
303 .field("storage_options", &self.storage_options)
304 .field("manifest_enabled", &self.manifest_enabled)
305 .field("dir_listing_enabled", &self.dir_listing_enabled)
306 .field(
307 "inline_optimization_enabled",
308 &self.inline_optimization_enabled,
309 )
310 .field(
311 "table_version_tracking_enabled",
312 &self.table_version_tracking_enabled,
313 )
314 .field(
315 "dir_listing_to_manifest_migration_enabled",
316 &self.dir_listing_to_manifest_migration_enabled,
317 )
318 .field(
319 "context_provider",
320 &self.context_provider.as_ref().map(|_| "Some(...)"),
321 )
322 .field(
323 "vend_input_storage_options",
324 &self.vend_input_storage_options,
325 )
326 .field(
327 "vend_input_storage_options_refresh_interval_millis",
328 &self.vend_input_storage_options_refresh_interval_millis,
329 )
330 .field("ops_metrics_enabled", &self.ops_metrics_enabled)
331 .finish()
332 }
333}
334
335impl DirectoryNamespaceBuilder {
336 pub fn new(root: impl Into<String>) -> Self {
342 Self {
343 root: root.into().trim_end_matches('/').to_string(),
344 storage_options: None,
345 session: None,
346 manifest_enabled: true,
347 dir_listing_enabled: true, inline_optimization_enabled: false,
349 table_version_tracking_enabled: false, dir_listing_to_manifest_migration_enabled: false, credential_vendor_properties: HashMap::new(),
352 context_provider: None,
353 commit_retries: None,
354 vend_input_storage_options: false,
355 vend_input_storage_options_refresh_interval_millis: None,
356 ops_metrics_enabled: false,
357 }
358 }
359
360 pub fn manifest_enabled(mut self, enabled: bool) -> Self {
365 self.manifest_enabled = enabled;
366 self
367 }
368
369 pub fn dir_listing_enabled(mut self, enabled: bool) -> Self {
374 self.dir_listing_enabled = enabled;
375 self
376 }
377
378 pub fn dir_listing_to_manifest_migration_enabled(mut self, enabled: bool) -> Self {
385 self.dir_listing_to_manifest_migration_enabled = enabled;
386 self
387 }
388
389 pub fn inline_optimization_enabled(mut self, enabled: bool) -> Self {
394 self.inline_optimization_enabled = enabled;
395 self
396 }
397
398 pub fn table_version_tracking_enabled(mut self, enabled: bool) -> Self {
406 self.table_version_tracking_enabled = enabled;
407 self
408 }
409
410 pub fn from_properties(
478 properties: HashMap<String, String>,
479 session: Option<Arc<Session>>,
480 ) -> Result<Self> {
481 let root = properties.get("root").cloned().ok_or_else(|| {
483 lance_core::Error::from(NamespaceError::InvalidInput {
484 message: "Missing required property 'root' for directory namespace".to_string(),
485 })
486 })?;
487
488 let storage_options: HashMap<String, String> = properties
490 .iter()
491 .filter_map(|(k, v)| {
492 k.strip_prefix("storage.")
493 .map(|key| (key.to_string(), v.clone()))
494 })
495 .collect();
496
497 let storage_options = if storage_options.is_empty() {
498 None
499 } else {
500 Some(storage_options)
501 };
502
503 let manifest_enabled = properties
505 .get("manifest_enabled")
506 .and_then(|v| str_to_bool(v))
507 .unwrap_or(true);
508
509 let dir_listing_enabled = properties
511 .get("dir_listing_enabled")
512 .and_then(|v| str_to_bool(v))
513 .unwrap_or(true);
514
515 let inline_optimization_enabled = properties
517 .get("inline_optimization_enabled")
518 .and_then(|v| str_to_bool(v))
519 .unwrap_or(false);
520
521 let table_version_tracking_enabled = properties
523 .get("table_version_tracking_enabled")
524 .and_then(|v| str_to_bool(v))
525 .unwrap_or(false);
526
527 let dir_listing_to_manifest_migration_enabled = properties
529 .get("dir_listing_to_manifest_migration_enabled")
530 .and_then(|v| str_to_bool(v))
531 .unwrap_or(false);
532
533 let credential_vendor_properties: HashMap<String, String> = properties
537 .iter()
538 .filter_map(|(k, v)| {
539 k.strip_prefix("credential_vendor.")
540 .map(|key| (key.to_string(), v.clone()))
541 })
542 .collect();
543
544 let commit_retries = properties
545 .get("commit_retries")
546 .and_then(|v| v.parse::<u32>().ok());
547
548 let vend_input_storage_options = properties
550 .get("vend_input_storage_options")
551 .and_then(|v| str_to_bool(v))
552 .unwrap_or(false);
553
554 let vend_input_storage_options_refresh_interval_millis = properties
556 .get("vend_input_storage_options_refresh_interval_millis")
557 .and_then(|v| v.parse::<u64>().ok());
558
559 let ops_metrics_enabled = properties
561 .get("ops_metrics_enabled")
562 .and_then(|v| str_to_bool(v))
563 .unwrap_or(false);
564
565 Ok(Self {
566 root: root.trim_end_matches('/').to_string(),
567 storage_options,
568 session,
569 manifest_enabled,
570 dir_listing_enabled,
571 inline_optimization_enabled,
572 table_version_tracking_enabled,
573 dir_listing_to_manifest_migration_enabled,
574 credential_vendor_properties,
575 context_provider: None,
576 commit_retries,
577 vend_input_storage_options,
578 vend_input_storage_options_refresh_interval_millis,
579 ops_metrics_enabled,
580 })
581 }
582
583 pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
590 self.storage_options
591 .get_or_insert_with(HashMap::new)
592 .insert(key.into(), value.into());
593 self
594 }
595
596 pub fn storage_options(mut self, options: HashMap<String, String>) -> Self {
602 self.storage_options
603 .get_or_insert_with(HashMap::new)
604 .extend(options);
605 self
606 }
607
608 pub fn session(mut self, session: Arc<Session>) -> Self {
618 self.session = Some(session);
619 self
620 }
621
622 pub fn commit_retries(mut self, retries: u32) -> Self {
625 self.commit_retries = Some(retries);
626 self
627 }
628
629 pub fn credential_vendor_property(
657 mut self,
658 key: impl Into<String>,
659 value: impl Into<String>,
660 ) -> Self {
661 self.credential_vendor_properties
662 .insert(key.into(), value.into());
663 self
664 }
665
666 pub fn credential_vendor_properties(mut self, properties: HashMap<String, String>) -> Self {
674 self.credential_vendor_properties.extend(properties);
675 self
676 }
677
678 pub fn context_provider(mut self, provider: Arc<dyn DynamicContextProvider>) -> Self {
688 self.context_provider = Some(provider);
689 self
690 }
691
692 pub fn vend_input_storage_options(mut self, enabled: bool) -> Self {
701 self.vend_input_storage_options = enabled;
702 self
703 }
704
705 pub fn vend_input_storage_options_refresh_interval_millis(
717 mut self,
718 interval_millis: u64,
719 ) -> Self {
720 self.vend_input_storage_options_refresh_interval_millis = Some(interval_millis);
721 self
722 }
723
724 pub fn ops_metrics_enabled(mut self, enabled: bool) -> Self {
732 self.ops_metrics_enabled = enabled;
733 self
734 }
735
736 pub async fn build(self) -> Result<DirectoryNamespace> {
749 let (object_store, base_path) =
750 Self::initialize_object_store(&self.root, &self.storage_options, &self.session).await?;
751
752 let manifest_ns = if self.manifest_enabled {
753 match manifest::ManifestNamespace::open_from_directory(
754 self.root.clone(),
755 self.storage_options.clone(),
756 self.session.clone(),
757 object_store.clone(),
758 base_path.clone(),
759 self.dir_listing_enabled,
760 self.inline_optimization_enabled,
761 self.commit_retries,
762 )
763 .await
764 {
765 Ok(ns) => Some(Arc::new(ns)),
766 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
767 return Err(e);
771 }
772 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => {
773 log::debug!("Manifest namespace does not exist yet: {}", e);
774 None
775 }
776 Err(e) => return Err(e),
777 }
778 } else {
779 None
780 };
781 let manifest_cell = OnceCell::new();
782 if let Some(manifest_ns) = manifest_ns {
783 let _ = manifest_cell.set(manifest_ns);
784 }
785
786 let credential_vendor = if has_credential_vendor_config(&self.credential_vendor_properties)
788 {
789 create_credential_vendor_for_location(&self.root, &self.credential_vendor_properties)
790 .await?
791 .map(Arc::from)
792 } else {
793 None
794 };
795
796 let ops_metrics = if self.ops_metrics_enabled {
797 Some(Arc::new(OpsMetrics::default()))
798 } else {
799 None
800 };
801
802 Ok(DirectoryNamespace {
803 root: self.root,
804 storage_options: self.storage_options,
805 session: self.session,
806 object_store,
807 base_path,
808 manifest_ns: manifest_cell,
809 write_manifest_ns: OnceCell::new(),
810 manifest_enabled: self.manifest_enabled,
811 dir_listing_enabled: self.dir_listing_enabled,
812 inline_optimization_enabled: self.inline_optimization_enabled,
813 commit_retries: self.commit_retries,
814 dir_listing_to_manifest_migration_enabled: self
815 .dir_listing_to_manifest_migration_enabled,
816 table_version_tracking_enabled: self.table_version_tracking_enabled,
817 credential_vendor,
818 context_provider: self.context_provider,
819 vend_input_storage_options: self.vend_input_storage_options,
820 vend_input_storage_options_refresh_interval_millis: self
821 .vend_input_storage_options_refresh_interval_millis,
822 ops_metrics,
823 })
824 }
825
826 async fn initialize_object_store(
828 root: &str,
829 storage_options: &Option<HashMap<String, String>>,
830 session: &Option<Arc<Session>>,
831 ) -> Result<(Arc<ObjectStore>, Path)> {
832 let accessor = storage_options.clone().map(|opts| {
834 Arc::new(lance_io::object_store::StorageOptionsAccessor::with_static_options(opts))
835 });
836 let params = ObjectStoreParams {
837 storage_options_accessor: accessor,
838 ..Default::default()
839 };
840
841 let registry = if let Some(session) = session {
843 session.store_registry()
844 } else {
845 Arc::new(ObjectStoreRegistry::default())
846 };
847
848 let (object_store, base_path) = ObjectStore::from_uri_and_params(registry, root, ¶ms)
850 .await
851 .map_err(|e| {
852 lance_core::Error::from(NamespaceError::Internal {
853 message: format!("Failed to create object store: {:?}", e),
854 })
855 })?;
856
857 Ok((object_store, base_path))
858 }
859}
860
861pub struct DirectoryNamespace {
885 root: String,
886 storage_options: Option<HashMap<String, String>>,
887 session: Option<Arc<Session>>,
888 object_store: Arc<ObjectStore>,
889 base_path: Path,
890 manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
891 write_manifest_ns: OnceCell<Arc<manifest::ManifestNamespace>>,
892 manifest_enabled: bool,
893 dir_listing_enabled: bool,
894 inline_optimization_enabled: bool,
895 commit_retries: Option<u32>,
896 dir_listing_to_manifest_migration_enabled: bool,
900 table_version_tracking_enabled: bool,
903 credential_vendor: Option<Arc<dyn CredentialVendor>>,
906 #[allow(dead_code)]
909 context_provider: Option<Arc<dyn DynamicContextProvider>>,
910 vend_input_storage_options: bool,
912 vend_input_storage_options_refresh_interval_millis: Option<u64>,
915 ops_metrics: Option<Arc<OpsMetrics>>,
917}
918
919impl std::fmt::Debug for DirectoryNamespace {
920 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
921 write!(f, "{}", self.namespace_id())
922 }
923}
924
925impl std::fmt::Display for DirectoryNamespace {
926 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
927 write!(f, "{}", self.namespace_id())
928 }
929}
930
931struct ExistingTableVersionResolve<'a> {
933 staging_path: &'a Path,
934 final_path: &'a Path,
935 version: u64,
936 table_uri: &'a str,
937 final_meta: &'a ObjectMeta,
938 request_manifest_size: Option<i64>,
939}
940
941struct TableDeleteEntry {
944 table_id: Option<Vec<String>>,
945 ranges: Vec<(i64, i64)>,
946}
947
948#[derive(Debug, Clone, Default)]
958struct TransactionAlteration {
959 status: Option<String>,
961 properties: HashMap<String, String>,
965 removed_properties: std::collections::HashSet<String>,
969}
970
971impl TransactionAlteration {
972 const F_STATUS: &'static str = "status";
974 const F_PROPERTIES: &'static str = "properties";
975 const F_REMOVED_PROPERTIES: &'static str = "removed_properties";
976
977 fn to_json_bytes(&self) -> serde_json::Result<Vec<u8>> {
984 serde_json::to_vec(&serde_json::json!({
985 Self::F_STATUS: self.status,
986 Self::F_PROPERTIES: self.properties,
987 Self::F_REMOVED_PROPERTIES: self.removed_properties,
988 }))
989 }
990
991 fn from_json_slice(bytes: &[u8]) -> serde_json::Result<Self> {
996 let mut obj: serde_json::Map<String, serde_json::Value> = serde_json::from_slice(bytes)?;
997 Ok(Self {
998 status: serde_json::from_value(
999 obj.remove(Self::F_STATUS)
1000 .unwrap_or(serde_json::Value::Null),
1001 )?,
1002 properties: serde_json::from_value(
1003 obj.remove(Self::F_PROPERTIES)
1004 .unwrap_or(serde_json::Value::Null),
1005 )
1006 .unwrap_or_default(),
1007 removed_properties: serde_json::from_value(
1008 obj.remove(Self::F_REMOVED_PROPERTIES)
1009 .unwrap_or(serde_json::Value::Null),
1010 )
1011 .unwrap_or_default(),
1012 })
1013 }
1014}
1015
1016fn apply_probe_bounds(
1017 scanner: &mut Scanner,
1018 nprobes: Option<i32>,
1019 minimum_nprobes: Option<i32>,
1020 maximum_nprobes: Option<i32>,
1021) -> Result<()> {
1022 let parse_probe_count = |name: &str, value: i32| {
1023 usize::try_from(value)
1024 .map_err(|_| Error::invalid_input(format!("{name} must be non-negative")))
1025 };
1026
1027 if let Some(nprobes) = nprobes {
1028 scanner.nprobes(parse_probe_count("nprobes", nprobes)?);
1029 }
1030 if let Some(minimum_nprobes) = minimum_nprobes {
1031 scanner.minimum_nprobes(parse_probe_count("minimum_nprobes", minimum_nprobes)?);
1032 }
1033 if let Some(maximum_nprobes) = maximum_nprobes {
1034 scanner.maximum_nprobes(parse_probe_count("maximum_nprobes", maximum_nprobes)?);
1035 }
1036 Ok(())
1037}
1038
1039impl DirectoryNamespace {
1040 fn manifest_ns_for_read(&self) -> Option<&Arc<manifest::ManifestNamespace>> {
1041 self.write_manifest_ns
1042 .get()
1043 .or_else(|| self.manifest_ns.get())
1044 }
1045
1046 async fn manifest_ns_for_write(&self) -> Result<Option<Arc<manifest::ManifestNamespace>>> {
1047 if !self.manifest_enabled {
1048 return Ok(None);
1049 }
1050
1051 let manifest_ns = self
1052 .write_manifest_ns
1053 .get_or_try_init(|| async {
1054 manifest::ManifestNamespace::from_directory(
1055 self.root.clone(),
1056 self.storage_options.clone(),
1057 self.session.clone(),
1058 self.object_store.clone(),
1059 self.base_path.clone(),
1060 self.dir_listing_enabled,
1061 self.inline_optimization_enabled,
1062 self.commit_retries,
1063 )
1064 .await
1065 .map(Arc::new)
1066 })
1067 .await?;
1068 Ok(Some(manifest_ns.clone()))
1069 }
1070
1071 async fn ensure_read_manifest(&self) -> Result<()> {
1083 if !self.manifest_enabled
1084 || self.manifest_ns.get().is_some()
1085 || self.write_manifest_ns.get().is_some()
1086 {
1087 return Ok(());
1088 }
1089 match self
1090 .manifest_ns
1091 .get_or_try_init(|| async {
1092 manifest::ManifestNamespace::open_from_directory(
1093 self.root.clone(),
1094 self.storage_options.clone(),
1095 self.session.clone(),
1096 self.object_store.clone(),
1097 self.base_path.clone(),
1098 self.dir_listing_enabled,
1099 self.inline_optimization_enabled,
1100 self.commit_retries,
1101 )
1102 .await
1103 .map(Arc::new)
1104 })
1105 .await
1106 {
1107 Ok(_) => Ok(()),
1108 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(()),
1113 Err(e) => Err(e),
1114 }
1115 }
1116
1117 fn child_namespace_requires_manifest_error(&self) -> Error {
1118 if self.manifest_enabled {
1119 NamespaceError::NamespaceNotFound {
1120 message: "Child namespace reads require an existing __manifest dataset".to_string(),
1121 }
1122 .into()
1123 } else {
1124 NamespaceError::Unsupported {
1125 message: "Child namespaces are only supported when manifest mode is enabled"
1126 .to_string(),
1127 }
1128 .into()
1129 }
1130 }
1131
1132 fn apply_pagination(
1145 names: &mut Vec<String>,
1146 page_token: Option<String>,
1147 limit: Option<i32>,
1148 ) -> Option<String> {
1149 names.sort();
1151
1152 if let Some(start_after) = page_token {
1154 if let Some(index) = names
1155 .iter()
1156 .position(|name| name.as_str() > start_after.as_str())
1157 {
1158 names.drain(0..index);
1159 } else {
1160 names.clear();
1161 }
1162 }
1163
1164 if let Some(limit) = limit
1166 && limit >= 0
1167 {
1168 let limit = limit as usize;
1169 if names.len() > limit {
1170 let next_page_token = if limit > 0 {
1171 Some(names[limit - 1].clone())
1172 } else {
1173 None
1174 };
1175 names.truncate(limit);
1176 return next_page_token;
1177 }
1178 }
1179
1180 None
1181 }
1182
1183 const LIST_DIRECTORY_PAGE_SIZE: usize = 1000;
1187
1188 async fn list_directory_tables(&self) -> Result<Vec<String>> {
1190 let mut tables = Vec::new();
1191 let mut page_token = None;
1192
1193 loop {
1194 let page = self
1195 .object_store
1196 .read_dir_page(
1197 self.base_path.clone(),
1198 ReadDirOptions {
1199 page_token,
1200 limit: Some(Self::LIST_DIRECTORY_PAGE_SIZE),
1205 },
1206 )
1207 .await
1208 .map_err(|e| {
1209 lance_core::Error::from(NamespaceError::Internal {
1210 message: format!("Failed to list directory: {:?}", e),
1211 })
1212 })?;
1213
1214 let candidates: Vec<String> = page
1215 .result
1216 .common_prefixes
1217 .iter()
1218 .chain(page.result.objects.iter().map(|o| &o.location))
1219 .filter_map(|p| {
1220 p.filename()?
1221 .trim_end_matches('/')
1222 .strip_suffix(".lance")
1223 .map(|name| name.to_string())
1224 })
1225 .collect();
1226
1227 let mut stream =
1231 futures::stream::iter(candidates.into_iter().map(|table_name| async move {
1232 let status = self.check_table_status(&table_name).await?;
1233 Ok::<Option<String>, Error>((!status.is_deregistered).then_some(table_name))
1234 }))
1235 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1236
1237 while let Some(result) = stream.next().await {
1238 if let Some(table_name) = result? {
1239 tables.push(table_name);
1240 }
1241 }
1242
1243 page_token = page.page_token;
1248 if page_token.is_none() {
1249 break;
1250 }
1251 }
1252
1253 Ok(tables)
1254 }
1255
1256 fn validate_root_namespace_id(id: &Option<Vec<String>>) -> Result<()> {
1258 if let Some(id) = id
1259 && !id.is_empty()
1260 {
1261 return Err(NamespaceError::Unsupported {
1262 message: format!(
1263 "Directory namespace only supports root namespace operations, but got namespace ID: {:?}. Expected empty ID.",
1264 id
1265 ),
1266 }
1267 .into());
1268 }
1269 Ok(())
1270 }
1271
1272 fn table_name_from_id(id: &Option<Vec<String>>) -> Result<String> {
1274 let id = id.as_ref().ok_or_else(|| {
1275 lance_core::Error::from(NamespaceError::InvalidInput {
1276 message: "Directory namespace table ID cannot be empty".to_string(),
1277 })
1278 })?;
1279
1280 if id.len() != 1 {
1281 return Err(NamespaceError::Unsupported {
1282 message: format!(
1283 "Multi-level table IDs are only supported when manifest mode is enabled, but got: {:?}",
1284 id
1285 ),
1286 }
1287 .into());
1288 }
1289
1290 Ok(id[0].clone())
1291 }
1292
1293 fn format_table_id(table_id: &[String]) -> String {
1294 format!(
1295 "table id '{}'",
1296 manifest::ManifestNamespace::str_object_id(table_id)
1297 )
1298 }
1299
1300 fn format_table_id_from_request(id: &Option<Vec<String>>) -> String {
1301 id.as_ref()
1302 .map(|table_id| Self::format_table_id(table_id))
1303 .unwrap_or_else(|| "table id '<unknown>'".to_string())
1304 }
1305
1306 async fn resolve_table_location(&self, id: &Option<Vec<String>>) -> Result<String> {
1307 let mut describe_req = DescribeTableRequest::new();
1308 describe_req.id = id.clone();
1309 describe_req.load_detailed_metadata = Some(false);
1310
1311 let describe_resp = self.describe_table_impl(describe_req).await?;
1313
1314 describe_resp.location.ok_or_else(|| {
1315 lance_core::Error::from(NamespaceError::TableNotFound {
1316 message: format!("Table location not found for: {:?}", id),
1317 })
1318 })
1319 }
1320
1321 fn map_tag_error(err: lance_core::Error, tag: &str, table_uri: &str) -> lance_core::Error {
1324 match err {
1325 lance_core::Error::RefNotFound { .. } => NamespaceError::TableTagNotFound {
1326 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1327 }
1328 .into(),
1329 lance_core::Error::RefConflict { .. } => NamespaceError::TableTagAlreadyExists {
1330 message: format!("tag '{}' for table at '{}'", tag, table_uri),
1331 }
1332 .into(),
1333 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1334 message: format!("invalid tag '{}': {}", tag, message),
1335 }
1336 .into(),
1337 lance_core::Error::VersionNotFound { message } => {
1338 NamespaceError::TableVersionNotFound {
1339 message: format!(
1340 "version referenced by tag '{}' not found for table at '{}': {}",
1341 tag, table_uri, message
1342 ),
1343 }
1344 .into()
1345 }
1346 other => NamespaceError::Internal {
1347 message: format!(
1348 "tag operation failed for tag '{}' on table at '{}': {}",
1349 tag, table_uri, other
1350 ),
1351 }
1352 .into(),
1353 }
1354 }
1355
1356 fn map_branch_error(
1362 err: lance_core::Error,
1363 branch: &str,
1364 table_uri: &str,
1365 ) -> lance_core::Error {
1366 match err {
1367 lance_core::Error::RefNotFound { .. } => NamespaceError::TableBranchNotFound {
1368 message: format!("branch '{}' for table at '{}'", branch, table_uri),
1369 }
1370 .into(),
1371 lance_core::Error::InvalidRef { message } => NamespaceError::InvalidInput {
1372 message: format!("invalid branch '{}': {}", branch, message),
1373 }
1374 .into(),
1375 lance_core::Error::VersionNotFound { message } => {
1376 NamespaceError::TableVersionNotFound {
1377 message: format!(
1378 "source version for branch '{}' not found for table at '{}': {}",
1379 branch, table_uri, message
1380 ),
1381 }
1382 .into()
1383 }
1384 other => NamespaceError::Internal {
1385 message: format!(
1386 "branch operation failed for branch '{}' on table at '{}': {}",
1387 branch, table_uri, other
1388 ),
1389 }
1390 .into(),
1391 }
1392 }
1393
1394 fn map_mutation_error(
1406 err: lance_core::Error,
1407 operation: &str,
1408 table_uri: &str,
1409 ) -> lance_core::Error {
1410 let detail = err.to_string();
1411 let ns_err = match &err {
1412 lance_core::Error::InvalidInput { .. }
1413 | lance_core::Error::Unprocessable { .. }
1414 | lance_core::Error::InvalidRef { .. } => NamespaceError::InvalidInput {
1415 message: format!(
1416 "Invalid input for {} on table at '{}': {}",
1417 operation, table_uri, detail
1418 ),
1419 },
1420 lance_core::Error::NotFound { .. } | lance_core::Error::DatasetNotFound { .. } => {
1421 NamespaceError::TableNotFound {
1422 message: format!(
1423 "Table at '{}' not found while running {}: {}",
1424 table_uri, operation, detail
1425 ),
1426 }
1427 }
1428 lance_core::Error::SchemaMismatch { .. } | lance_core::Error::Schema { .. } => {
1429 NamespaceError::TableSchemaValidationError {
1430 message: format!(
1431 "Schema validation failed for {} on table at '{}': {}",
1432 operation, table_uri, detail
1433 ),
1434 }
1435 }
1436 lance_core::Error::CommitConflict { .. } => NamespaceError::Throttling {
1440 message: format!(
1441 "Too many concurrent writes for {} on table at '{}', please retry later: {}",
1442 operation, table_uri, detail
1443 ),
1444 },
1445 lance_core::Error::TooMuchWriteContention { .. }
1449 | lance_core::Error::RetryableCommitConflict { .. }
1450 | lance_core::Error::IncompatibleTransaction { .. }
1451 | lance_core::Error::VersionConflict { .. } => NamespaceError::ConcurrentModification {
1452 message: format!(
1453 "Concurrent modification detected for {} on table at '{}': {}",
1454 operation, table_uri, detail
1455 ),
1456 },
1457 lance_core::Error::NotSupported { .. } => NamespaceError::Unsupported {
1458 message: format!(
1459 "{} is not supported on table at '{}': {}",
1460 operation, table_uri, detail
1461 ),
1462 },
1463 _ => NamespaceError::Internal {
1464 message: format!(
1465 "Failed to run {} on table at '{}': {}",
1466 operation, table_uri, detail
1467 ),
1468 },
1469 };
1470 ns_err.into()
1471 }
1472
1473 async fn table_has_actual_manifests(&self, table_name: &str) -> Result<bool> {
1474 manifest::ManifestNamespace::path_has_actual_manifests(
1475 &self.object_store,
1476 &self.table_path(table_name),
1477 )
1478 .await
1479 }
1480
1481 async fn filter_declared_tables(
1482 &self,
1483 tables: Vec<String>,
1484 include_declared: bool,
1485 ) -> Result<Vec<String>> {
1486 if include_declared {
1487 return Ok(tables);
1488 }
1489
1490 let mut stream = futures::stream::iter(tables.into_iter().map(|table_name| async move {
1491 if self.table_has_actual_manifests(&table_name).await? {
1495 Ok::<Option<String>, Error>(Some(table_name))
1496 } else {
1497 Ok::<Option<String>, Error>(None)
1498 }
1499 }))
1500 .buffered(manifest::DECLARED_FILTER_CONCURRENCY);
1501
1502 let mut filtered = Vec::new();
1503 while let Some(result) = stream.next().await {
1504 if let Some(table_name) = result? {
1505 filtered.push(table_name);
1506 }
1507 }
1508 Ok(filtered)
1509 }
1510
1511 fn ipc_reader_from_request_data(
1512 request_data: &Bytes,
1513 operation: &str,
1514 ) -> Result<(
1515 Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1516 usize,
1517 )> {
1518 if request_data.is_empty() {
1519 return Err(NamespaceError::InvalidInput {
1520 message: format!(
1521 "Request data (Arrow IPC stream) is required for {}",
1522 operation
1523 ),
1524 }
1525 .into());
1526 }
1527
1528 let cursor = Cursor::new(request_data.as_ref());
1529 let stream_reader =
1530 StreamReader::try_new(cursor, None).map_err(|e| NamespaceError::InvalidInput {
1531 message: format!("Invalid Arrow IPC stream: {}", e),
1532 })?;
1533 let arrow_schema = stream_reader.schema();
1534
1535 let mut num_rows = 0usize;
1536 let mut batches = Vec::new();
1537 for batch_result in stream_reader {
1538 let batch = batch_result.map_err(|e| NamespaceError::Internal {
1539 message: format!("Failed to read batch from IPC stream: {}", e),
1540 })?;
1541 num_rows += batch.num_rows();
1542 batches.push(batch);
1543 }
1544
1545 let reader: Box<dyn arrow::record_batch::RecordBatchReader + Send> = if batches.is_empty() {
1546 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
1547 Box::new(RecordBatchIterator::new(vec![Ok(batch)], arrow_schema))
1548 } else {
1549 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
1550 Box::new(RecordBatchIterator::new(batch_results, arrow_schema))
1551 };
1552
1553 Ok((reader, num_rows))
1554 }
1555
1556 async fn table_uri_has_actual_manifests(&self, table_uri: &str) -> Result<bool> {
1557 let table_path = self.object_store_path_from_uri(table_uri)?;
1558 manifest::ManifestNamespace::path_has_actual_manifests(&self.object_store, &table_path)
1559 .await
1560 }
1561
1562 fn object_store_path_from_uri(&self, uri: &str) -> Result<Path> {
1563 let registry = self
1564 .session
1565 .as_ref()
1566 .map(|session| session.store_registry())
1567 .unwrap_or_else(|| Arc::new(ObjectStoreRegistry::default()));
1568 ObjectStore::extract_path_from_uri(registry, uri)
1569 }
1570
1571 fn normalized_branch(branch: Option<&str>) -> Result<Option<&str>> {
1576 match branch.filter(|b| !b.is_empty() && *b != "main") {
1577 Some(branch) => {
1578 check_valid_branch(branch).map_err(|e| {
1579 lance_core::Error::from(NamespaceError::InvalidInput {
1580 message: format!("invalid branch name '{}': {}", branch, e),
1581 })
1582 })?;
1583 Ok(Some(branch))
1584 }
1585 None => Ok(None),
1586 }
1587 }
1588
1589 async fn open_validated_branch(&self, table_uri: &str, branch: &str) -> Result<Dataset> {
1590 let dataset = self
1591 .configured_builder(table_uri)
1592 .with_branch(branch, None)
1593 .load()
1594 .await
1595 .map_err(|e| {
1596 let message = format!(
1597 "branch '{}' not found for table at '{}': {}",
1598 branch, table_uri, e
1599 );
1600 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1601 })?;
1602 dataset.branches().get(branch).await.map_err(|e| {
1603 Self::map_open_error(
1604 e,
1605 NamespaceError::TableNotFound {
1606 message: format!("branch '{}' not found for table at '{}'", branch, table_uri),
1607 },
1608 )
1609 })?;
1610 Ok(dataset)
1611 }
1612
1613 async fn resolve_branch_location(&self, table_uri: &str, branch: &str) -> Result<String> {
1614 Ok(self
1615 .open_validated_branch(table_uri, branch)
1616 .await?
1617 .branch_location()
1618 .uri)
1619 }
1620
1621 async fn resolve_branch_for_commit(
1631 &self,
1632 table_uri: &str,
1633 branch: &str,
1634 ) -> Result<(String, Path, Option<u64>)> {
1635 let main = self
1636 .configured_builder(table_uri)
1637 .load()
1638 .await
1639 .map_err(|e| {
1640 let message = format!("table at '{}' not found: {}", table_uri, e);
1641 Self::map_open_error(e, NamespaceError::TableNotFound { message })
1642 })?;
1643 let branch_location = main.branch_location().find_branch(Some(branch))?;
1644 match main.branches().get(branch).await {
1645 Ok(contents) => Ok((
1646 branch_location.uri,
1647 branch_location.path,
1648 Some(contents.parent_version),
1649 )),
1650 Err(lance_core::Error::RefNotFound { .. }) => {
1651 if self
1652 .branch_has_committed_versions(&branch_location.path)
1653 .await?
1654 {
1655 return Err(NamespaceError::TableNotFound {
1656 message: format!(
1657 "branch '{}' not found for table at '{}'",
1658 branch, table_uri
1659 ),
1660 }
1661 .into());
1662 }
1663 Ok((branch_location.uri, branch_location.path, None))
1664 }
1665 Err(e) => Err(e),
1666 }
1667 }
1668
1669 async fn branch_has_committed_versions(&self, branch_path: &Path) -> Result<bool> {
1670 Ok(!self
1671 .list_versions_under(branch_path, false, Some(1))
1672 .await?
1673 .is_empty())
1674 }
1675
1676 fn validate_dir_only_properties(
1677 properties: Option<&HashMap<String, String>>,
1678 operation: &str,
1679 ) -> Result<()> {
1680 if properties.is_some_and(|properties| !properties.is_empty()) {
1684 return Err(NamespaceError::Unsupported {
1685 message: format!(
1686 "{} with non-empty table properties requires manifest_enabled=true",
1687 operation
1688 ),
1689 }
1690 .into());
1691 }
1692 Ok(())
1693 }
1694
1695 async fn write_reader_to_table(
1696 &self,
1697 table_uri: &str,
1698 reader: Box<dyn arrow::record_batch::RecordBatchReader + Send>,
1699 mode: WriteMode,
1700 extra_storage_options: Option<HashMap<String, String>>,
1701 ) -> Result<Dataset> {
1702 let mut merged_storage_options = self.storage_options.clone().unwrap_or_default();
1705 if let Some(extra_storage_options) = extra_storage_options {
1706 merged_storage_options.extend(extra_storage_options);
1707 }
1708 let store_params = (!merged_storage_options.is_empty()).then(|| ObjectStoreParams {
1709 storage_options_accessor: Some(Arc::new(
1710 lance_io::object_store::StorageOptionsAccessor::with_static_options(
1711 merged_storage_options,
1712 ),
1713 )),
1714 ..Default::default()
1715 });
1716
1717 let write_params = WriteParams {
1718 mode,
1719 store_params,
1720 session: self.session.clone(),
1721 ..Default::default()
1722 };
1723
1724 let dataset = Dataset::write(reader, table_uri, Some(write_params))
1725 .await
1726 .map_err(|e| NamespaceError::Internal {
1727 message: format!("Failed to write table at '{}': {}", table_uri, e),
1728 })?;
1729
1730 Ok(dataset)
1731 }
1732
1733 fn manifest_version_from_filename(filename: &str) -> Option<u64> {
1737 ManifestNamingScheme::detect_scheme(filename)?.parse_version(filename)
1738 }
1739
1740 fn create_table_version_response(
1742 version: u64,
1743 final_path: &Path,
1744 final_meta: &ObjectMeta,
1745 ) -> CreateTableVersionResponse {
1746 CreateTableVersionResponse {
1747 transaction_id: None,
1748 version: Some(Box::new(TableVersion {
1749 version: version as i64,
1750 manifest_path: final_path.to_string(),
1751 manifest_size: Some(final_meta.size as i64),
1752 e_tag: final_meta.e_tag.clone(),
1753 timestamp_millis: None,
1754 metadata: None,
1755 })),
1756 ..Default::default()
1757 }
1758 }
1759
1760 async fn staging_matches_final_manifest(
1768 &self,
1769 staging_path: &Path,
1770 final_path: &Path,
1771 final_meta: &ObjectMeta,
1772 request_manifest_size: Option<i64>,
1773 ) -> Result<bool> {
1774 if let Some(size) = request_manifest_size
1775 && size != final_meta.size as i64
1776 {
1777 return Ok(false);
1778 }
1779
1780 let staging_bytes = match self.object_store.inner.get(staging_path).await {
1781 Ok(r) => r.bytes().await.map_err(|e| {
1782 lance_core::Error::from(NamespaceError::Internal {
1783 message: format!(
1784 "Failed to read staging manifest at '{}': {}",
1785 staging_path, e
1786 ),
1787 })
1788 })?,
1789 Err(ObjectStoreError::NotFound { .. }) => return Ok(false),
1790 Err(e) => {
1791 return Err(lance_core::Error::from(NamespaceError::Internal {
1792 message: format!(
1793 "Failed to read staging manifest at '{}': {}",
1794 staging_path, e
1795 ),
1796 }));
1797 }
1798 };
1799
1800 let final_bytes = self
1801 .object_store
1802 .inner
1803 .get(final_path)
1804 .await
1805 .map_err(|e| {
1806 lance_core::Error::from(NamespaceError::Internal {
1807 message: format!(
1808 "Failed to read existing version manifest at '{}': {}",
1809 final_path, e
1810 ),
1811 })
1812 })?
1813 .bytes()
1814 .await
1815 .map_err(|e| {
1816 lance_core::Error::from(NamespaceError::Internal {
1817 message: format!(
1818 "Failed to read existing version manifest bytes at '{}': {}",
1819 final_path, e
1820 ),
1821 })
1822 })?;
1823
1824 Ok(staging_bytes.as_ref() == final_bytes.as_ref())
1825 }
1826
1827 async fn resolve_existing_table_version(
1829 &self,
1830 args: ExistingTableVersionResolve<'_>,
1831 ) -> Result<CreateTableVersionResponse> {
1832 if self
1833 .staging_matches_final_manifest(
1834 args.staging_path,
1835 args.final_path,
1836 args.final_meta,
1837 args.request_manifest_size,
1838 )
1839 .await?
1840 {
1841 if let Err(e) = self.object_store.inner.delete(args.staging_path).await {
1843 log::warn!(
1844 "Failed to delete staging manifest at '{}': {:?}",
1845 args.staging_path,
1846 e
1847 );
1848 }
1849 return Ok(Self::create_table_version_response(
1850 args.version,
1851 args.final_path,
1852 args.final_meta,
1853 ));
1854 }
1855
1856 Err(lance_core::Error::from(
1857 NamespaceError::ConcurrentModification {
1858 message: format!(
1859 "Version {} already exists for table at '{}' with different content",
1860 args.version, args.table_uri
1861 ),
1862 },
1863 ))
1864 }
1865
1866 async fn enforce_create_table_version_cas(
1875 &self,
1876 table_path: &Path,
1877 version: u64,
1878 table_uri: &str,
1879 is_branch: bool,
1880 branch_parent_version: Option<u64>,
1881 ) -> Result<()> {
1882 let latest = self.list_versions_under(table_path, true, Some(1)).await?;
1883 let expected = match latest.first() {
1884 Some(v) => (v.version as u64).checked_add(1).ok_or_else(|| {
1885 lance_core::Error::from(NamespaceError::ConcurrentModification {
1886 message: format!(
1887 "Version overflow computing next version for table at '{}': \
1888 latest version {} cannot advance",
1889 table_uri, v.version
1890 ),
1891 })
1892 })?,
1893 None => {
1894 if is_branch {
1895 match branch_parent_version {
1898 Some(parent_version) => parent_version,
1899 None => version,
1900 }
1901 } else {
1902 1
1903 }
1904 }
1905 };
1906 if version != expected {
1907 let latest_display = latest
1908 .first()
1909 .map(|v| v.version.to_string())
1910 .unwrap_or_else(|| "none".to_string());
1911 return Err(lance_core::Error::from(
1912 NamespaceError::ConcurrentModification {
1913 message: format!(
1914 "Version CAS failed for table at '{}': requested {}, expected {} (latest {})",
1915 table_uri, version, expected, latest_display
1916 ),
1917 },
1918 ));
1919 }
1920 Ok(())
1921 }
1922
1923 async fn materialize_version_manifest_create(
1925 &self,
1926 staging_path: &Path,
1927 final_path: &Path,
1928 staging_manifest_path: &str,
1929 ) -> std::result::Result<(), ObjectStoreError> {
1930 match self
1931 .object_store
1932 .inner
1933 .copy_if_not_exists(staging_path, final_path)
1934 .await
1935 {
1936 Ok(()) => Ok(()),
1937 Err(ObjectStoreError::NotImplemented { .. })
1938 | Err(ObjectStoreError::NotSupported { .. }) => {
1939 let manifest_data = self
1940 .object_store
1941 .inner
1942 .get(staging_path)
1943 .await?
1944 .bytes()
1945 .await
1946 .map_err(|e| ObjectStoreError::Generic {
1947 store: "DirectoryNamespace",
1948 source: Box::new(std::io::Error::other(format!(
1949 "Failed to read staging manifest bytes at '{}': {}",
1950 staging_manifest_path, e
1951 ))),
1952 })?;
1953 self.object_store
1954 .inner
1955 .put_opts(
1956 final_path,
1957 manifest_data.into(),
1958 PutOptions {
1959 mode: PutMode::Create,
1960 ..Default::default()
1961 },
1962 )
1963 .await
1964 .map(|_| ())
1965 }
1966 Err(e) => Err(e),
1967 }
1968 }
1969
1970 async fn list_table_versions_from_storage(
1971 &self,
1972 table_uri: &str,
1973 descending: bool,
1974 limit: Option<i32>,
1975 ) -> Result<Vec<TableVersion>> {
1976 let table_path = self.object_store_path_from_uri(table_uri)?;
1977 self.list_versions_under(&table_path, descending, limit)
1978 .await
1979 }
1980
1981 async fn list_versions_under(
1985 &self,
1986 table_path: &Path,
1987 descending: bool,
1988 limit: Option<i32>,
1989 ) -> Result<Vec<TableVersion>> {
1990 let versions_dir = table_path.clone().join(VERSIONS_DIR);
1991 let mut stream = self.object_store.read_dir_all(&versions_dir, None);
1992 let list_err = |e: lance_core::Error| {
1993 lance_core::Error::from(NamespaceError::Internal {
1994 message: format!(
1995 "Failed to list manifest files under '{}': {}",
1996 versions_dir, e
1997 ),
1998 })
1999 };
2000
2001 let limit = limit
2002 .filter(|limit| *limit >= 0)
2003 .map(|limit| limit as usize);
2004
2005 let mut table_versions: Vec<TableVersion> = Vec::new();
2006 let push_meta = |meta: ObjectMeta, out: &mut Vec<TableVersion>| -> bool {
2007 let Some(filename) = meta.location.filename() else {
2008 return false;
2009 };
2010 let Some(actual_version) = Self::manifest_version_from_filename(filename) else {
2011 return false;
2012 };
2013 out.push(TableVersion {
2014 version: actual_version as i64,
2015 manifest_path: meta.location.to_string(),
2016 manifest_size: Some(meta.size as i64),
2017 e_tag: meta.e_tag,
2018 timestamp_millis: Some(meta.last_modified.timestamp_millis()),
2019 metadata: None,
2020 });
2021 true
2022 };
2023
2024 let mut first_manifest_filename_len = None;
2029 while first_manifest_filename_len.is_none() {
2030 match stream.try_next().await.map_err(list_err)? {
2031 Some(meta) => {
2032 let filename_len = meta.location.filename().map(|f| f.len());
2033 if push_meta(meta, &mut table_versions) {
2034 first_manifest_filename_len = filename_len;
2035 }
2036 }
2037 None => break,
2038 }
2039 }
2040 let is_v2_naming = first_manifest_filename_len == Some(29);
2041
2042 let list_is_ordered = self.object_store.list_is_lexically_ordered;
2047 let stream_matches_request = list_is_ordered
2048 && if is_v2_naming {
2049 descending
2050 } else {
2051 !descending
2052 };
2053 let early_stop_at = limit.filter(|_| stream_matches_request);
2054
2055 while early_stop_at.is_none_or(|n| table_versions.len() < n) {
2056 match stream.try_next().await.map_err(list_err)? {
2057 Some(meta) => {
2058 push_meta(meta, &mut table_versions);
2059 }
2060 None => break,
2061 }
2062 }
2063
2064 if let Some(n) = early_stop_at {
2067 table_versions.truncate(n);
2068 } else {
2069 if descending {
2070 table_versions.sort_by_key(|v| std::cmp::Reverse(v.version));
2071 } else {
2072 table_versions.sort_by_key(|v| v.version);
2073 }
2074 if let Some(limit) = limit {
2075 table_versions.truncate(limit);
2076 }
2077 }
2078
2079 Ok(table_versions)
2080 }
2081
2082 async fn describe_table_impl(
2086 &self,
2087 request: DescribeTableRequest,
2088 ) -> Result<DescribeTableResponse> {
2089 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
2090 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
2091 let skip_manifest_for_root = self.dir_listing_enabled
2092 && is_root_level
2093 && !self.dir_listing_to_manifest_migration_enabled;
2094 if is_child_table
2102 || !self.dir_listing_enabled
2103 || self.dir_listing_to_manifest_migration_enabled
2104 {
2105 self.ensure_read_manifest().await?;
2106 }
2107 if let Some(manifest_ns) = self.manifest_ns_for_read()
2108 && !skip_manifest_for_root
2109 {
2110 match manifest_ns.describe_table(request.clone()).await {
2111 Ok(mut response) => {
2112 if let Some(ref table_uri) = response.table_uri {
2113 let vend = request.vend_credentials.unwrap_or(true);
2115 let identity = request.identity.as_deref();
2116 response.storage_options = self
2117 .get_storage_options_for_table(table_uri, vend, identity)
2118 .await?;
2119 }
2120 if self.table_version_tracking_enabled {
2122 response.managed_versioning = Some(true);
2123 }
2124 return Ok(response);
2125 }
2126 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
2127 return Err(e);
2130 }
2131 Err(e) if self.dir_listing_enabled && is_root_level => {
2132 if !Self::is_manifest_table_absent_error(&e) {
2136 return Err(Self::classify_storage_error(e));
2137 }
2138 }
2139 Err(e) => return Err(e),
2140 }
2141 }
2142 if is_child_table {
2143 return Err(self.child_namespace_requires_manifest_error());
2144 }
2145
2146 let table_name = Self::table_name_from_id(&request.id)?;
2147 let table_id = Self::format_table_id_from_request(&request.id);
2148 if !self.dir_listing_enabled {
2149 return Err(NamespaceError::TableNotFound { message: table_id }.into());
2150 }
2151
2152 let table_uri = self.table_full_uri(&table_name);
2153
2154 let status = self.check_table_status(&table_name).await?;
2156
2157 if !status.exists {
2158 return Err(NamespaceError::TableNotFound {
2159 message: table_id.clone(),
2160 }
2161 .into());
2162 }
2163
2164 if status.is_deregistered {
2165 return Err(NamespaceError::TableNotFound {
2166 message: format!("Table is deregistered: {}", table_id),
2167 }
2168 .into());
2169 }
2170
2171 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
2172 let should_check_declared =
2173 load_detailed_metadata || request.check_declared.unwrap_or(false);
2174 let vend_credentials = request.vend_credentials.unwrap_or(true);
2176 let identity = request.identity.as_deref();
2177 let is_only_declared = if should_check_declared {
2178 if status.has_reserved_file {
2179 Some(!self.table_has_actual_manifests(&table_name).await?)
2180 } else {
2181 Some(false)
2182 }
2183 } else {
2184 None
2185 };
2186
2187 if !load_detailed_metadata {
2188 let storage_options = self
2189 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2190 .await?;
2191 return Ok(DescribeTableResponse {
2192 table: Some(table_name),
2193 namespace: request.id.as_ref().map(|id| {
2194 if id.len() > 1 {
2195 id[..id.len() - 1].to_vec()
2196 } else {
2197 vec![]
2198 }
2199 }),
2200 location: Some(table_uri.clone()),
2201 table_uri: Some(table_uri),
2202 storage_options,
2203 is_only_declared,
2204 managed_versioning: if self.table_version_tracking_enabled {
2205 Some(true)
2206 } else {
2207 None
2208 },
2209 ..Default::default()
2210 });
2211 }
2212
2213 if is_only_declared == Some(true) {
2214 let storage_options = self
2215 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2216 .await?;
2217 return Ok(DescribeTableResponse {
2218 table: Some(table_name),
2219 namespace: request.id.as_ref().map(|id| {
2220 if id.len() > 1 {
2221 id[..id.len() - 1].to_vec()
2222 } else {
2223 vec![]
2224 }
2225 }),
2226 location: Some(table_uri.clone()),
2227 table_uri: Some(table_uri),
2228 storage_options,
2229 is_only_declared,
2230 managed_versioning: if self.table_version_tracking_enabled {
2231 Some(true)
2232 } else {
2233 None
2234 },
2235 ..Default::default()
2236 });
2237 }
2238
2239 let mut builder = DatasetBuilder::from_uri(&table_uri);
2242 if let Some(opts) = &self.storage_options {
2243 builder = builder.with_storage_options(opts.clone());
2244 }
2245 if let Some(sess) = &self.session {
2246 builder = builder.with_session(sess.clone());
2247 }
2248 match builder.load().await {
2249 Ok(mut dataset) => {
2250 if let Some(requested_version) = request.version {
2252 dataset = dataset
2253 .checkout_version(requested_version as u64)
2254 .await
2255 .map_err(|e| {
2256 let message = format!(
2257 "Version {} not found for table '{}': {}",
2258 requested_version, table_name, e
2259 );
2260 Self::map_open_error(
2261 e,
2262 NamespaceError::TableVersionNotFound { message },
2263 )
2264 })?;
2265 }
2266
2267 let version_info = dataset.version();
2268 let lance_schema = dataset.schema();
2269 let arrow_schema: arrow_schema::Schema = lance_schema.into();
2270 let json_schema = arrow_schema_to_json(&arrow_schema)?;
2271 let storage_options = self
2272 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2273 .await?;
2274
2275 let metadata: std::collections::HashMap<String, String> =
2277 version_info.metadata.into_iter().collect();
2278
2279 Ok(DescribeTableResponse {
2280 table: Some(table_name),
2281 namespace: request.id.as_ref().map(|id| {
2282 if id.len() > 1 {
2283 id[..id.len() - 1].to_vec()
2284 } else {
2285 vec![]
2286 }
2287 }),
2288 version: Some(version_info.version as i64),
2289 location: Some(table_uri.clone()),
2290 table_uri: Some(table_uri),
2291 schema: Some(Box::new(json_schema)),
2292 storage_options,
2293 metadata: Some(metadata),
2294 is_only_declared,
2295 managed_versioning: if self.table_version_tracking_enabled {
2296 Some(true)
2297 } else {
2298 None
2299 },
2300 ..Default::default()
2301 })
2302 }
2303 Err(err) => {
2304 if manifest::ManifestNamespace::is_not_found_load_error(&err)
2305 && is_only_declared == Some(true)
2306 {
2307 let storage_options = self
2308 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
2309 .await?;
2310 Ok(DescribeTableResponse {
2311 table: Some(table_name),
2312 namespace: request.id.as_ref().map(|id| {
2313 if id.len() > 1 {
2314 id[..id.len() - 1].to_vec()
2315 } else {
2316 vec![]
2317 }
2318 }),
2319 location: Some(table_uri.clone()),
2320 table_uri: Some(table_uri),
2321 storage_options,
2322 is_only_declared,
2323 managed_versioning: if self.table_version_tracking_enabled {
2324 Some(true)
2325 } else {
2326 None
2327 },
2328 ..Default::default()
2329 })
2330 } else {
2331 Err(NamespaceError::Internal {
2332 message: format!(
2333 "Table directory exists but cannot load dataset {}: {:?}",
2334 table_name, err
2335 ),
2336 }
2337 .into())
2338 }
2339 }
2340 }
2341 }
2342
2343 fn configured_builder(&self, table_uri: &str) -> DatasetBuilder {
2346 let mut builder = DatasetBuilder::from_uri(table_uri);
2347 if let Some(opts) = &self.storage_options {
2348 builder = builder.with_storage_options(opts.clone());
2349 }
2350 if let Some(sess) = &self.session {
2351 builder = builder.with_session(sess.clone());
2352 }
2353 builder
2354 }
2355
2356 async fn load_dataset(
2357 &self,
2358 table_uri: &str,
2359 version: Option<i64>,
2360 operation: &str,
2361 ) -> Result<Dataset> {
2362 if let Some(version) = version
2363 && version < 0
2364 {
2365 return Err(NamespaceError::InvalidInput {
2366 message: format!(
2367 "Table version for {} must be non-negative, got {}",
2368 operation, version
2369 ),
2370 }
2371 .into());
2372 }
2373
2374 let builder = self.configured_builder(table_uri);
2375
2376 let dataset = builder.load().await.map_err(|e| {
2377 let message = format!(
2378 "Failed to open table at '{}' for {}: {}",
2379 table_uri, operation, e
2380 );
2381 Self::map_open_error(e, NamespaceError::TableNotFound { message })
2382 })?;
2383
2384 if let Some(version) = version {
2385 return dataset.checkout_version(version as u64).await.map_err(|e| {
2386 let message = format!(
2387 "Failed to checkout version {} for table at '{}' during {}: {}",
2388 version, table_uri, operation, e
2389 );
2390 Self::map_open_error(e, NamespaceError::TableVersionNotFound { message })
2391 });
2392 }
2393
2394 Ok(dataset)
2395 }
2396
2397 fn parse_index_type(index_type: &str) -> Result<IndexType> {
2398 match index_type.trim().to_ascii_uppercase().as_str() {
2399 "SCALAR" | "BTREE" => Ok(IndexType::BTree),
2400 "BITMAP" => Ok(IndexType::Bitmap),
2401 "LABEL_LIST" | "LABELLIST" => Ok(IndexType::LabelList),
2402 "INVERTED" | "FTS" => Ok(IndexType::Inverted),
2403 "NGRAM" => Ok(IndexType::NGram),
2404 "ZONEMAP" | "ZONE_MAP" => Ok(IndexType::ZoneMap),
2405 "BLOOMFILTER" | "BLOOM_FILTER" => Ok(IndexType::BloomFilter),
2406 "RTREE" | "R_TREE" => Ok(IndexType::RTree),
2407 "VECTOR" | "IVF_PQ" => Ok(IndexType::IvfPq),
2408 "IVF_FLAT" => Ok(IndexType::IvfFlat),
2409 "IVF_SQ" => Ok(IndexType::IvfSq),
2410 "IVF_RQ" => Ok(IndexType::IvfRq),
2411 "IVF_HNSW_FLAT" => Ok(IndexType::IvfHnswFlat),
2412 "IVF_HNSW_SQ" => Ok(IndexType::IvfHnswSq),
2413 "IVF_HNSW_PQ" => Ok(IndexType::IvfHnswPq),
2414 other => Err(NamespaceError::InvalidInput {
2415 message: format!("Unsupported index_type '{}'", other),
2416 }
2417 .into()),
2418 }
2419 }
2420
2421 fn parse_metric_type(distance_type: Option<&str>) -> Result<MetricType> {
2422 let distance_type = distance_type.unwrap_or("l2");
2423 MetricType::try_from(distance_type).map_err(|e| {
2424 lance_core::Error::from(NamespaceError::InvalidInput {
2425 message: format!(
2426 "Unsupported distance_type '{}' for vector index: {}",
2427 distance_type, e
2428 ),
2429 })
2430 })
2431 }
2432
2433 fn build_index_params(request: &CreateTableIndexRequest) -> Result<DirectoryIndexParams> {
2434 let index_type = Self::parse_index_type(&request.index_type)?;
2435 Ok(match index_type {
2436 IndexType::BTree => DirectoryIndexParams::Scalar {
2437 index_type,
2438 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
2439 },
2440 IndexType::Bitmap => DirectoryIndexParams::Scalar {
2441 index_type,
2442 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
2443 },
2444 IndexType::LabelList => DirectoryIndexParams::Scalar {
2445 index_type,
2446 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
2447 },
2448 IndexType::NGram => DirectoryIndexParams::Scalar {
2449 index_type,
2450 params: ScalarIndexParams::for_builtin(BuiltinIndexType::NGram),
2451 },
2452 IndexType::ZoneMap => DirectoryIndexParams::Scalar {
2453 index_type,
2454 params: ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap),
2455 },
2456 IndexType::BloomFilter => DirectoryIndexParams::Scalar {
2457 index_type,
2458 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BloomFilter),
2459 },
2460 IndexType::RTree => DirectoryIndexParams::Scalar {
2461 index_type,
2462 params: ScalarIndexParams::for_builtin(BuiltinIndexType::RTree),
2463 },
2464 IndexType::Inverted => {
2465 let mut params = InvertedIndexParams::default();
2466 if let Some(with_position) = request.with_position {
2467 params = params.with_position(with_position);
2468 }
2469 if let Some(base_tokenizer) = &request.base_tokenizer {
2470 params = params.base_tokenizer(base_tokenizer.clone());
2471 }
2472 if let Some(language) = &request.language {
2473 params = params.language(language)?;
2474 }
2475 if let Some(max_token_length) = request.max_token_length {
2476 if max_token_length < 0 {
2477 return Err(NamespaceError::InvalidInput {
2478 message: format!(
2479 "FTS max_token_length must be non-negative, got {}",
2480 max_token_length
2481 ),
2482 }
2483 .into());
2484 }
2485 params = params.max_token_length(Some(max_token_length as usize));
2486 }
2487 if let Some(lower_case) = request.lower_case {
2488 params = params.lower_case(lower_case);
2489 }
2490 if let Some(stem) = request.stem {
2491 params = params.stem(stem);
2492 }
2493 if let Some(remove_stop_words) = request.remove_stop_words {
2494 params = params.remove_stop_words(remove_stop_words);
2495 }
2496 if let Some(ascii_folding) = request.ascii_folding {
2497 params = params.ascii_folding(ascii_folding);
2498 }
2499 DirectoryIndexParams::Inverted(params)
2500 }
2501 IndexType::IvfFlat => DirectoryIndexParams::Vector {
2502 index_type,
2503 params: VectorIndexParams::with_ivf_flat_params(
2504 Self::parse_metric_type(request.distance_type.as_deref())?,
2505 IvfBuildParams::default(),
2506 ),
2507 },
2508 IndexType::IvfPq => DirectoryIndexParams::Vector {
2509 index_type,
2510 params: VectorIndexParams::with_ivf_pq_params(
2511 Self::parse_metric_type(request.distance_type.as_deref())?,
2512 IvfBuildParams::default(),
2513 PQBuildParams::default(),
2514 ),
2515 },
2516 IndexType::IvfSq => DirectoryIndexParams::Vector {
2517 index_type,
2518 params: VectorIndexParams::with_ivf_sq_params(
2519 Self::parse_metric_type(request.distance_type.as_deref())?,
2520 IvfBuildParams::default(),
2521 SQBuildParams::default(),
2522 ),
2523 },
2524 IndexType::IvfRq => {
2525 let rq_params = if let Some(requested_num_bits) = request.num_bits {
2526 let invalid_num_bits = || NamespaceError::InvalidInput {
2527 message: format!(
2528 "IVF_RQ num_bits must be in {}..={}, got {}",
2529 RABIT_MIN_NUM_BITS, RABIT_MAX_NUM_BITS, requested_num_bits
2530 ),
2531 };
2532 let num_bits =
2533 u8::try_from(requested_num_bits).map_err(|_| invalid_num_bits())?;
2534 validate_supported_rq_num_bits(num_bits).map_err(|_| invalid_num_bits())?;
2535 RQBuildParams::new(num_bits)
2536 } else {
2537 RQBuildParams::default()
2538 };
2539 DirectoryIndexParams::Vector {
2540 index_type,
2541 params: VectorIndexParams::with_ivf_rq_params(
2542 Self::parse_metric_type(request.distance_type.as_deref())?,
2543 IvfBuildParams::default(),
2544 rq_params,
2545 ),
2546 }
2547 }
2548 IndexType::IvfHnswFlat => DirectoryIndexParams::Vector {
2549 index_type,
2550 params: VectorIndexParams::ivf_hnsw(
2551 Self::parse_metric_type(request.distance_type.as_deref())?,
2552 IvfBuildParams::default(),
2553 HnswBuildParams::default(),
2554 ),
2555 },
2556 IndexType::IvfHnswSq => DirectoryIndexParams::Vector {
2557 index_type,
2558 params: VectorIndexParams::with_ivf_hnsw_sq_params(
2559 Self::parse_metric_type(request.distance_type.as_deref())?,
2560 IvfBuildParams::default(),
2561 HnswBuildParams::default(),
2562 SQBuildParams::default(),
2563 ),
2564 },
2565 IndexType::IvfHnswPq => DirectoryIndexParams::Vector {
2566 index_type,
2567 params: VectorIndexParams::with_ivf_hnsw_pq_params(
2568 Self::parse_metric_type(request.distance_type.as_deref())?,
2569 IvfBuildParams::default(),
2570 HnswBuildParams::default(),
2571 PQBuildParams::default(),
2572 ),
2573 },
2574 other => {
2575 return Err(NamespaceError::InvalidInput {
2576 message: format!("Unsupported index type for namespace API: {}", other),
2577 }
2578 .into());
2579 }
2580 })
2581 }
2582
2583 fn paginate_indices(
2584 indices: &mut Vec<IndexContent>,
2585 page_token: Option<String>,
2586 limit: Option<i32>,
2587 ) -> Option<String> {
2588 indices.sort_by(|a, b| a.index_name.cmp(&b.index_name));
2589
2590 if let Some(start_after) = page_token {
2591 if let Some(index) = indices
2592 .iter()
2593 .position(|index| index.index_name.as_str() > start_after.as_str())
2594 {
2595 indices.drain(0..index);
2596 } else {
2597 indices.clear();
2598 }
2599 }
2600
2601 let mut next_page_token = None;
2602 if let Some(limit) = limit
2603 && limit >= 0
2604 {
2605 let limit = limit as usize;
2606 if limit > 0 && indices.len() > limit {
2607 next_page_token = Some(indices[limit - 1].index_name.clone());
2608 }
2609 indices.truncate(limit);
2610 }
2611 if indices.is_empty() {
2612 None
2613 } else {
2614 next_page_token
2615 }
2616 }
2617
2618 fn transaction_operation_name(transaction: &Transaction) -> String {
2619 match &transaction.operation {
2620 Operation::CreateIndex {
2621 new_indices,
2622 removed_indices,
2623 ..
2624 } if new_indices.is_empty() && !removed_indices.is_empty() => "DropIndex".to_string(),
2625 _ => transaction.operation.to_string(),
2626 }
2627 }
2628
2629 fn transaction_response(
2630 version: u64,
2631 transaction: &Transaction,
2632 alteration: Option<TransactionAlteration>,
2633 ) -> DescribeTransactionResponse {
2634 let mut properties = transaction
2635 .transaction_properties
2636 .as_ref()
2637 .map(|properties| (**properties).clone())
2638 .unwrap_or_default();
2639
2640 let mut effective_status = "SUCCEEDED".to_string();
2643 if let Some(alteration) = alteration {
2644 for key in &alteration.removed_properties {
2645 properties.remove(key);
2646 }
2647 for (key, value) in alteration.properties {
2648 properties.insert(key, value);
2649 }
2650 if let Some(status) = alteration.status {
2651 effective_status = status;
2652 }
2653 }
2654
2655 properties.insert("uuid".to_string(), transaction.uuid.clone());
2656 properties.insert("version".to_string(), version.to_string());
2657 properties.insert(
2658 "read_version".to_string(),
2659 transaction.read_version.to_string(),
2660 );
2661 properties.insert(
2662 "operation".to_string(),
2663 Self::transaction_operation_name(transaction),
2664 );
2665 if let Some(tag) = &transaction.tag {
2666 properties.insert("tag".to_string(), tag.clone());
2667 }
2668
2669 DescribeTransactionResponse {
2670 status: effective_status,
2671 properties: Some(properties),
2672 ..Default::default()
2673 }
2674 }
2675
2676 fn describe_table_index_stats_response(
2677 stats: &serde_json::Value,
2678 ) -> DescribeTableIndexStatsResponse {
2679 let get_i64 = |key: &str| {
2680 stats.get(key).and_then(|value| {
2681 value
2682 .as_i64()
2683 .or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
2684 })
2685 };
2686
2687 DescribeTableIndexStatsResponse {
2688 distance_type: stats
2689 .get("distance_type")
2690 .and_then(|value| value.as_str())
2691 .map(str::to_string),
2692 index_type: stats
2693 .get("index_type")
2694 .and_then(|value| value.as_str())
2695 .map(str::to_string),
2696 num_indexed_rows: get_i64("num_indexed_rows"),
2697 num_unindexed_rows: get_i64("num_unindexed_rows"),
2698 num_indices: get_i64("num_indices").and_then(|value| i32::try_from(value).ok()),
2699 ..Default::default()
2700 }
2701 }
2702
2703 async fn find_transaction(&self, dataset: &Dataset, id: &str) -> Result<(u64, Transaction)> {
2708 if let Ok(version) = id.parse::<u64>() {
2709 let transaction = dataset
2710 .read_transaction_by_version(version)
2711 .await
2712 .map_err(|e| {
2713 lance_core::Error::from(NamespaceError::TransactionNotFound {
2714 message: format!(
2715 "Failed to read transaction for version {}: {}",
2716 version, e
2717 ),
2718 })
2719 })?
2720 .ok_or_else(|| {
2721 lance_core::Error::from(NamespaceError::TransactionNotFound {
2722 message: format!("version {}", version),
2723 })
2724 })?;
2725 return Ok((version, transaction));
2726 }
2727
2728 let versions = dataset.versions().await.map_err(|e| {
2729 lance_core::Error::from(NamespaceError::Internal {
2730 message: format!(
2731 "Failed to list table versions while resolving transaction '{}': {}",
2732 id, e
2733 ),
2734 })
2735 })?;
2736
2737 for version in versions.into_iter().rev() {
2738 if let Some(transaction) = dataset
2739 .read_transaction_by_version(version.version)
2740 .await
2741 .map_err(|e| {
2742 lance_core::Error::from(NamespaceError::Internal {
2743 message: format!(
2744 "Failed to read transaction for version {} while resolving '{}': {}",
2745 version.version, id, e
2746 ),
2747 })
2748 })?
2749 && transaction.uuid == id
2750 {
2751 return Ok((version.version, transaction));
2752 }
2753 }
2754
2755 Err(NamespaceError::TransactionNotFound {
2756 message: id.to_string(),
2757 }
2758 .into())
2759 }
2760
2761 const TRANSACTION_ALTERATIONS_DIR: &'static str = "_alter_transactions";
2765
2766 fn transaction_alteration_path(&self, table_uri: &str, txn_uuid: &str) -> Result<Path> {
2767 let table_path = self.object_store_path_from_uri(table_uri)?;
2768 Ok(table_path
2769 .join(Self::TRANSACTION_ALTERATIONS_DIR)
2770 .join(format!("{}.json", txn_uuid).as_str()))
2771 }
2772
2773 async fn load_transaction_alteration(
2774 &self,
2775 table_uri: &str,
2776 txn_uuid: &str,
2777 ) -> Result<Option<TransactionAlteration>> {
2778 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2779 match self.object_store.inner.get(&path).await {
2780 Ok(get_result) => {
2781 let bytes = get_result.bytes().await.map_err(|e| {
2782 lance_core::Error::from(NamespaceError::Internal {
2783 message: format!(
2784 "Failed to read alter_transaction sidecar for '{}': {}",
2785 txn_uuid, e
2786 ),
2787 })
2788 })?;
2789 let alteration = TransactionAlteration::from_json_slice(&bytes).map_err(|e| {
2790 lance_core::Error::from(NamespaceError::Internal {
2791 message: format!(
2792 "Failed to parse alter_transaction sidecar for '{}': {}",
2793 txn_uuid, e
2794 ),
2795 })
2796 })?;
2797 Ok(Some(alteration))
2798 }
2799 Err(ObjectStoreError::NotFound { .. }) => Ok(None),
2800 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2801 message: format!(
2802 "Failed to load alter_transaction sidecar for '{}': {}",
2803 txn_uuid, e
2804 ),
2805 })),
2806 }
2807 }
2808
2809 async fn save_transaction_alteration(
2810 &self,
2811 table_uri: &str,
2812 txn_uuid: &str,
2813 alteration: &TransactionAlteration,
2814 ) -> Result<()> {
2815 let path = self.transaction_alteration_path(table_uri, txn_uuid)?;
2816 let bytes = alteration.to_json_bytes().map_err(|e| {
2817 lance_core::Error::from(NamespaceError::Internal {
2818 message: format!(
2819 "Failed to serialize alter_transaction sidecar for '{}': {}",
2820 txn_uuid, e
2821 ),
2822 })
2823 })?;
2824 self.object_store
2825 .inner
2826 .put(&path, bytes.into())
2827 .await
2828 .map_err(|e| {
2829 lance_core::Error::from(NamespaceError::Internal {
2830 message: format!(
2831 "Failed to persist alter_transaction sidecar for '{}': {}",
2832 txn_uuid, e
2833 ),
2834 })
2835 })?;
2836 Ok(())
2837 }
2838
2839 fn table_full_uri(&self, table_name: &str) -> String {
2840 format!("{}/{}.lance", self.root, table_name)
2841 }
2842
2843 fn table_path(&self, table_name: &str) -> Path {
2845 self.base_path
2846 .clone()
2847 .join(format!("{}.lance", table_name).as_str())
2848 }
2849
2850 fn table_reserved_file_path(&self, table_name: &str) -> Path {
2852 self.base_path
2853 .clone()
2854 .join(format!("{}.lance", table_name).as_str())
2855 .join(".lance-reserved")
2856 }
2857
2858 fn table_deregistered_file_path(&self, table_name: &str) -> Path {
2860 self.base_path
2861 .clone()
2862 .join(format!("{}.lance", table_name).as_str())
2863 .join(".lance-deregistered")
2864 }
2865
2866 pub(crate) async fn check_table_status(&self, table_name: &str) -> Result<TableStatus> {
2872 let table_path = self.table_path(table_name);
2873 match self.object_store.read_dir(table_path).await {
2874 Ok(entries) => {
2875 let exists = !entries.is_empty();
2876 let is_deregistered = entries.iter().any(|e| e.ends_with(".lance-deregistered"));
2877 let has_reserved_file = entries.iter().any(|e| e.ends_with(".lance-reserved"));
2878 Ok(TableStatus {
2879 exists,
2880 is_deregistered,
2881 has_reserved_file,
2882 })
2883 }
2884 Err(e) if manifest::ManifestNamespace::is_not_found_load_error(&e) => Ok(TableStatus {
2887 exists: false,
2888 is_deregistered: false,
2889 has_reserved_file: false,
2890 }),
2891 Err(e) => Err(Self::classify_storage_error(e)),
2895 }
2896 }
2897
2898 fn classify_storage_error(err: Error) -> Error {
2903 if matches!(&err, Error::Namespace { .. }) {
2904 return err;
2905 }
2906 let detail = err.to_string();
2907 if let Error::IO { source, .. } = &err
2908 && let Some(os_err) = source.downcast_ref::<ObjectStoreError>()
2909 {
2910 if is_throttle_error(os_err) {
2911 return NamespaceError::Throttling {
2912 message: format!(
2913 "Storage request was throttled while resolving table: {detail}"
2914 ),
2915 }
2916 .into();
2917 }
2918 if Self::is_service_unavailable_error(os_err) {
2919 return NamespaceError::ServiceUnavailable {
2920 message: format!("Storage service unavailable while resolving table: {detail}"),
2921 }
2922 .into();
2923 }
2924 }
2925 NamespaceError::Internal {
2926 message: format!("Storage error while resolving table: {detail}"),
2927 }
2928 .into()
2929 }
2930
2931 fn is_service_unavailable_error(err: &ObjectStoreError) -> bool {
2935 if let ObjectStoreError::Generic { source, .. } = err {
2936 let message = source.to_string().to_ascii_lowercase();
2937 message.contains("503 service unavailable")
2938 || message.contains("502 bad gateway")
2939 || message.contains("504 gateway timeout")
2940 } else {
2941 false
2942 }
2943 }
2944
2945 fn is_manifest_table_absent_error(err: &Error) -> bool {
2949 if manifest::ManifestNamespace::is_not_found_load_error(err) {
2950 return true;
2951 }
2952 if let Error::Namespace { source, .. } = err
2953 && let Some(ns_err) = source.downcast_ref::<NamespaceError>()
2954 {
2955 return matches!(ns_err, NamespaceError::TableNotFound { .. });
2956 }
2957 false
2958 }
2959
2960 fn map_open_error(err: Error, not_found: NamespaceError) -> Error {
2964 if matches!(&err, Error::IO { .. })
2965 && !manifest::ManifestNamespace::is_not_found_load_error(&err)
2966 {
2967 return Self::classify_storage_error(err);
2968 }
2969 not_found.into()
2970 }
2971
2972 async fn get_storage_options_for_table(
2992 &self,
2993 table_uri: &str,
2994 vend_credentials: bool,
2995 identity: Option<&Identity>,
2996 ) -> Result<Option<HashMap<String, String>>> {
2997 if vend_credentials && let Some(ref vendor) = self.credential_vendor {
2998 let vended = vendor.vend_credentials(table_uri, identity).await?;
2999 return Ok(Some(vended.storage_options));
3000 }
3001 if self.vend_input_storage_options {
3004 let mut options = self.storage_options.clone().unwrap_or_default();
3005 if let Some(refresh_interval_millis) =
3007 self.vend_input_storage_options_refresh_interval_millis
3008 {
3009 let now_millis = std::time::SystemTime::now()
3010 .duration_since(std::time::UNIX_EPOCH)
3011 .unwrap()
3012 .as_millis() as u64;
3013 let expires_at_millis = now_millis + refresh_interval_millis;
3014 options.insert(
3015 "expires_at_millis".to_string(),
3016 expires_at_millis.to_string(),
3017 );
3018 }
3019 return Ok(Some(options));
3020 }
3021 Ok(None)
3024 }
3025
3026 pub async fn migrate(&self) -> Result<usize> {
3080 let Some(manifest_ns) = self.manifest_ns_for_write().await? else {
3082 return Ok(0); };
3084
3085 let manifest_locations = manifest_ns.list_manifest_table_locations().await?;
3087
3088 let dir_tables = self
3091 .filter_declared_tables(self.list_directory_tables().await?, false)
3092 .await?;
3093
3094 let mut migrated_count = 0;
3099 for table_name in dir_tables {
3100 let dir_name = format!("{}.lance", table_name);
3102 if !manifest_locations.contains(&dir_name) {
3103 manifest_ns.register_table(&table_name, dir_name).await?;
3104 migrated_count += 1;
3105 }
3106 }
3107
3108 Ok(migrated_count)
3109 }
3110
3111 async fn delete_physical_version_files(
3119 &self,
3120 table_entries: &[TableDeleteEntry],
3121 branch: Option<&str>,
3122 ) -> Result<i64> {
3123 let mut deleted_count = 0i64;
3124 for te in table_entries {
3125 let table_uri = self.resolve_table_location(&te.table_id).await?;
3126 let table_uri = match branch {
3127 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
3128 None => table_uri,
3129 };
3130 let table_path = self.object_store_path_from_uri(&table_uri)?;
3131 let versions_dir_path = table_path.clone().join(VERSIONS_DIR);
3132
3133 let manifest_metas: Vec<_> = self
3135 .object_store
3136 .read_dir_all(&versions_dir_path, None)
3137 .try_collect()
3138 .await
3139 .map_err(|e| {
3140 lance_core::Error::from(NamespaceError::Internal {
3141 message: format!(
3142 "Failed to list manifest files for table at '{}': {}",
3143 table_uri, e
3144 ),
3145 })
3146 })?;
3147 let location_by_version: HashMap<u64, Path> = manifest_metas
3148 .into_iter()
3149 .filter_map(|meta| {
3150 let version = Self::manifest_version_from_filename(meta.location.filename()?)?;
3151 Some((version, meta.location))
3152 })
3153 .collect();
3154
3155 for (&v, version_path) in &location_by_version {
3156 let vi = v as i64;
3157 if !te.ranges.iter().any(|&(s, e)| vi >= s && (e < 0 || vi < e)) {
3158 continue;
3159 }
3160 match self.object_store.inner.delete(version_path).await {
3161 Ok(_) => {
3162 deleted_count += 1;
3163 }
3164 Err(object_store::Error::NotFound { .. }) => {}
3165 Err(e) => {
3166 return Err(NamespaceError::Internal {
3167 message: format!(
3168 "Failed to delete version {} for table at '{}': {}",
3169 v, table_uri, e
3170 ),
3171 }
3172 .into());
3173 }
3174 }
3175 }
3176 }
3177 Ok(deleted_count)
3178 }
3179
3180 #[allow(clippy::too_many_arguments)]
3185 fn apply_query_params_to_scanner(
3186 scanner: &mut Scanner,
3187 filter: Option<&str>,
3188 columns: Option<&QueryTableRequestColumns>,
3189 vector_column: Option<&str>,
3190 vector: &QueryTableRequestVector,
3191 k: i32,
3192 offset: Option<i32>,
3193 prefilter: Option<bool>,
3194 bypass_vector_index: Option<bool>,
3195 nprobes: Option<i32>,
3196 minimum_nprobes: Option<i32>,
3197 maximum_nprobes: Option<i32>,
3198 ef: Option<i32>,
3199 refine_factor: Option<i32>,
3200 distance_type: Option<&str>,
3201 fast_search_flag: Option<bool>,
3202 with_row_id: Option<bool>,
3203 lower_bound: Option<f32>,
3204 upper_bound: Option<f32>,
3205 operation: &str,
3206 ) -> Result<()> {
3207 if let Some(pf) = prefilter {
3209 scanner.prefilter(pf);
3210 }
3211
3212 if let Some(filter) = filter {
3213 scanner.filter(filter).map_err(|e| {
3214 Error::invalid_input_source(
3215 format!("Invalid filter expression for {}: {}", operation, e).into(),
3216 )
3217 })?;
3218 }
3219
3220 if let Some(cols) = columns {
3221 if let Some(ref names) = cols.column_names {
3222 scanner.project(names.as_slice()).map_err(|e| {
3223 Error::invalid_input_source(
3224 format!("Invalid column projection for {}: {}", operation, e).into(),
3225 )
3226 })?;
3227 } else if let Some(ref aliases) = cols.column_aliases {
3228 let pairs: Vec<(&str, &str)> = aliases
3230 .iter()
3231 .map(|(alias, src)| (alias.as_str(), src.as_str()))
3232 .collect();
3233 scanner.project_with_transform(&pairs).map_err(|e| {
3234 Error::invalid_input_source(
3235 format!("Invalid column aliases for {}: {}", operation, e).into(),
3236 )
3237 })?;
3238 }
3239 }
3240
3241 let query_vec: Option<Vec<f32>> = vector
3243 .single_vector
3244 .as_ref()
3245 .filter(|v| !v.is_empty())
3246 .cloned()
3247 .or_else(|| {
3248 vector
3249 .multi_vector
3250 .as_ref()
3251 .and_then(|mv| mv.first())
3252 .filter(|v| !v.is_empty())
3253 .cloned()
3254 });
3255
3256 if let Some(q_vec) = query_vec {
3257 let col = vector_column.unwrap_or("vector");
3258 let q = Arc::new(Float32Array::from(q_vec));
3259 scanner
3260 .nearest(col, q.as_ref(), k.max(1) as usize)
3261 .map_err(|e| {
3262 Error::invalid_input_source(
3263 format!("Invalid vector query for {}: {}", operation, e).into(),
3264 )
3265 })?;
3266
3267 apply_probe_bounds(scanner, nprobes, minimum_nprobes, maximum_nprobes)?;
3269 if let Some(e) = ef {
3270 scanner.ef(e.max(1) as usize);
3271 }
3272 if let Some(rf) = refine_factor {
3273 scanner.refine(rf.max(0) as u32);
3274 }
3275 if let Some(true) = bypass_vector_index {
3277 scanner.use_index(false);
3278 }
3279 if let Some(true) = fast_search_flag {
3280 scanner.fast_search();
3281 }
3282 if lower_bound.is_some() || upper_bound.is_some() {
3283 scanner.distance_range(lower_bound, upper_bound);
3284 }
3285 if let Some(dt) = distance_type {
3286 let metric = Self::parse_metric_type(Some(dt))?;
3287 scanner.distance_metric(metric);
3288 }
3289 if let Some(off) = offset.filter(|&o| o > 0) {
3291 scanner.limit(None, Some(off as i64)).map_err(|e| {
3292 Error::invalid_input_source(
3293 format!("Invalid offset for {}: {}", operation, e).into(),
3294 )
3295 })?;
3296 }
3297 } else {
3298 let limit = if k > 0 { Some(k as i64) } else { None };
3300 scanner
3301 .limit(limit, offset.map(|o| o as i64))
3302 .map_err(|e| {
3303 Error::invalid_input_source(
3304 format!("Invalid limit/offset for {}: {}", operation, e).into(),
3305 )
3306 })?;
3307 }
3308
3309 if let Some(true) = with_row_id {
3310 scanner.with_row_id();
3311 }
3312
3313 Ok(())
3314 }
3315
3316 pub fn retrieve_ops_metrics(&self) -> HashMap<String, u64> {
3323 self.ops_metrics
3324 .as_ref()
3325 .map(|m| m.retrieve())
3326 .unwrap_or_default()
3327 }
3328
3329 pub fn reset_ops_metrics(&self) {
3333 if let Some(ref metrics) = self.ops_metrics {
3334 metrics.reset();
3335 }
3336 }
3337
3338 fn record_op(&self, operation: &str) {
3340 if let Some(ref metrics) = self.ops_metrics {
3341 metrics.increment(operation);
3342 }
3343 }
3344}
3345
3346#[async_trait]
3347impl LanceNamespace for DirectoryNamespace {
3348 async fn list_namespaces(
3349 &self,
3350 request: ListNamespacesRequest,
3351 ) -> Result<ListNamespacesResponse> {
3352 self.record_op("list_namespaces");
3353 self.ensure_read_manifest().await?;
3354 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3355 return manifest_ns.list_namespaces(request).await;
3356 }
3357
3358 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3359 return Err(self.child_namespace_requires_manifest_error());
3360 }
3361 Self::validate_root_namespace_id(&request.id)?;
3362 Ok(ListNamespacesResponse::new(vec![]))
3363 }
3364
3365 async fn describe_namespace(
3366 &self,
3367 request: DescribeNamespaceRequest,
3368 ) -> Result<DescribeNamespaceResponse> {
3369 self.record_op("describe_namespace");
3370 self.ensure_read_manifest().await?;
3371 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3372 return manifest_ns.describe_namespace(request).await;
3373 }
3374
3375 if request.id.as_ref().is_some_and(|id| !id.is_empty()) {
3376 return Err(self.child_namespace_requires_manifest_error());
3377 }
3378 Self::validate_root_namespace_id(&request.id)?;
3379 #[allow(clippy::needless_update)]
3380 Ok(DescribeNamespaceResponse {
3381 properties: Some(HashMap::new()),
3382 ..Default::default()
3383 })
3384 }
3385
3386 async fn create_namespace(
3387 &self,
3388 request: CreateNamespaceRequest,
3389 ) -> Result<CreateNamespaceResponse> {
3390 self.record_op("create_namespace");
3391 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3392 return manifest_ns.create_namespace(request).await;
3393 }
3394
3395 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3396 return Err(NamespaceError::NamespaceAlreadyExists {
3397 message: "root namespace".to_string(),
3398 }
3399 .into());
3400 }
3401
3402 Err(NamespaceError::Unsupported {
3403 message: "Child namespaces are only supported when manifest mode is enabled"
3404 .to_string(),
3405 }
3406 .into())
3407 }
3408
3409 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
3410 self.record_op("drop_namespace");
3411 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3412 return manifest_ns.drop_namespace(request).await;
3413 }
3414
3415 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3416 return Err(NamespaceError::InvalidInput {
3417 message: "Root namespace cannot be dropped".to_string(),
3418 }
3419 .into());
3420 }
3421
3422 Err(NamespaceError::Unsupported {
3423 message: "Child namespaces are only supported when manifest mode is enabled"
3424 .to_string(),
3425 }
3426 .into())
3427 }
3428
3429 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
3430 self.record_op("namespace_exists");
3431 self.ensure_read_manifest().await?;
3432 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3433 return manifest_ns.namespace_exists(request).await;
3434 }
3435
3436 if request.id.is_none() || request.id.as_ref().unwrap().is_empty() {
3437 return Ok(());
3438 }
3439
3440 Err(self.child_namespace_requires_manifest_error())
3441 }
3442
3443 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
3444 self.record_op("list_tables");
3445 let namespace_id = request.id.as_ref().ok_or_else(|| {
3447 lance_core::Error::from(NamespaceError::InvalidInput {
3448 message: "Namespace ID is required".to_string(),
3449 })
3450 })?;
3451
3452 if !namespace_id.is_empty()
3458 || !self.dir_listing_enabled
3459 || self.dir_listing_to_manifest_migration_enabled
3460 {
3461 self.ensure_read_manifest().await?;
3462 }
3463
3464 if !namespace_id.is_empty() {
3466 if let Some(manifest_ns) = self.manifest_ns_for_read() {
3467 return manifest_ns.list_tables(request).await;
3468 }
3469 return Err(self.child_namespace_requires_manifest_error());
3470 }
3471
3472 if let Some(manifest_ns) = self.manifest_ns_for_read()
3474 && !self.dir_listing_enabled
3475 {
3476 return manifest_ns.list_tables(request).await;
3477 }
3478 if !self.dir_listing_enabled {
3479 return Ok(ListTablesResponse::new(vec![]));
3480 }
3481
3482 let mut tables = if self.manifest_ns_for_read().is_some()
3485 && self.dir_listing_enabled
3486 && self.dir_listing_to_manifest_migration_enabled
3487 {
3488 let manifest_locations = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3490 manifest_ns.list_manifest_table_locations().await?
3491 } else {
3492 std::collections::HashSet::new()
3493 };
3494
3495 let mut manifest_request = request.clone();
3497 manifest_request.limit = None;
3498 manifest_request.page_token = None;
3499 let manifest_tables = if let Some(manifest_ns) = self.manifest_ns_for_read() {
3500 let manifest_response = manifest_ns.list_tables(manifest_request).await?;
3501 manifest_response.tables
3502 } else {
3503 vec![]
3504 };
3505
3506 let mut all_tables: Vec<String> = manifest_tables;
3509 let dir_tables = self.list_directory_tables().await?;
3510 for table_name in dir_tables {
3511 let full_location = format!("{}/{}.lance", self.root, table_name);
3514 let relative_location = format!("{}.lance", table_name);
3515 if !manifest_locations.contains(&full_location)
3516 && !manifest_locations.contains(&relative_location)
3517 {
3518 all_tables.push(table_name);
3519 }
3520 }
3521
3522 all_tables
3523 } else {
3524 self.list_directory_tables().await?
3525 };
3526
3527 tables = self
3528 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
3529 .await?;
3530
3531 let next_page_token =
3533 Self::apply_pagination(&mut tables, request.page_token, request.limit);
3534 let mut response = ListTablesResponse::new(tables);
3535 response.page_token = next_page_token;
3536 Ok(response)
3537 }
3538
3539 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
3540 self.record_op("describe_table");
3541 self.describe_table_impl(request).await
3542 }
3543
3544 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
3545 self.record_op("table_exists");
3546 let is_root_level = request.id.as_ref().is_some_and(|id| id.len() == 1);
3547 let is_child_table = request.id.as_ref().is_some_and(|id| id.len() > 1);
3548 let skip_manifest_for_root = self.dir_listing_enabled
3549 && is_root_level
3550 && !self.dir_listing_to_manifest_migration_enabled;
3551 if is_child_table
3554 || !self.dir_listing_enabled
3555 || self.dir_listing_to_manifest_migration_enabled
3556 {
3557 self.ensure_read_manifest().await?;
3558 }
3559 if let Some(manifest_ns) = self.manifest_ns_for_read()
3560 && !skip_manifest_for_root
3561 {
3562 match manifest_ns.table_exists(request.clone()).await {
3563 Ok(()) => return Ok(()),
3564 Err(e) if manifest_feature_flags::is_incompatible_manifest_error(&e) => {
3565 return Err(e);
3568 }
3569 Err(e) if self.dir_listing_enabled && is_root_level => {
3570 if !Self::is_manifest_table_absent_error(&e) {
3574 return Err(Self::classify_storage_error(e));
3575 }
3576 }
3577 Err(e) => return Err(e),
3578 }
3579 }
3580 if is_child_table {
3581 return Err(self.child_namespace_requires_manifest_error());
3582 }
3583
3584 let table_name = Self::table_name_from_id(&request.id)?;
3585 let table_id = Self::format_table_id_from_request(&request.id);
3586 if !self.dir_listing_enabled {
3587 return Err(NamespaceError::TableNotFound { message: table_id }.into());
3588 }
3589
3590 let status = self.check_table_status(&table_name).await?;
3592
3593 if !status.exists {
3594 return Err(NamespaceError::TableNotFound {
3595 message: table_id.clone(),
3596 }
3597 .into());
3598 }
3599
3600 if status.is_deregistered {
3601 return Err(NamespaceError::TableNotFound {
3602 message: format!("Table is deregistered: {}", table_id),
3603 }
3604 .into());
3605 }
3606
3607 Ok(())
3608 }
3609
3610 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3611 self.record_op("drop_table");
3612 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3613 return manifest_ns.drop_table(request).await;
3614 }
3615
3616 let table_name = Self::table_name_from_id(&request.id)?;
3617 let table_uri = self.table_full_uri(&table_name);
3618 let table_path = self.table_path(&table_name);
3619
3620 self.object_store
3621 .remove_dir_all(table_path)
3622 .await
3623 .map_err(|e| {
3624 lance_core::Error::from(NamespaceError::Internal {
3625 message: format!("Failed to drop table {}: {:?}", table_name, e),
3626 })
3627 })?;
3628
3629 Ok(DropTableResponse {
3630 id: request.id,
3631 location: Some(table_uri),
3632 ..Default::default()
3633 })
3634 }
3635
3636 async fn create_table(
3637 &self,
3638 request: CreateTableRequest,
3639 request_data: Bytes,
3640 ) -> Result<CreateTableResponse> {
3641 self.record_op("create_table");
3642 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3643 return manifest_ns.create_table(request, request_data).await;
3644 }
3645
3646 Self::validate_dir_only_properties(request.properties.as_ref(), "create_table")?;
3647
3648 let table_name = Self::table_name_from_id(&request.id)?;
3649 let table_uri = self.table_full_uri(&table_name);
3650 let status = self.check_table_status(&table_name).await?;
3651 let (reader, _num_rows) =
3652 Self::ipc_reader_from_request_data(&request_data, "create_table")?;
3653
3654 if status.exists && self.table_has_actual_manifests(&table_name).await? {
3655 return Err(NamespaceError::TableAlreadyExists {
3656 message: table_name,
3657 }
3658 .into());
3659 }
3660
3661 let write_result = self
3662 .write_reader_to_table(
3663 &table_uri,
3664 reader,
3665 WriteMode::Create,
3666 request.storage_options.clone(),
3667 )
3668 .await;
3669 if let Err(err) = write_result {
3670 if self.table_uri_has_actual_manifests(&table_uri).await? {
3671 return Err(NamespaceError::TableAlreadyExists {
3672 message: table_name,
3673 }
3674 .into());
3675 }
3676 return Err(err);
3677 }
3678 Ok(CreateTableResponse {
3679 version: Some(1),
3680 location: Some(table_uri),
3681 storage_options: self.storage_options.clone(),
3682 properties: request.properties,
3683 ..Default::default()
3684 })
3685 }
3686
3687 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3688 self.record_op("declare_table");
3689 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3690 let mut response = manifest_ns.declare_table(request.clone()).await?;
3691 if let Some(ref location) = response.location {
3692 let vend = request.vend_credentials.unwrap_or(true);
3694 let identity = request.identity.as_deref();
3695 response.storage_options = self
3696 .get_storage_options_for_table(location, vend, identity)
3697 .await?;
3698 }
3699 if self.table_version_tracking_enabled {
3701 response.managed_versioning = Some(true);
3702 }
3703 return Ok(response);
3704 }
3705
3706 Self::validate_dir_only_properties(request.properties.as_ref(), "declare_table")?;
3707
3708 let table_name = Self::table_name_from_id(&request.id)?;
3709 let table_uri = self.table_full_uri(&table_name);
3710
3711 if let Some(location) = &request.location {
3713 let location = location.trim_end_matches('/');
3714 if location != table_uri {
3715 return Err(NamespaceError::InvalidInput {
3716 message: format!(
3717 "Cannot declare table {} at location {}, must be at location {}",
3718 table_name, location, table_uri
3719 ),
3720 }
3721 .into());
3722 }
3723 }
3724
3725 let status = self.check_table_status(&table_name).await?;
3729 if status.exists && !status.has_reserved_file {
3730 return Err(NamespaceError::TableAlreadyExists {
3732 message: table_name.to_string(),
3733 }
3734 .into());
3735 }
3736
3737 let reserved_file_path = self.table_reserved_file_path(&table_name);
3741
3742 put_marker_file_atomic(
3743 &self.object_store,
3744 &reserved_file_path,
3745 &format!("table {}", table_name),
3746 )
3747 .await
3748 .map_err(|e| match e {
3749 MarkerFileError::AlreadyExists { .. } => {
3750 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3751 message: table_name.to_string(),
3752 })
3753 }
3754 MarkerFileError::Other { message } => {
3755 lance_core::Error::from(NamespaceError::Internal { message })
3756 }
3757 })?;
3758
3759 let vend_credentials = request.vend_credentials.unwrap_or(true);
3761 let identity = request.identity.as_deref();
3762 let storage_options = self
3763 .get_storage_options_for_table(&table_uri, vend_credentials, identity)
3764 .await?;
3765
3766 Ok(DeclareTableResponse {
3767 location: Some(table_uri),
3768 storage_options,
3769 properties: request.properties,
3770 managed_versioning: if self.table_version_tracking_enabled {
3771 Some(true)
3772 } else {
3773 None
3774 },
3775 ..Default::default()
3776 })
3777 }
3778
3779 async fn register_table(
3780 &self,
3781 request: lance_namespace::models::RegisterTableRequest,
3782 ) -> Result<lance_namespace::models::RegisterTableResponse> {
3783 self.record_op("register_table");
3784 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3786 return LanceNamespace::register_table(manifest_ns.as_ref(), request).await;
3787 }
3788
3789 Err(NamespaceError::Unsupported {
3791 message: "register_table is only supported when manifest mode is enabled".to_string(),
3792 }
3793 .into())
3794 }
3795
3796 async fn deregister_table(
3797 &self,
3798 request: lance_namespace::models::DeregisterTableRequest,
3799 ) -> Result<lance_namespace::models::DeregisterTableResponse> {
3800 self.record_op("deregister_table");
3801 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3803 return LanceNamespace::deregister_table(manifest_ns.as_ref(), request).await;
3804 }
3805
3806 let table_name = Self::table_name_from_id(&request.id)?;
3808 let table_uri = self.table_full_uri(&table_name);
3809
3810 let status = self.check_table_status(&table_name).await?;
3813
3814 if !status.exists {
3815 return Err(NamespaceError::TableNotFound {
3816 message: table_name.to_string(),
3817 }
3818 .into());
3819 }
3820
3821 if status.is_deregistered {
3822 return Err(NamespaceError::TableNotFound {
3823 message: format!("Table is already deregistered: {}", table_name),
3824 }
3825 .into());
3826 }
3827
3828 let deregistered_path = self.table_deregistered_file_path(&table_name);
3834 put_marker_file_atomic(
3835 &self.object_store,
3836 &deregistered_path,
3837 &format!("deregistration marker for table {}", table_name),
3838 )
3839 .await
3840 .map_err(|e| match e {
3841 MarkerFileError::AlreadyExists { .. } => {
3842 lance_core::Error::from(NamespaceError::InvalidTableState {
3843 message: format!("Table is already deregistered: {}", table_name),
3844 })
3845 }
3846 MarkerFileError::Other { message } => {
3847 lance_core::Error::from(NamespaceError::Internal { message })
3848 }
3849 })?;
3850
3851 Ok(lance_namespace::models::DeregisterTableResponse {
3852 id: request.id,
3853 location: Some(table_uri),
3854 ..Default::default()
3855 })
3856 }
3857
3858 async fn alter_table_add_columns(
3859 &self,
3860 request: AlterTableAddColumnsRequest,
3861 ) -> Result<AlterTableAddColumnsResponse> {
3862 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3863 return manifest_ns.alter_table_add_columns(request).await;
3864 }
3865
3866 let table_name = Self::table_name_from_id(&request.id)?;
3868 let table_uri = self.table_full_uri(&table_name);
3869
3870 let status = self.check_table_status(&table_name).await?;
3872 if !status.exists {
3873 return Err(NamespaceError::TableNotFound {
3874 message: table_name,
3875 }
3876 .into());
3877 }
3878 if status.is_deregistered {
3879 return Err(NamespaceError::TableNotFound {
3880 message: format!("Table is deregistered: {}", table_name),
3881 }
3882 .into());
3883 }
3884
3885 let mut dataset = self
3886 .configured_builder(&table_uri)
3887 .load()
3888 .await
3889 .map_err(|e| {
3890 Error::io_source(box_error(std::io::Error::other(format!(
3891 "Failed to open dataset: {}",
3892 e
3893 ))))
3894 })?;
3895
3896 let sql_expressions = build_sql_expressions(&request.new_columns)?;
3897
3898 dataset
3899 .add_columns(
3900 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3901 None,
3902 None,
3903 )
3904 .await
3905 .map_err(|e| {
3906 Error::io_source(box_error(std::io::Error::other(format!(
3907 "Failed to add columns: {}",
3908 e
3909 ))))
3910 })?;
3911
3912 let version = dataset.version().version as i64;
3913 Ok(AlterTableAddColumnsResponse::new(version))
3914 }
3915
3916 async fn alter_table_alter_columns(
3917 &self,
3918 request: AlterTableAlterColumnsRequest,
3919 ) -> Result<AlterTableAlterColumnsResponse> {
3920 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3921 return manifest_ns.alter_table_alter_columns(request).await;
3922 }
3923
3924 let table_name = Self::table_name_from_id(&request.id)?;
3925 let table_uri = self.table_full_uri(&table_name);
3926
3927 let status = self.check_table_status(&table_name).await?;
3929 if !status.exists {
3930 return Err(NamespaceError::TableNotFound {
3931 message: table_name,
3932 }
3933 .into());
3934 }
3935 if status.is_deregistered {
3936 return Err(NamespaceError::TableNotFound {
3937 message: format!("Table is deregistered: {}", table_name),
3938 }
3939 .into());
3940 }
3941
3942 let mut dataset = self
3943 .configured_builder(&table_uri)
3944 .load()
3945 .await
3946 .map_err(|e| {
3947 Error::io_source(box_error(std::io::Error::other(format!(
3948 "Failed to open dataset: {}",
3949 e
3950 ))))
3951 })?;
3952
3953 let alterations = build_column_alterations(&request.alterations)?;
3954
3955 dataset.alter_columns(&alterations).await.map_err(|e| {
3956 Error::io_source(box_error(std::io::Error::other(format!(
3957 "Failed to alter columns: {}",
3958 e
3959 ))))
3960 })?;
3961
3962 let version = dataset.version().version as i64;
3963 Ok(AlterTableAlterColumnsResponse::new(version))
3964 }
3965
3966 async fn alter_table_drop_columns(
3967 &self,
3968 request: AlterTableDropColumnsRequest,
3969 ) -> Result<AlterTableDropColumnsResponse> {
3970 if let Some(manifest_ns) = self.manifest_ns_for_write().await? {
3971 return manifest_ns.alter_table_drop_columns(request).await;
3972 }
3973
3974 let table_name = Self::table_name_from_id(&request.id)?;
3975 let table_uri = self.table_full_uri(&table_name);
3976
3977 let status = self.check_table_status(&table_name).await?;
3979 if !status.exists {
3980 return Err(NamespaceError::TableNotFound {
3981 message: table_name,
3982 }
3983 .into());
3984 }
3985 if status.is_deregistered {
3986 return Err(NamespaceError::TableNotFound {
3987 message: format!("Table is deregistered: {}", table_name),
3988 }
3989 .into());
3990 }
3991
3992 let mut dataset = self
3993 .configured_builder(&table_uri)
3994 .load()
3995 .await
3996 .map_err(|e| {
3997 Error::io_source(box_error(std::io::Error::other(format!(
3998 "Failed to open dataset: {}",
3999 e
4000 ))))
4001 })?;
4002
4003 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
4004 dataset.drop_columns(&columns).await.map_err(|e| {
4005 Error::io_source(box_error(std::io::Error::other(format!(
4006 "Failed to drop columns: {}",
4007 e
4008 ))))
4009 })?;
4010
4011 let version = dataset.version().version as i64;
4012 Ok(AlterTableDropColumnsResponse::new(version))
4013 }
4014
4015 async fn list_table_versions(
4016 &self,
4017 request: ListTableVersionsRequest,
4018 ) -> Result<ListTableVersionsResponse> {
4019 self.record_op("list_table_versions");
4020 let branch = Self::normalized_branch(request.branch.as_deref())?;
4021 let table_uri = self.resolve_table_location(&request.id).await?;
4022 let table_uri = match branch {
4023 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
4024 None => table_uri,
4025 };
4026 let want_descending = request.descending == Some(true);
4027 let table_versions = self
4028 .list_table_versions_from_storage(&table_uri, want_descending, request.limit)
4029 .await?;
4030
4031 Ok(ListTableVersionsResponse {
4032 versions: table_versions,
4033 page_token: None,
4034 ..Default::default()
4035 })
4036 }
4037
4038 async fn create_table_version(
4039 &self,
4040 request: CreateTableVersionRequest,
4041 ) -> Result<CreateTableVersionResponse> {
4042 self.record_op("create_table_version");
4043 let branch = Self::normalized_branch(request.branch.as_deref())?;
4044 let table_uri = self.resolve_table_location(&request.id).await?;
4045 let (table_uri, table_path, branch_parent_version) = match branch {
4046 Some(b) => self.resolve_branch_for_commit(&table_uri, b).await?,
4047 None => {
4048 let table_path = self.object_store_path_from_uri(&table_uri)?;
4049 (table_uri, table_path, None)
4050 }
4051 };
4052
4053 let staging_manifest_path = &request.manifest_path;
4054 let version = request.version as u64;
4055
4056 let naming_scheme = match request.naming_scheme.as_deref() {
4058 Some("V1") => ManifestNamingScheme::V1,
4059 _ => ManifestNamingScheme::V2,
4060 };
4061
4062 let final_path = naming_scheme.manifest_path(&table_path, version);
4064
4065 let staging_path = Path::parse(staging_manifest_path).map_err(|e| {
4066 lance_core::Error::from(NamespaceError::InvalidInput {
4067 message: format!(
4068 "Invalid staging manifest path '{}': {}",
4069 staging_manifest_path, e
4070 ),
4071 })
4072 })?;
4073
4074 match self.object_store.inner.head(&final_path).await {
4076 Ok(existing_meta) => {
4077 return self
4078 .resolve_existing_table_version(ExistingTableVersionResolve {
4079 staging_path: &staging_path,
4080 final_path: &final_path,
4081 version,
4082 table_uri: &table_uri,
4083 final_meta: &existing_meta,
4084 request_manifest_size: request.manifest_size,
4085 })
4086 .await;
4087 }
4088 Err(ObjectStoreError::NotFound { .. }) => {}
4089 Err(e) => {
4090 return Err(lance_core::Error::from(NamespaceError::Internal {
4091 message: format!(
4092 "Failed to stat version {} for table at '{}': {}",
4093 version, table_uri, e
4094 ),
4095 }));
4096 }
4097 }
4098
4099 let is_branch = branch.is_some();
4102 self.enforce_create_table_version_cas(
4103 &table_path,
4104 version,
4105 &table_uri,
4106 is_branch,
4107 branch_parent_version,
4108 )
4109 .await?;
4110
4111 let copy_result = self
4113 .materialize_version_manifest_create(&staging_path, &final_path, staging_manifest_path)
4114 .await;
4115
4116 match copy_result {
4117 Ok(()) => {}
4118 Err(ObjectStoreError::AlreadyExists { .. })
4119 | Err(ObjectStoreError::Precondition { .. }) => {
4120 let existing_meta = self.object_store.inner.head(&final_path).await.map_err(|e| {
4122 lance_core::Error::from(NamespaceError::Internal {
4123 message: format!(
4124 "Version {} conflict for table at '{}' but failed to stat winner: {}",
4125 version, table_uri, e
4126 ),
4127 })
4128 })?;
4129 return self
4130 .resolve_existing_table_version(ExistingTableVersionResolve {
4131 staging_path: &staging_path,
4132 final_path: &final_path,
4133 version,
4134 table_uri: &table_uri,
4135 final_meta: &existing_meta,
4136 request_manifest_size: request.manifest_size,
4137 })
4138 .await;
4139 }
4140 Err(ObjectStoreError::NotFound { .. }) => {
4141 return Err(lance_core::Error::from(NamespaceError::InvalidInput {
4142 message: format!(
4143 "Staging manifest not found at '{}' for version {} of table at '{}'",
4144 staging_manifest_path, version, table_uri
4145 ),
4146 }));
4147 }
4148 Err(e) => {
4149 return Err(lance_core::Error::from(NamespaceError::Internal {
4150 message: format!(
4151 "Failed to create version {} for table at '{}': {}",
4152 version, table_uri, e
4153 ),
4154 }));
4155 }
4156 }
4157
4158 let final_meta = self
4159 .object_store
4160 .inner
4161 .head(&final_path)
4162 .await
4163 .map_err(|e| {
4164 lance_core::Error::from(NamespaceError::Internal {
4165 message: format!(
4166 "Failed to stat created version {} for table at '{}': {}",
4167 version, table_uri, e
4168 ),
4169 })
4170 })?;
4171
4172 if let Err(e) = self.object_store.inner.delete(&staging_path).await {
4174 log::warn!(
4175 "Failed to delete staging manifest at '{}': {:?}",
4176 staging_path,
4177 e
4178 );
4179 }
4180
4181 Ok(Self::create_table_version_response(
4182 version,
4183 &final_path,
4184 &final_meta,
4185 ))
4186 }
4187
4188 async fn describe_table_version(
4189 &self,
4190 request: DescribeTableVersionRequest,
4191 ) -> Result<DescribeTableVersionResponse> {
4192 self.record_op("describe_table_version");
4193 let branch = Self::normalized_branch(request.branch.as_deref())?;
4194 let table_uri = self.resolve_table_location(&request.id).await?;
4195 let table_uri = match branch {
4196 Some(b) => self.resolve_branch_location(&table_uri, b).await?,
4197 None => table_uri,
4198 };
4199 let versions = self
4200 .list_table_versions_from_storage(&table_uri, true, None)
4201 .await?;
4202 let table_version = if let Some(requested_version) = request.version {
4203 versions
4204 .into_iter()
4205 .find(|version| version.version == requested_version)
4206 .ok_or_else(|| {
4207 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4208 message: format!(
4209 "version {} for table {}",
4210 requested_version,
4211 Self::format_table_id_from_request(&request.id)
4212 ),
4213 })
4214 })?
4215 } else {
4216 versions.into_iter().next().ok_or_else(|| {
4217 lance_core::Error::from(NamespaceError::TableVersionNotFound {
4218 message: format!(
4219 "latest version for table {}",
4220 Self::format_table_id_from_request(&request.id)
4221 ),
4222 })
4223 })?
4224 };
4225
4226 Ok(DescribeTableVersionResponse {
4227 version: Box::new(table_version),
4228 ..Default::default()
4229 })
4230 }
4231
4232 async fn batch_delete_table_versions(
4233 &self,
4234 request: BatchDeleteTableVersionsRequest,
4235 ) -> Result<BatchDeleteTableVersionsResponse> {
4236 self.record_op("batch_delete_table_versions");
4237 let branch = Self::normalized_branch(request.branch.as_deref())?;
4238 let ranges: Vec<(i64, i64)> = request
4241 .ranges
4242 .iter()
4243 .map(|r| (r.start_version, r.end_version))
4244 .collect();
4245
4246 const MAX_VERSIONS_PER_REQUEST: i128 = 1_000_000;
4250 let requested: i128 = ranges
4251 .iter()
4252 .map(|(s, e)| {
4253 if *e < 0 {
4254 0
4255 } else {
4256 (*e as i128 - *s as i128).max(0)
4257 }
4258 })
4259 .sum();
4260 if requested > MAX_VERSIONS_PER_REQUEST {
4261 return Err(NamespaceError::InvalidInput {
4262 message: format!(
4263 "batch_delete requested {} versions; limit is {}",
4264 requested, MAX_VERSIONS_PER_REQUEST
4265 ),
4266 }
4267 .into());
4268 }
4269
4270 let table_entries = vec![TableDeleteEntry {
4271 table_id: request.id.clone(),
4272 ranges,
4273 }];
4274
4275 let total_deleted_count = self
4276 .delete_physical_version_files(&table_entries, branch)
4277 .await?;
4278
4279 Ok(BatchDeleteTableVersionsResponse {
4280 deleted_count: Some(total_deleted_count),
4281 transaction_id: None,
4282 ..Default::default()
4283 })
4284 }
4285
4286 async fn create_table_index(
4287 &self,
4288 request: CreateTableIndexRequest,
4289 ) -> Result<CreateTableIndexResponse> {
4290 self.record_op("create_table_index");
4291 let table_uri = self.resolve_table_location(&request.id).await?;
4292 let mut dataset = self
4293 .load_dataset(&table_uri, None, "create_table_index")
4294 .await?;
4295 let index_request = Self::build_index_params(&request)?;
4296
4297 dataset
4298 .create_index(
4299 &[request.column.as_str()],
4300 index_request.index_type(),
4301 request.name.clone(),
4302 index_request.params(),
4303 false,
4304 )
4305 .await
4306 .map_err(|e| {
4307 let err_msg = format!("{}", e);
4308 let ns_err = if err_msg.contains("already exists") {
4309 NamespaceError::TableIndexAlreadyExists {
4310 message: format!(
4311 "Index '{}' already exists on table '{}': {:?}",
4312 request.name.as_deref().unwrap_or("<auto-generated>"),
4313 table_uri,
4314 e
4315 ),
4316 }
4317 } else if err_msg.contains("not found") || err_msg.contains("does not exist") {
4318 NamespaceError::TableColumnNotFound {
4319 message: format!(
4320 "Column '{}' not found for table '{}': {:?}",
4321 request.column, table_uri, e
4322 ),
4323 }
4324 } else {
4325 NamespaceError::Internal {
4326 message: format!(
4327 "Failed to create {} index '{}' on column '{}' for table '{}': {:?}",
4328 request.index_type,
4329 request.name.as_deref().unwrap_or("<auto-generated>"),
4330 request.column,
4331 table_uri,
4332 e
4333 ),
4334 }
4335 };
4336 lance_core::Error::from(ns_err)
4337 })?;
4338
4339 let transaction_id = dataset
4340 .read_transaction()
4341 .await
4342 .map_err(|e| {
4343 lance_core::Error::from(NamespaceError::Internal {
4344 message: format!(
4345 "Failed to read committed transaction after creating index on '{}': {}",
4346 table_uri, e
4347 ),
4348 })
4349 })?
4350 .map(|transaction| transaction.uuid);
4351
4352 Ok(CreateTableIndexResponse {
4353 transaction_id,
4354 ..Default::default()
4355 })
4356 }
4357
4358 async fn list_table_indices(
4359 &self,
4360 request: ListTableIndicesRequest,
4361 ) -> Result<ListTableIndicesResponse> {
4362 self.record_op("list_table_indices");
4363 let table_uri = self.resolve_table_location(&request.id).await?;
4364 let dataset = self
4365 .load_dataset(&table_uri, request.version, "list_table_indices")
4366 .await?;
4367 let total_rows = dataset.count_rows(None).await.map_err(|e| {
4368 lance_core::Error::from(NamespaceError::Internal {
4369 message: format!("Failed to count rows for table '{}': {:?}", table_uri, e),
4370 })
4371 })? as u64;
4372 let mut indices = dataset
4373 .describe_indices(None)
4374 .await
4375 .map_err(|e| {
4376 lance_core::Error::from(NamespaceError::Internal {
4377 message: format!("Failed to describe table indices for '{}': {:?}", table_uri, e),
4378 })
4379 })?
4380 .into_iter()
4381 .filter(|description| {
4382 description
4383 .metadata()
4384 .first()
4385 .map(|metadata| !is_system_index(metadata))
4386 .unwrap_or(false)
4387 })
4388 .map(|description| {
4389 let columns = description
4390 .field_ids()
4391 .iter()
4392 .map(|field_id| {
4393 dataset
4394 .schema()
4395 .field_path_minimal(i32::try_from(*field_id).map_err(|e| {
4396 lance_core::Error::from(NamespaceError::Internal {
4397 message: format!(
4398 "Field id {} does not fit in i32 for table '{}': {}",
4399 field_id, table_uri, e
4400 ),
4401 })
4402 })?)
4403 .map_err(|e| {
4404 lance_core::Error::from(NamespaceError::Internal {
4405 message: format!(
4406 "Failed to resolve field path for field_id {} in table '{}': {}",
4407 field_id, table_uri, e
4408 ),
4409 })
4410 })
4411 })
4412 .collect::<Result<Vec<_>>>()?;
4413
4414 let segments = description.segments();
4415 let created_at = segments
4416 .iter()
4417 .filter_map(|segment| segment.created_at)
4418 .min()
4419 .map(|ts| ts.to_rfc3339());
4420
4421 #[allow(clippy::needless_update)]
4424 let content = IndexContent {
4425 index_name: description.name().to_string(),
4426 index_uuid: description.metadata()[0].uuid.to_string(),
4427 columns,
4428 status: "SUCCEEDED".to_string(),
4429 index_type: Some(description.index_type().to_string()),
4430 type_url: Some(description.type_url().to_string()),
4431 num_indexed_rows: Some(description.rows_indexed() as i64),
4432 num_unindexed_rows: Some(
4433 total_rows.saturating_sub(description.rows_indexed()) as i64,
4434 ),
4435 size_bytes: description.total_size_bytes().map(|size| size as i64),
4436 num_segments: Some(segments.len() as i32),
4437 created_at,
4438 index_version: segments.first().map(|segment| segment.index_version),
4439 index_details: description.details().ok(),
4440 ..Default::default()
4441 };
4442 Ok(content)
4443 })
4444 .collect::<Result<Vec<_>>>()?;
4445
4446 let page_token = Self::paginate_indices(&mut indices, request.page_token, request.limit);
4447 Ok(ListTableIndicesResponse {
4448 indexes: indices,
4449 page_token,
4450 ..Default::default()
4451 })
4452 }
4453
4454 async fn describe_table_index_stats(
4455 &self,
4456 request: DescribeTableIndexStatsRequest,
4457 ) -> Result<DescribeTableIndexStatsResponse> {
4458 self.record_op("describe_table_index_stats");
4459 let table_uri = self.resolve_table_location(&request.id).await?;
4460 let dataset = self
4461 .load_dataset(&table_uri, request.version, "describe_table_index_stats")
4462 .await?;
4463 let index_name = request.index_name.as_deref().ok_or_else(|| {
4464 lance_core::Error::from(NamespaceError::InvalidInput {
4465 message: "Index name is required for describe_table_index_stats".to_string(),
4466 })
4467 })?;
4468 let metadatas = dataset
4469 .load_indices_by_name(index_name)
4470 .await
4471 .map_err(|e| {
4472 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4473 message: format!(
4474 "Failed to load index '{}' metadata for table '{}': {}",
4475 index_name, table_uri, e
4476 ),
4477 })
4478 })?;
4479 if metadatas.first().is_some_and(is_system_index) {
4480 return Err(NamespaceError::Unsupported {
4481 message: format!("System index '{}' is not exposed by this API", index_name),
4482 }
4483 .into());
4484 }
4485
4486 let stats = <Dataset as DatasetIndexExt>::index_statistics(&dataset, index_name)
4487 .await
4488 .map_err(|e| {
4489 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4490 message: format!(
4491 "Failed to describe index statistics for '{}' on table '{}': {}",
4492 index_name, table_uri, e
4493 ),
4494 })
4495 })?;
4496 let stats: serde_json::Value = serde_json::from_str(&stats).map_err(|e| {
4497 lance_core::Error::from(NamespaceError::Internal {
4498 message: format!(
4499 "Failed to parse index statistics for '{}' on table '{}': {}",
4500 index_name, table_uri, e
4501 ),
4502 })
4503 })?;
4504
4505 Ok(Self::describe_table_index_stats_response(&stats))
4506 }
4507
4508 async fn describe_transaction(
4509 &self,
4510 request: DescribeTransactionRequest,
4511 ) -> Result<DescribeTransactionResponse> {
4512 self.record_op("describe_transaction");
4513 let mut request_id = request.id.ok_or_else(|| {
4514 lance_core::Error::from(NamespaceError::InvalidInput {
4515 message: "Transaction id must include table id and transaction identifier"
4516 .to_string(),
4517 })
4518 })?;
4519 if request_id.len() < 2 {
4520 return Err(NamespaceError::InvalidInput {
4521 message: format!(
4522 "Transaction request id must include table id and transaction identifier, got {:?}",
4523 request_id
4524 ),
4525 }
4526 .into());
4527 }
4528
4529 let id = request_id.pop().expect("request_id len checked above");
4530 let table_id = Some(request_id);
4531 let table_uri = self.resolve_table_location(&table_id).await?;
4532 let dataset = self
4533 .load_dataset(&table_uri, None, "describe_transaction")
4534 .await?;
4535 let (version, transaction) = self.find_transaction(&dataset, &id).await?;
4536
4537 let sidecar = self
4540 .load_transaction_alteration(&table_uri, &transaction.uuid)
4541 .await?;
4542
4543 Ok(Self::transaction_response(version, &transaction, sidecar))
4544 }
4545
4546 async fn alter_transaction(
4547 &self,
4548 request: AlterTransactionRequest,
4549 ) -> Result<AlterTransactionResponse> {
4550 self.record_op("alter_transaction");
4551
4552 let mut request_id = request.id.ok_or_else(|| {
4554 lance_core::Error::from(NamespaceError::InvalidInput {
4555 message: "Transaction id must include table id and transaction identifier"
4556 .to_string(),
4557 })
4558 })?;
4559 if request_id.len() < 2 {
4560 return Err(NamespaceError::InvalidInput {
4561 message: format!(
4562 "Transaction request id must include table id and transaction identifier, got {:?}",
4563 request_id
4564 ),
4565 }
4566 .into());
4567 }
4568
4569 let txn_id = request_id.pop().expect("request_id len checked above");
4570 let table_id = Some(request_id);
4571 let table_uri = self.resolve_table_location(&table_id).await?;
4572 let dataset = self
4573 .load_dataset(&table_uri, None, "alter_transaction")
4574 .await?;
4575 let (version, transaction) = self.find_transaction(&dataset, &txn_id).await?;
4576
4577 const RESERVED_KEYS: &[&str] = &["uuid", "version", "read_version", "operation", "tag"];
4581 let is_reserved = |key: &str| RESERVED_KEYS.contains(&key);
4582
4583 let mut sidecar = self
4586 .load_transaction_alteration(&table_uri, &transaction.uuid)
4587 .await?
4588 .unwrap_or_default();
4589
4590 for action in &request.actions {
4591 if let Some(ref set_status) = action.set_status_action
4592 && let Some(ref status) = set_status.status
4593 {
4594 let normalized = status.to_lowercase().replace('_', "");
4596 match normalized.as_str() {
4597 "queued" | "running" | "succeeded" | "failed" | "canceled" => {
4598 sidecar.status = Some(status.clone());
4599 }
4600 _ => {
4601 return Err(NamespaceError::InvalidInput {
4602 message: format!(
4603 "Invalid transaction status '{}'. Valid values are: Queued, Running, Succeeded, Failed, Canceled",
4604 status
4605 ),
4606 }
4607 .into());
4608 }
4609 }
4610 }
4611
4612 if let Some(ref set_property) = action.set_property_action
4613 && let (Some(key), Some(value)) = (&set_property.key, &set_property.value)
4614 {
4615 if is_reserved(key) {
4616 return Err(NamespaceError::InvalidInput {
4617 message: format!("Property '{}' is reserved and cannot be modified", key),
4618 }
4619 .into());
4620 }
4621 let mode = set_property
4622 .mode
4623 .as_deref()
4624 .unwrap_or("Overwrite")
4625 .to_lowercase();
4626 match mode.as_str() {
4627 "overwrite" => {
4628 sidecar.properties.insert(key.clone(), value.clone());
4629 }
4630 "fail" => {
4631 let exists = sidecar.properties.contains_key(key)
4634 || transaction
4635 .transaction_properties
4636 .as_ref()
4637 .is_some_and(|props| props.contains_key(key));
4638 if exists {
4639 return Err(NamespaceError::ConcurrentModification {
4640 message: format!(
4641 "Property '{}' already exists and mode is 'Fail'",
4642 key
4643 ),
4644 }
4645 .into());
4646 }
4647 sidecar.properties.insert(key.clone(), value.clone());
4648 }
4649 "skip" => {
4650 let exists = sidecar.properties.contains_key(key)
4651 || transaction
4652 .transaction_properties
4653 .as_ref()
4654 .is_some_and(|props| props.contains_key(key));
4655 if !exists {
4656 sidecar.properties.insert(key.clone(), value.clone());
4657 }
4658 }
4659 _ => {
4660 return Err(NamespaceError::InvalidInput {
4661 message: format!(
4662 "Invalid set_property mode '{}'. Valid values are: Overwrite, Fail, Skip",
4663 mode
4664 ),
4665 }
4666 .into());
4667 }
4668 }
4669 }
4670
4671 if let Some(ref unset_property) = action.unset_property_action
4672 && let Some(ref key) = unset_property.key
4673 {
4674 if is_reserved(key) {
4675 return Err(NamespaceError::InvalidInput {
4676 message: format!("Property '{}' is reserved and cannot be modified", key),
4677 }
4678 .into());
4679 }
4680 let mode = unset_property
4681 .mode
4682 .as_deref()
4683 .unwrap_or("Skip")
4684 .to_lowercase();
4685 let exists_in_transaction = transaction
4686 .transaction_properties
4687 .as_ref()
4688 .is_some_and(|props| props.contains_key(key));
4689 match mode.as_str() {
4690 "skip" => {
4691 sidecar.properties.remove(key);
4692 if exists_in_transaction {
4693 sidecar.removed_properties.insert(key.clone());
4696 }
4697 }
4698 "fail" => {
4699 if !sidecar.properties.contains_key(key) && !exists_in_transaction {
4700 return Err(NamespaceError::InvalidInput {
4701 message: format!(
4702 "Property '{}' does not exist and mode is 'Fail'",
4703 key
4704 ),
4705 }
4706 .into());
4707 }
4708 sidecar.properties.remove(key);
4709 if exists_in_transaction {
4710 sidecar.removed_properties.insert(key.clone());
4711 }
4712 }
4713 _ => {
4714 return Err(NamespaceError::InvalidInput {
4715 message: format!(
4716 "Invalid unset_property mode '{}'. Valid values are: Skip, Fail",
4717 mode
4718 ),
4719 }
4720 .into());
4721 }
4722 }
4723 }
4724 }
4725
4726 self.save_transaction_alteration(&table_uri, &transaction.uuid, &sidecar)
4730 .await?;
4731
4732 let final_status = sidecar
4735 .status
4736 .clone()
4737 .unwrap_or_else(|| "SUCCEEDED".to_string());
4738 let response = Self::transaction_response(version, &transaction, Some(sidecar));
4739 Ok(AlterTransactionResponse {
4740 status: final_status,
4741 properties: response.properties,
4742 ..Default::default()
4743 })
4744 }
4745
4746 async fn create_table_scalar_index(
4747 &self,
4748 request: CreateTableIndexRequest,
4749 ) -> Result<CreateTableScalarIndexResponse> {
4750 self.record_op("create_table_scalar_index");
4751 let index_type = Self::parse_index_type(&request.index_type)?;
4752 if !index_type.is_scalar() {
4753 return Err(NamespaceError::InvalidInput {
4754 message: format!(
4755 "create_table_scalar_index only supports scalar index types, got {}",
4756 request.index_type
4757 ),
4758 }
4759 .into());
4760 }
4761
4762 let response = self.create_table_index(request).await?;
4763 Ok(CreateTableScalarIndexResponse {
4764 transaction_id: response.transaction_id,
4765 ..Default::default()
4766 })
4767 }
4768
4769 async fn drop_table_index(
4770 &self,
4771 request: DropTableIndexRequest,
4772 ) -> Result<DropTableIndexResponse> {
4773 self.record_op("drop_table_index");
4774 let table_uri = self.resolve_table_location(&request.id).await?;
4775 let index_name = request.index_name.as_deref().ok_or_else(|| {
4776 lance_core::Error::from(NamespaceError::InvalidInput {
4777 message: "Index name is required for drop_table_index".to_string(),
4778 })
4779 })?;
4780 let mut dataset = self
4781 .load_dataset(&table_uri, None, "drop_table_index")
4782 .await?;
4783 let metadatas = dataset
4784 .load_indices_by_name(index_name)
4785 .await
4786 .map_err(|e| {
4787 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4788 message: format!(
4789 "Failed to load index '{}' before dropping it from table '{}': {}",
4790 index_name, table_uri, e
4791 ),
4792 })
4793 })?;
4794 if metadatas.first().is_some_and(is_system_index) {
4795 return Err(NamespaceError::Unsupported {
4796 message: format!(
4797 "System index '{}' cannot be dropped via this API",
4798 index_name
4799 ),
4800 }
4801 .into());
4802 }
4803
4804 dataset.drop_index(index_name).await.map_err(|e| {
4805 lance_core::Error::from(NamespaceError::TableIndexNotFound {
4806 message: format!(
4807 "Failed to drop index '{}' from table '{}': {}",
4808 index_name, table_uri, e
4809 ),
4810 })
4811 })?;
4812
4813 let transaction_id = dataset
4814 .read_transaction()
4815 .await
4816 .map_err(|e| {
4817 lance_core::Error::from(NamespaceError::Internal {
4818 message: format!(
4819 "Failed to read committed transaction after dropping index '{}' from '{}': {}",
4820 index_name, table_uri, e
4821 ),
4822 })
4823 })?
4824 .map(|transaction| transaction.uuid);
4825
4826 Ok(DropTableIndexResponse {
4827 transaction_id,
4828 ..Default::default()
4829 })
4830 }
4831
4832 async fn list_all_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
4833 let mut tables = self.list_directory_tables().await?;
4836 tables = self
4837 .filter_declared_tables(tables, request.include_declared.unwrap_or(true))
4838 .await?;
4839 Self::apply_pagination(&mut tables, request.page_token, request.limit);
4840 Ok(ListTablesResponse::new(tables))
4841 }
4842
4843 async fn restore_table(&self, request: RestoreTableRequest) -> Result<RestoreTableResponse> {
4844 let version = request.version;
4845 if version < 0 {
4846 return Err(Error::invalid_input_source(
4847 format!(
4848 "Table version for restore_table must be non-negative, got {}",
4849 version
4850 )
4851 .into(),
4852 ));
4853 }
4854
4855 let branch = Self::normalized_branch(request.branch.as_deref())?;
4856 let table_uri = self.resolve_table_location(&request.id).await?;
4857 let mut dataset = match branch {
4858 Some(branch) => self.open_validated_branch(&table_uri, branch).await?,
4859 None => self.load_dataset(&table_uri, None, "restore_table").await?,
4860 };
4861
4862 dataset = dataset
4863 .checkout_version(version as u64)
4864 .await
4865 .map_err(|e| {
4866 Error::namespace_source(
4867 format!(
4868 "Failed to checkout version {} for restore at '{}': {}",
4869 version, table_uri, e
4870 )
4871 .into(),
4872 )
4873 })?;
4874
4875 dataset.restore().await.map_err(|e| {
4876 Error::namespace_source(
4877 format!(
4878 "Failed to restore table at '{}' to version {}: {}",
4879 table_uri, version, e
4880 )
4881 .into(),
4882 )
4883 })?;
4884
4885 let transaction_id = dataset
4886 .read_transaction()
4887 .await
4888 .map_err(|e| {
4889 Error::namespace_source(
4890 format!(
4891 "Failed to read transaction after restoring '{}': {}",
4892 table_uri, e
4893 )
4894 .into(),
4895 )
4896 })?
4897 .map(|t| t.uuid);
4898
4899 Ok(RestoreTableResponse {
4900 transaction_id,
4901 ..Default::default()
4902 })
4903 }
4904
4905 async fn update_table_schema_metadata(
4906 &self,
4907 request: UpdateTableSchemaMetadataRequest,
4908 ) -> Result<UpdateTableSchemaMetadataResponse> {
4909 let table_uri = self.resolve_table_location(&request.id).await?;
4910 let mut dataset = self
4911 .load_dataset(&table_uri, None, "update_table_schema_metadata")
4912 .await?;
4913
4914 let new_metadata = request.metadata.unwrap_or_default();
4915 let updated_metadata = dataset
4916 .update_schema_metadata(new_metadata.iter().map(|(k, v)| (k.as_str(), v.as_str())))
4917 .await
4918 .map_err(|e| {
4919 Error::namespace_source(
4920 format!(
4921 "Failed to update schema metadata for table at '{}': {}",
4922 table_uri, e
4923 )
4924 .into(),
4925 )
4926 })?;
4927
4928 let transaction_id = dataset
4929 .read_transaction()
4930 .await
4931 .map_err(|e| {
4932 Error::namespace_source(
4933 format!(
4934 "Failed to read transaction after updating metadata for '{}': {}",
4935 table_uri, e
4936 )
4937 .into(),
4938 )
4939 })?
4940 .map(|t| t.uuid);
4941
4942 Ok(UpdateTableSchemaMetadataResponse {
4943 metadata: Some(updated_metadata),
4944 transaction_id,
4945 ..Default::default()
4946 })
4947 }
4948
4949 async fn get_table_stats(
4950 &self,
4951 request: GetTableStatsRequest,
4952 ) -> Result<GetTableStatsResponse> {
4953 let table_uri = self.resolve_table_location(&request.id).await?;
4954 let dataset = Arc::new(
4955 self.load_dataset(&table_uri, None, "get_table_stats")
4956 .await?,
4957 );
4958
4959 let data_stats = dataset.calculate_data_stats().await.map_err(|e| {
4961 Error::namespace_source(
4962 format!(
4963 "Failed to calculate data statistics for table at '{}': {}",
4964 table_uri, e
4965 )
4966 .into(),
4967 )
4968 })?;
4969 let total_bytes: i64 = data_stats
4970 .fields
4971 .iter()
4972 .map(|f| f.bytes_on_disk as i64)
4973 .sum();
4974
4975 let fragment_row_futures: Vec<_> = dataset
4977 .get_fragments()
4978 .into_iter()
4979 .map(|f| async move { f.physical_rows().await })
4980 .collect();
4981 let fragment_row_results = futures::future::join_all(fragment_row_futures).await;
4982 let mut fragment_row_counts: Vec<i64> = fragment_row_results
4983 .into_iter()
4984 .filter_map(|r| r.ok())
4985 .map(|r| r as i64)
4986 .collect();
4987
4988 let num_fragments = fragment_row_counts.len() as i64;
4989 let num_rows: i64 = fragment_row_counts.iter().sum();
4990
4991 const SMALL_FRAGMENT_THRESHOLD: i64 = 1024 * 1024;
4994 let num_small_fragments = fragment_row_counts
4995 .iter()
4996 .filter(|&&r| r < SMALL_FRAGMENT_THRESHOLD)
4997 .count() as i64;
4998
4999 fragment_row_counts.sort_unstable();
5001 let lengths = if fragment_row_counts.is_empty() {
5002 FragmentSummary::new(0, 0, 0, 0, 0, 0, 0)
5003 } else {
5004 let len = fragment_row_counts.len();
5005 let min = fragment_row_counts[0];
5006 let max = fragment_row_counts[len - 1];
5007 let mean = num_rows / num_fragments;
5008 let pct = |p: f64| fragment_row_counts[((len - 1) as f64 * p) as usize];
5009 FragmentSummary::new(min, max, mean, pct(0.25), pct(0.50), pct(0.75), pct(0.99))
5010 };
5011
5012 let indices = dataset.load_indices().await.map_err(|e| {
5014 Error::namespace_source(
5015 format!("Failed to load indices for table at '{}': {}", table_uri, e).into(),
5016 )
5017 })?;
5018 let num_indices = indices.iter().filter(|m| !is_system_index(m)).count() as i64;
5019
5020 let fragment_stats = FragmentStats::new(num_fragments, num_small_fragments, lengths);
5021 Ok(GetTableStatsResponse::new(
5022 total_bytes,
5023 num_rows,
5024 num_indices,
5025 fragment_stats,
5026 ))
5027 }
5028
5029 async fn explain_table_query_plan(
5030 &self,
5031 request: ExplainTableQueryPlanRequest,
5032 ) -> Result<String> {
5033 let table_uri = self.resolve_table_location(&request.id).await?;
5034 let dataset = self
5035 .load_dataset(
5036 &table_uri,
5037 request.query.version,
5038 "explain_table_query_plan",
5039 )
5040 .await?;
5041 let verbose = request.verbose.unwrap_or(false);
5042
5043 let mut scanner = dataset.scan();
5044 Self::apply_query_params_to_scanner(
5045 &mut scanner,
5046 request.query.filter.as_deref(),
5047 request.query.columns.as_deref(),
5048 request.query.vector_column.as_deref(),
5049 &request.query.vector,
5050 request.query.k,
5051 request.query.offset,
5052 request.query.prefilter,
5053 request.query.bypass_vector_index,
5054 request.query.nprobes,
5055 request.query.minimum_nprobes,
5056 request.query.maximum_nprobes,
5057 request.query.ef,
5058 request.query.refine_factor,
5059 request.query.distance_type.as_deref(),
5060 request.query.fast_search,
5061 request.query.with_row_id,
5062 request.query.lower_bound,
5063 request.query.upper_bound,
5064 "explain_table_query_plan",
5065 )?;
5066
5067 scanner.explain_plan(verbose).await.map_err(|e| {
5068 Error::namespace_source(
5069 format!(
5070 "Failed to explain query plan for table at '{}': {}",
5071 table_uri, e
5072 )
5073 .into(),
5074 )
5075 })
5076 }
5077
5078 async fn analyze_table_query_plan(
5079 &self,
5080 request: AnalyzeTableQueryPlanRequest,
5081 ) -> Result<String> {
5082 let table_uri = self.resolve_table_location(&request.id).await?;
5083 let dataset = self
5084 .load_dataset(&table_uri, request.version, "analyze_table_query_plan")
5085 .await?;
5086
5087 let mut scanner = dataset.scan();
5088 Self::apply_query_params_to_scanner(
5089 &mut scanner,
5090 request.filter.as_deref(),
5091 request.columns.as_deref(),
5092 request.vector_column.as_deref(),
5093 &request.vector,
5094 request.k,
5095 request.offset,
5096 request.prefilter,
5097 request.bypass_vector_index,
5098 request.nprobes,
5099 request.minimum_nprobes,
5100 request.maximum_nprobes,
5101 request.ef,
5102 request.refine_factor,
5103 request.distance_type.as_deref(),
5104 request.fast_search,
5105 request.with_row_id,
5106 request.lower_bound,
5107 request.upper_bound,
5108 "analyze_table_query_plan",
5109 )?;
5110
5111 scanner.analyze_plan().await.map_err(|e| {
5112 Error::namespace_source(
5113 format!(
5114 "Failed to analyze query plan for table at '{}': {}",
5115 table_uri, e
5116 )
5117 .into(),
5118 )
5119 })
5120 }
5121
5122 async fn count_table_rows(&self, request: CountTableRowsRequest) -> Result<i64> {
5123 self.record_op("count_table_rows");
5124 let table_uri = self.resolve_table_location(&request.id).await?;
5125 let dataset = self
5126 .load_dataset(&table_uri, request.version, "count_table_rows")
5127 .await?;
5128
5129 let count =
5130 dataset
5131 .count_rows(request.predicate)
5132 .await
5133 .map_err(|e| NamespaceError::Internal {
5134 message: format!("Failed to count rows for table at '{}': {:?}", table_uri, e),
5135 })?;
5136
5137 Ok(count as i64)
5138 }
5139
5140 async fn insert_into_table(
5141 &self,
5142 request: InsertIntoTableRequest,
5143 request_data: Bytes,
5144 ) -> Result<InsertIntoTableResponse> {
5145 self.record_op("insert_into_table");
5146 let table_uri = self.resolve_table_location(&request.id).await?;
5147 let (reader, _num_rows) =
5148 Self::ipc_reader_from_request_data(&request_data, "insert_into_table")?;
5149
5150 let mode = match request.mode.as_deref() {
5151 Some(m) if m.eq_ignore_ascii_case("overwrite") => WriteMode::Overwrite,
5152 Some(m) if m.eq_ignore_ascii_case("append") => WriteMode::Append,
5153 None => WriteMode::Append,
5154 Some(m) => {
5155 return Err(lance_namespace::error::NamespaceError::InvalidInput {
5156 message: format!(
5157 "Unsupported write mode '{}'. Supported modes are: 'append', 'overwrite'",
5158 m
5159 ),
5160 }
5161 .into());
5162 }
5163 };
5164
5165 if !self.table_uri_has_actual_manifests(&table_uri).await? {
5166 self.write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
5167 .await?;
5168 } else {
5169 self.write_reader_to_table(&table_uri, reader, mode, None)
5170 .await?;
5171 }
5172
5173 Ok(InsertIntoTableResponse {
5174 transaction_id: None,
5175 ..Default::default()
5176 })
5177 }
5178
5179 async fn merge_insert_into_table(
5180 &self,
5181 request: MergeInsertIntoTableRequest,
5182 request_data: Bytes,
5183 ) -> Result<MergeInsertIntoTableResponse> {
5184 self.record_op("merge_insert_into_table");
5185 let table_uri = self.resolve_table_location(&request.id).await?;
5186 let on = merge_insert_on_columns(request.on.as_deref(), "merge_insert_into_table")?;
5187
5188 let table_has_manifests = self.table_uri_has_actual_manifests(&table_uri).await?;
5189 let (reader, num_rows) =
5190 Self::ipc_reader_from_request_data(&request_data, "merge_insert_into_table")?;
5191
5192 if !table_has_manifests {
5193 let dataset = self
5194 .write_reader_to_table(&table_uri, reader, WriteMode::Create, None)
5195 .await?;
5196 let version = dataset.version().version as i64;
5197 return Ok(MergeInsertIntoTableResponse {
5198 transaction_id: None,
5199 num_updated_rows: Some(0),
5200 num_inserted_rows: Some(num_rows as i64),
5201 num_deleted_rows: Some(0),
5202 version: Some(version),
5203 ..Default::default()
5204 });
5205 }
5206
5207 let dataset = Arc::new(
5208 self.load_dataset(&table_uri, None, "merge_insert_into_table")
5209 .await?,
5210 );
5211
5212 let mut merge_builder =
5213 MergeInsertBuilder::try_new(dataset.clone(), on.to_vec()).map_err(|e| {
5214 lance_core::Error::from(NamespaceError::InvalidInput {
5215 message: format!("Failed to create merge_insert_into_table builder: {}", e),
5216 })
5217 })?;
5218
5219 if let Some(filter) = request.when_matched_update_all_filt.as_deref() {
5220 let behavior = WhenMatched::update_if(dataset.as_ref(), filter).map_err(|e| {
5221 lance_core::Error::from(NamespaceError::InvalidInput {
5222 message: format!(
5223 "Invalid when_matched_update_all_filt for merge_insert_into_table: {}",
5224 e
5225 ),
5226 })
5227 })?;
5228 merge_builder.when_matched(behavior);
5229 } else if request.when_matched_update_all.unwrap_or(false) {
5230 merge_builder.when_matched(WhenMatched::UpdateAll);
5231 }
5232
5233 if matches!(request.when_not_matched_insert_all, Some(false)) {
5234 merge_builder.when_not_matched(WhenNotMatched::DoNothing);
5235 } else {
5236 merge_builder.when_not_matched(WhenNotMatched::InsertAll);
5237 }
5238
5239 if let Some(filter) = request.when_not_matched_by_source_delete_filt.as_deref() {
5240 let behavior = WhenNotMatchedBySource::delete_if(dataset.as_ref(), filter).map_err(|e| {
5241 lance_core::Error::from(NamespaceError::InvalidInput {
5242 message: format!(
5243 "Invalid when_not_matched_by_source_delete_filt for merge_insert_into_table: {}",
5244 e
5245 ),
5246 })
5247 })?;
5248 merge_builder.when_not_matched_by_source(behavior);
5249 } else if request.when_not_matched_by_source_delete.unwrap_or(false) {
5250 merge_builder.when_not_matched_by_source(WhenNotMatchedBySource::Delete);
5251 }
5252
5253 if let Some(use_index) = request.use_index {
5254 merge_builder.use_index(use_index);
5255 }
5256
5257 let (dataset, stats) = merge_builder
5258 .try_build()
5259 .map_err(|e| {
5260 lance_core::Error::from(NamespaceError::InvalidInput {
5261 message: format!("Failed to build merge_insert_into_table job: {}", e),
5262 })
5263 })?
5264 .execute_reader(reader)
5265 .await
5266 .map_err(|e| Self::map_mutation_error(e, "merge_insert_into_table", &table_uri))?;
5267
5268 Ok(MergeInsertIntoTableResponse {
5269 transaction_id: None,
5270 num_updated_rows: Some(stats.num_updated_rows as i64),
5271 num_inserted_rows: Some(stats.num_inserted_rows as i64),
5272 num_deleted_rows: Some(stats.num_deleted_rows as i64),
5273 version: Some(dataset.version().version as i64),
5274 ..Default::default()
5275 })
5276 }
5277
5278 async fn update_table(&self, request: UpdateTableRequest) -> Result<UpdateTableResponse> {
5279 self.record_op("update_table");
5280
5281 if request.updates.is_empty() {
5282 return Err(NamespaceError::InvalidInput {
5283 message: "update_table requires at least one [column, expression] pair".to_string(),
5284 }
5285 .into());
5286 }
5287
5288 let mut seen_columns: HashMap<String, ()> = HashMap::with_capacity(request.updates.len());
5291 for (idx, pair) in request.updates.iter().enumerate() {
5292 if pair.len() != 2 {
5293 return Err(NamespaceError::InvalidInput {
5294 message: format!(
5295 "update_table updates[{}] must be a [column, expression] pair, got {} elements",
5296 idx,
5297 pair.len()
5298 ),
5299 }
5300 .into());
5301 }
5302 let column = &pair[0];
5303 if column.trim().is_empty() {
5304 return Err(NamespaceError::InvalidInput {
5305 message: format!("update_table updates[{}] has an empty column name", idx),
5306 }
5307 .into());
5308 }
5309 if seen_columns.insert(column.clone(), ()).is_some() {
5310 return Err(NamespaceError::InvalidInput {
5311 message: format!(
5312 "update_table cannot update column '{}' more than once",
5313 column
5314 ),
5315 }
5316 .into());
5317 }
5318 }
5319
5320 let table_uri = self.resolve_table_location(&request.id).await?;
5321 let dataset = Arc::new(self.load_dataset(&table_uri, None, "update_table").await?);
5322
5323 let mut builder = UpdateBuilder::new(dataset);
5324 for pair in &request.updates {
5325 builder = builder.set(&pair[0], &pair[1]).map_err(|e| {
5327 lance_core::Error::from(NamespaceError::InvalidInput {
5328 message: format!("Invalid update expression for column '{}': {}", pair[0], e),
5329 })
5330 })?;
5331 }
5332 if let Some(predicate) = request.predicate.as_deref()
5333 && !predicate.trim().is_empty()
5334 {
5335 builder = builder.update_where(predicate).map_err(|e| {
5336 lance_core::Error::from(NamespaceError::InvalidInput {
5337 message: format!("Invalid update_table predicate '{}': {}", predicate, e),
5338 })
5339 })?;
5340 }
5341
5342 let job = builder.build().map_err(|e| {
5343 lance_core::Error::from(NamespaceError::InvalidInput {
5344 message: format!("Failed to build update_table job: {}", e),
5345 })
5346 })?;
5347
5348 let result = job
5349 .execute()
5350 .await
5351 .map_err(|e| Self::map_mutation_error(e, "update_table", &table_uri))?;
5352
5353 let version = result.new_dataset.version().version as i64;
5354 Ok(UpdateTableResponse {
5355 transaction_id: None,
5356 updated_rows: result.rows_updated as i64,
5357 version,
5358 properties: None,
5359 ..Default::default()
5360 })
5361 }
5362
5363 async fn delete_from_table(
5364 &self,
5365 request: DeleteFromTableRequest,
5366 ) -> Result<DeleteFromTableResponse> {
5367 self.record_op("delete_from_table");
5368
5369 if request.predicate.trim().is_empty() {
5370 return Err(NamespaceError::InvalidInput {
5371 message: "delete_from_table requires a non-empty predicate".to_string(),
5372 }
5373 .into());
5374 }
5375
5376 let table_uri = self.resolve_table_location(&request.id).await?;
5377 let mut dataset = self
5378 .load_dataset(&table_uri, None, "delete_from_table")
5379 .await?;
5380
5381 let result = dataset
5382 .delete(&request.predicate)
5383 .await
5384 .map_err(|e| Self::map_mutation_error(e, "delete_from_table", &table_uri))?;
5385
5386 Ok(DeleteFromTableResponse {
5387 transaction_id: None,
5388 version: Some(result.new_dataset.version().version as i64),
5389 ..Default::default()
5390 })
5391 }
5392
5393 async fn query_table(&self, request: QueryTableRequest) -> Result<Bytes> {
5394 use arrow::ipc::writer::FileWriter;
5395
5396 self.record_op("query_table");
5397 let table_uri = self.resolve_table_location(&request.id).await?;
5398 let dataset = self
5399 .load_dataset(&table_uri, request.version, "query_table")
5400 .await?;
5401
5402 let mut scanner = dataset.scan();
5404
5405 let has_vector_query = request
5408 .vector
5409 .single_vector
5410 .as_ref()
5411 .map(|sv| !sv.is_empty())
5412 .unwrap_or(false)
5413 || request
5414 .vector
5415 .multi_vector
5416 .as_ref()
5417 .map(|mv| !mv.is_empty())
5418 .unwrap_or(false);
5419
5420 if let Some(prefilter) = request.prefilter {
5422 scanner.prefilter(prefilter);
5423 }
5424
5425 if has_vector_query {
5427 let vector_column = request.vector_column.as_deref().unwrap_or("vector");
5428
5429 let query_vector: Vec<f32> = request
5431 .vector
5432 .single_vector
5433 .clone()
5434 .or_else(|| {
5435 request
5436 .vector
5437 .multi_vector
5438 .as_ref()
5439 .and_then(|mv| mv.first().cloned())
5440 })
5441 .unwrap_or_default();
5442
5443 if !query_vector.is_empty() {
5444 let k = if request.k > 0 {
5445 request.k as usize
5446 } else {
5447 10
5448 };
5449 let query_array = Float32Array::from(query_vector);
5450 scanner
5451 .nearest(vector_column, &query_array, k)
5452 .map_err(|e| NamespaceError::InvalidInput {
5453 message: format!("Invalid vector search: {:?}", e),
5454 })?;
5455
5456 if let Some(ref distance_type) = request.distance_type {
5458 let metric = match distance_type.to_lowercase().as_str() {
5459 "l2" | "euclidean" => MetricType::L2,
5460 "cosine" => MetricType::Cosine,
5461 "dot" | "inner_product" => MetricType::Dot,
5462 "hamming" => MetricType::Hamming,
5463 _ => {
5464 return Err(NamespaceError::InvalidInput {
5465 message: format!("Unknown distance type: {}", distance_type),
5466 }
5467 .into());
5468 }
5469 };
5470 scanner.distance_metric(metric);
5471 }
5472
5473 apply_probe_bounds(
5474 &mut scanner,
5475 request.nprobes,
5476 request.minimum_nprobes,
5477 request.maximum_nprobes,
5478 )?;
5479
5480 if let Some(ef) = request.ef {
5482 scanner.ef(ef as usize);
5483 }
5484
5485 if let Some(refine_factor) = request.refine_factor {
5487 scanner.refine(refine_factor as u32);
5488 }
5489
5490 if request.lower_bound.is_some() || request.upper_bound.is_some() {
5492 scanner.distance_range(request.lower_bound, request.upper_bound);
5493 }
5494
5495 if let Some(bypass) = request.bypass_vector_index {
5497 scanner.use_index(!bypass);
5498 }
5499
5500 if request.fast_search == Some(true) {
5502 scanner.fast_search();
5503 }
5504 }
5505 }
5506
5507 if let Some(ref fts_query) = request.full_text_query {
5509 if let Some(ref string_query) = fts_query.string_query {
5511 let mut fts = FullTextSearchQuery::new(string_query.query.clone());
5512
5513 if let Some(ref columns) = string_query.columns
5515 && !columns.is_empty()
5516 {
5517 fts = fts
5518 .with_columns(columns)
5519 .map_err(|e| NamespaceError::InvalidInput {
5520 message: format!("Invalid FTS columns: {:?}", e),
5521 })?;
5522 }
5523
5524 scanner
5525 .full_text_search(fts)
5526 .map_err(|e| NamespaceError::InvalidInput {
5527 message: format!("Invalid full text search: {:?}", e),
5528 })?;
5529 } else if let Some(ref structured_query) = fts_query.structured_query {
5530 let engine_query = build_engine_fts_query(&structured_query.query)?;
5532 let fts = FullTextSearchQuery::new_query(engine_query);
5533 scanner
5534 .full_text_search(fts)
5535 .map_err(|e| NamespaceError::InvalidInput {
5536 message: format!("Invalid full text search: {:?}", e),
5537 })?;
5538 }
5539 }
5540
5541 if let Some(ref columns) = request.columns {
5543 if let Some(ref column_names) = columns.column_names
5544 && !column_names.is_empty()
5545 {
5546 scanner
5547 .project(column_names)
5548 .map_err(|e| NamespaceError::InvalidInput {
5549 message: format!("Invalid column projection: {:?}", e),
5550 })?;
5551 } else if let Some(ref column_aliases) = columns.column_aliases
5552 && !column_aliases.is_empty()
5553 {
5554 let transform_pairs: Vec<(String, String)> = column_aliases
5556 .iter()
5557 .map(|(alias, sql)| (alias.clone(), sql.clone()))
5558 .collect();
5559 scanner
5560 .project_with_transform(
5561 &transform_pairs
5562 .iter()
5563 .map(|(a, s)| (a.as_str(), s.as_str()))
5564 .collect::<Vec<_>>(),
5565 )
5566 .map_err(|e| NamespaceError::InvalidInput {
5567 message: format!("Invalid column alias expression: {:?}", e),
5568 })?;
5569 }
5570 }
5571
5572 if let Some(ref filter) = request.filter
5574 && !filter.is_empty()
5575 {
5576 scanner
5577 .filter(filter)
5578 .map_err(|e| NamespaceError::InvalidInput {
5579 message: format!("Invalid filter expression: {:?}", e),
5580 })?;
5581 }
5582
5583 if request.with_row_id == Some(true) {
5585 scanner.with_row_id();
5586 }
5587
5588 if !has_vector_query && request.k > 0 {
5592 let offset = request.offset.map(|o| o as i64);
5593 scanner.limit(Some(request.k as i64), offset).map_err(|e| {
5594 NamespaceError::InvalidInput {
5595 message: format!("Invalid limit/offset: {:?}", e),
5596 }
5597 })?;
5598 } else if has_vector_query && request.offset.is_some() {
5599 let offset = request.offset.map(|o| o as i64);
5601 scanner
5602 .limit(None, offset)
5603 .map_err(|e| NamespaceError::InvalidInput {
5604 message: format!("Invalid offset: {:?}", e),
5605 })?;
5606 }
5607
5608 let batch = scanner
5610 .try_into_batch()
5611 .await
5612 .map_err(|e| NamespaceError::Internal {
5613 message: format!("Failed to execute query: {:?}", e),
5614 })?;
5615
5616 let schema = batch.schema();
5618 let mut buffer = Vec::new();
5619 {
5620 let mut writer = FileWriter::try_new(&mut buffer, &schema).map_err(|e| {
5621 NamespaceError::Internal {
5622 message: format!("Failed to create IPC writer: {:?}", e),
5623 }
5624 })?;
5625 writer.write(&batch).map_err(|e| NamespaceError::Internal {
5626 message: format!("Failed to write batch to IPC: {:?}", e),
5627 })?;
5628 writer.finish().map_err(|e| NamespaceError::Internal {
5629 message: format!("Failed to finish IPC writer: {:?}", e),
5630 })?;
5631 }
5632
5633 Ok(Bytes::from(buffer))
5634 }
5635
5636 async fn list_table_tags(
5637 &self,
5638 request: ListTableTagsRequest,
5639 ) -> Result<ListTableTagsResponse> {
5640 self.record_op("list_table_tags");
5641 let table_uri = self.resolve_table_location(&request.id).await?;
5642 let dataset = self
5643 .load_dataset(&table_uri, None, "list_table_tags")
5644 .await?;
5645
5646 let raw_tags = dataset.tags().list().await.map_err(|e| {
5647 lance_core::Error::from(NamespaceError::Internal {
5648 message: format!("Failed to list tags for table at '{}': {}", table_uri, e),
5649 })
5650 })?;
5651
5652 let tags = raw_tags
5653 .into_iter()
5654 .map(|(name, contents)| {
5655 let mut tag_model =
5656 ModelTagContents::new(contents.version as i64, contents.manifest_size as i64);
5657 tag_model.branch = contents.branch;
5658 (name, tag_model)
5659 })
5660 .collect();
5661
5662 Ok(ListTableTagsResponse {
5663 tags,
5664 page_token: None,
5665 ..Default::default()
5666 })
5667 }
5668
5669 async fn get_table_tag_version(
5670 &self,
5671 request: GetTableTagVersionRequest,
5672 ) -> Result<GetTableTagVersionResponse> {
5673 self.record_op("get_table_tag_version");
5674 if request.tag.is_empty() {
5675 return Err(NamespaceError::InvalidInput {
5676 message: "tag name must not be empty for get_table_tag_version".to_string(),
5677 }
5678 .into());
5679 }
5680
5681 let table_uri = self.resolve_table_location(&request.id).await?;
5682 let dataset = self
5683 .load_dataset(&table_uri, None, "get_table_tag_version")
5684 .await?;
5685
5686 let contents = dataset
5687 .tags()
5688 .get(&request.tag)
5689 .await
5690 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5691
5692 Ok(GetTableTagVersionResponse {
5693 version: contents.version as i64,
5694 branch: contents.branch,
5695 ..Default::default()
5696 })
5697 }
5698
5699 async fn create_table_tag(
5700 &self,
5701 request: CreateTableTagRequest,
5702 ) -> Result<CreateTableTagResponse> {
5703 self.record_op("create_table_tag");
5704 if request.tag.is_empty() {
5705 return Err(NamespaceError::InvalidInput {
5706 message: "tag name must not be empty for create_table_tag".to_string(),
5707 }
5708 .into());
5709 }
5710 if request.version <= 0 {
5711 return Err(NamespaceError::InvalidInput {
5712 message: format!(
5713 "tag version must be a positive integer, got {} for create_table_tag",
5714 request.version
5715 ),
5716 }
5717 .into());
5718 }
5719
5720 let table_uri = self.resolve_table_location(&request.id).await?;
5721 let dataset = self
5722 .load_dataset(&table_uri, None, "create_table_tag")
5723 .await?;
5724
5725 dataset
5726 .tags()
5727 .create(&request.tag, request.version as u64)
5728 .await
5729 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5730
5731 Ok(CreateTableTagResponse {
5732 transaction_id: None,
5733 ..Default::default()
5734 })
5735 }
5736
5737 async fn delete_table_tag(
5738 &self,
5739 request: DeleteTableTagRequest,
5740 ) -> Result<DeleteTableTagResponse> {
5741 self.record_op("delete_table_tag");
5742 if request.tag.is_empty() {
5743 return Err(NamespaceError::InvalidInput {
5744 message: "tag name must not be empty for delete_table_tag".to_string(),
5745 }
5746 .into());
5747 }
5748
5749 let table_uri = self.resolve_table_location(&request.id).await?;
5750 let dataset = self
5751 .load_dataset(&table_uri, None, "delete_table_tag")
5752 .await?;
5753
5754 dataset
5755 .tags()
5756 .delete(&request.tag)
5757 .await
5758 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5759
5760 Ok(DeleteTableTagResponse {
5761 transaction_id: None,
5762 ..Default::default()
5763 })
5764 }
5765
5766 async fn update_table_tag(
5767 &self,
5768 request: UpdateTableTagRequest,
5769 ) -> Result<UpdateTableTagResponse> {
5770 self.record_op("update_table_tag");
5771 if request.tag.is_empty() {
5772 return Err(NamespaceError::InvalidInput {
5773 message: "tag name must not be empty for update_table_tag".to_string(),
5774 }
5775 .into());
5776 }
5777 if request.version <= 0 {
5778 return Err(NamespaceError::InvalidInput {
5779 message: format!(
5780 "tag version must be a positive integer, got {} for update_table_tag",
5781 request.version
5782 ),
5783 }
5784 .into());
5785 }
5786
5787 let table_uri = self.resolve_table_location(&request.id).await?;
5788 let dataset = self
5789 .load_dataset(&table_uri, None, "update_table_tag")
5790 .await?;
5791
5792 dataset
5793 .tags()
5794 .update(&request.tag, request.version as u64)
5795 .await
5796 .map_err(|e| Self::map_tag_error(e, &request.tag, &table_uri))?;
5797
5798 Ok(UpdateTableTagResponse {
5799 transaction_id: None,
5800 ..Default::default()
5801 })
5802 }
5803
5804 async fn create_table_branch(
5805 &self,
5806 request: CreateTableBranchRequest,
5807 ) -> Result<CreateTableBranchResponse> {
5808 self.record_op("create_table_branch");
5809 if request.name.is_empty() {
5810 return Err(NamespaceError::InvalidInput {
5811 message: "branch name must not be empty for create_table_branch".to_string(),
5812 }
5813 .into());
5814 }
5815 let from_version = match request.from_version {
5816 Some(v) if v <= 0 => {
5817 return Err(NamespaceError::InvalidInput {
5818 message: format!(
5819 "from_version must be a positive integer, got {} for create_table_branch",
5820 v
5821 ),
5822 }
5823 .into());
5824 }
5825 Some(v) => Some(v as u64),
5826 None => None,
5827 };
5828
5829 let table_uri = self.resolve_table_location(&request.id).await?;
5830 let mut dataset = self
5831 .load_dataset(&table_uri, None, "create_table_branch")
5832 .await?;
5833
5834 if dataset.branches().get(&request.name).await.is_ok() {
5839 return Err(NamespaceError::TableBranchAlreadyExists {
5840 message: format!("branch '{}' for table at '{}'", request.name, table_uri),
5841 }
5842 .into());
5843 }
5844
5845 dataset
5846 .create_branch(
5847 &request.name,
5848 (request.from_branch.as_deref(), from_version),
5849 None,
5850 )
5851 .await
5852 .map_err(|e| {
5853 if matches!(e, lance_core::Error::DatasetNotFound { .. }) {
5856 NamespaceError::InvalidInput {
5857 message: format!(
5858 "from_branch/from_version for branch '{}' refers to a source that does not exist: {}",
5859 request.name, e
5860 ),
5861 }
5862 .into()
5863 } else {
5864 Self::map_branch_error(e, &request.name, &table_uri)
5865 }
5866 })?;
5867
5868 Ok(CreateTableBranchResponse {
5869 transaction_id: None,
5870 ..Default::default()
5871 })
5872 }
5873
5874 async fn list_table_branches(
5875 &self,
5876 request: ListTableBranchesRequest,
5877 ) -> Result<ListTableBranchesResponse> {
5878 self.record_op("list_table_branches");
5879 let table_uri = self.resolve_table_location(&request.id).await?;
5880 let dataset = self
5881 .load_dataset(&table_uri, None, "list_table_branches")
5882 .await?;
5883
5884 let raw_branches = dataset.list_branches().await.map_err(|e| {
5885 lance_core::Error::from(NamespaceError::Internal {
5886 message: format!(
5887 "Failed to list branches for table at '{}': {}",
5888 table_uri, e
5889 ),
5890 })
5891 })?;
5892
5893 let branches = raw_branches
5894 .into_iter()
5895 .map(|(name, contents)| {
5896 let mut branch_model = ModelBranchContents::new(
5899 contents.parent_version as i64,
5900 contents.create_at as i64,
5901 contents.manifest_size as i64,
5902 );
5903 branch_model.parent_branch = contents.parent_branch;
5904 branch_model.metadata = if contents.metadata.is_empty() {
5905 None
5906 } else {
5907 Some(contents.metadata)
5908 };
5909 (name, branch_model)
5910 })
5911 .collect();
5912
5913 Ok(ListTableBranchesResponse {
5914 branches,
5915 page_token: None,
5916 ..Default::default()
5917 })
5918 }
5919
5920 async fn delete_table_branch(
5921 &self,
5922 request: DeleteTableBranchRequest,
5923 ) -> Result<DeleteTableBranchResponse> {
5924 self.record_op("delete_table_branch");
5925 if request.name.is_empty() {
5926 return Err(NamespaceError::InvalidInput {
5927 message: "branch name must not be empty for delete_table_branch".to_string(),
5928 }
5929 .into());
5930 }
5931
5932 let table_uri = self.resolve_table_location(&request.id).await?;
5933 let mut dataset = self
5934 .load_dataset(&table_uri, None, "delete_table_branch")
5935 .await?;
5936
5937 dataset
5938 .delete_branch(&request.name)
5939 .await
5940 .map_err(|e| match e {
5941 lance_core::Error::RefConflict { message } => NamespaceError::InvalidInput {
5942 message: format!(
5943 "branch '{}' for table at '{}': {}",
5944 request.name, table_uri, message
5945 ),
5946 }
5947 .into(),
5948 other => Self::map_branch_error(other, &request.name, &table_uri),
5949 })?;
5950
5951 Ok(DeleteTableBranchResponse {
5952 transaction_id: None,
5953 ..Default::default()
5954 })
5955 }
5956
5957 fn namespace_id(&self) -> String {
5958 format!("DirectoryNamespace {{ root: {:?} }}", self.root)
5959 }
5960}
5961
5962#[derive(Debug)]
5964pub(crate) enum MarkerFileError {
5965 AlreadyExists { description: String },
5967 Other { message: String },
5969}
5970
5971impl std::fmt::Display for MarkerFileError {
5972 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
5973 match self {
5974 Self::AlreadyExists { description } => {
5975 write!(f, "{} already exists", description)
5976 }
5977 Self::Other { message } => write!(f, "{}", message),
5978 }
5979 }
5980}
5981
5982pub(crate) async fn put_marker_file_atomic(
6003 object_store: &ObjectStore,
6004 path: &Path,
6005 file_description: &str,
6006) -> std::result::Result<(), MarkerFileError> {
6007 let staging_name = format!("lance-marker.staging.{}", uuid::Uuid::new_v4().simple());
6008 let path_str = path.as_ref();
6009 let staging_path = match path_str.rfind('/') {
6010 Some(idx) => Path::from(format!("{}/{}", &path_str[..idx], staging_name)),
6011 None => Path::from(staging_name.as_str()),
6012 };
6013
6014 object_store
6015 .inner
6016 .put(&staging_path, bytes::Bytes::from_static(b"reserved").into())
6017 .await
6018 .map_err(|e| MarkerFileError::Other {
6019 message: format!("Failed to stage {}: {:?}", file_description, e),
6020 })?;
6021
6022 let mut staging_consumed = false;
6025 let publish_result = match object_store
6026 .inner
6027 .rename_if_not_exists(&staging_path, path)
6028 .await
6029 {
6030 Ok(()) => {
6031 staging_consumed = true;
6032 Ok(())
6033 }
6034 Err(ObjectStoreError::NotImplemented { .. })
6035 | Err(ObjectStoreError::NotSupported { .. }) => {
6036 match object_store
6037 .inner
6038 .copy_if_not_exists(&staging_path, path)
6039 .await
6040 {
6041 Ok(()) => Ok(()),
6042 Err(ObjectStoreError::NotImplemented { .. })
6043 | Err(ObjectStoreError::NotSupported { .. }) => object_store
6044 .inner
6045 .put_opts(
6046 path,
6047 bytes::Bytes::from_static(b"reserved").into(),
6048 PutOptions {
6049 mode: PutMode::Create,
6050 ..Default::default()
6051 },
6052 )
6053 .await
6054 .map(|_| ()),
6055 Err(e) => Err(e),
6056 }
6057 }
6058 Err(e) => Err(e),
6059 };
6060
6061 if !staging_consumed {
6062 delete_staging_marker_best_effort(object_store, &staging_path).await;
6063 }
6064
6065 match publish_result {
6066 Ok(()) => Ok(()),
6067 Err(ObjectStoreError::AlreadyExists { .. })
6068 | Err(ObjectStoreError::Precondition { .. }) => Err(MarkerFileError::AlreadyExists {
6069 description: file_description.to_string(),
6070 }),
6071 Err(e) => Err(MarkerFileError::Other {
6072 message: format!("Failed to create {}: {:?}", file_description, e),
6073 }),
6074 }
6075}
6076
6077async fn delete_staging_marker_best_effort(object_store: &ObjectStore, staging_path: &Path) {
6081 const MAX_ATTEMPTS: u32 = 3;
6082 const BACKOFF_MS: [u64; 2] = [20, 50];
6083
6084 let mut last_err: Option<ObjectStoreError> = None;
6085 for attempt in 0..MAX_ATTEMPTS {
6086 match object_store.inner.delete(staging_path).await {
6087 Ok(()) => return,
6088 Err(ObjectStoreError::NotFound { .. }) => return,
6089 Err(e) => {
6090 last_err = Some(e);
6091 if let Some(&delay_ms) = BACKOFF_MS.get(attempt as usize) {
6092 tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await;
6093 }
6094 }
6095 }
6096 }
6097 if let Some(del_err) = last_err {
6098 log::warn!(
6099 "Failed to delete staging marker at '{}' after {} attempts: {:?}",
6100 staging_path,
6101 MAX_ATTEMPTS,
6102 del_err
6103 );
6104 }
6105}
6106
6107fn build_engine_fts_query(
6111 query: &lance_namespace::models::FtsQuery,
6112) -> std::result::Result<FtsQuery, NamespaceError> {
6113 if let Some(ref m) = query.r#match {
6114 Ok(FtsQuery::Match(build_engine_match_query(m)?))
6115 } else if let Some(ref p) = query.phrase {
6116 let mut phrase = PhraseQuery::new(p.terms.clone());
6117 if let Some(ref column) = p.column {
6118 phrase = phrase.with_column(Some(column.clone()));
6119 }
6120 if let Some(slop) = p.slop {
6121 phrase = phrase.with_slop(slop as u32);
6122 }
6123 Ok(FtsQuery::Phrase(phrase))
6124 } else if let Some(ref mm) = query.multi_match {
6125 let match_queries = mm
6126 .match_queries
6127 .iter()
6128 .map(build_engine_match_query)
6129 .collect::<std::result::Result<Vec<_>, _>>()?;
6130 Ok(FtsQuery::MultiMatch(MultiMatchQuery { match_queries }))
6131 } else if let Some(ref b) = query.boolean {
6132 let mut clauses: Vec<(Occur, FtsQuery)> = Vec::new();
6133 for clause in &b.must {
6134 clauses.push((Occur::Must, build_engine_fts_query(clause)?));
6135 }
6136 for clause in &b.should {
6137 clauses.push((Occur::Should, build_engine_fts_query(clause)?));
6138 }
6139 for clause in &b.must_not {
6140 clauses.push((Occur::MustNot, build_engine_fts_query(clause)?));
6141 }
6142 Ok(FtsQuery::Boolean(BooleanQuery::new(clauses)))
6143 } else if let Some(ref boost) = query.boost {
6144 let positive = build_engine_fts_query(&boost.positive)?;
6145 let negative = build_engine_fts_query(&boost.negative)?;
6146 Ok(FtsQuery::Boost(BoostQuery::new(
6147 positive,
6148 negative,
6149 boost.negative_boost,
6150 )))
6151 } else {
6152 Err(NamespaceError::InvalidInput {
6153 message: "structured_query.query must set exactly one of match, phrase, multi_match, \
6154 boolean, or boost"
6155 .to_string(),
6156 })
6157 }
6158}
6159
6160fn build_engine_match_query(
6161 m: &lance_namespace::models::MatchQuery,
6162) -> std::result::Result<MatchQuery, NamespaceError> {
6163 let mut match_query = MatchQuery::new(m.terms.clone());
6164 if let Some(ref column) = m.column {
6165 match_query = match_query.with_column(Some(column.clone()));
6166 }
6167 if let Some(boost) = m.boost {
6168 match_query = match_query.with_boost(boost);
6169 }
6170 if let Some(fuzziness) = m.fuzziness {
6171 match_query = match_query.with_fuzziness(Some(fuzziness as u32));
6172 }
6173 if let Some(max_expansions) = m.max_expansions {
6174 match_query = match_query.with_max_expansions(max_expansions as usize);
6175 }
6176 if let Some(ref operator) = m.operator {
6177 let op =
6178 Operator::try_from(operator.as_str()).map_err(|e| NamespaceError::InvalidInput {
6179 message: format!("Invalid FTS operator: {:?}", e),
6180 })?;
6181 match_query = match_query.with_operator(op);
6182 }
6183 if let Some(prefix_length) = m.prefix_length {
6184 match_query = match_query.with_prefix_length(prefix_length as u32);
6185 }
6186 Ok(match_query)
6187}
6188
6189#[cfg(test)]
6190mod tests {
6191 use super::*;
6192 use arrow_ipc::reader::{FileReader, StreamReader};
6193 use lance::index::vector::StageParams;
6194 use rstest::rstest;
6195
6196 fn build_ivf_rq_num_bits(num_bits: Option<i32>) -> Result<u8> {
6197 let mut request = CreateTableIndexRequest::new("vector".to_string(), "IVF_RQ".to_string());
6198 request.num_bits = num_bits;
6199
6200 let DirectoryIndexParams::Vector {
6201 index_type: IndexType::IvfRq,
6202 params,
6203 } = DirectoryNamespace::build_index_params(&request)?
6204 else {
6205 panic!("expected IVF_RQ vector index params");
6206 };
6207 match params.stages.as_slice() {
6208 [StageParams::Ivf(_), StageParams::RQ(rq)] => Ok(rq.num_bits),
6209 stages => panic!("expected IVF and RQ stages, got {stages:?}"),
6210 }
6211 }
6212
6213 #[rstest]
6214 #[case::omitted(None, 5)]
6215 #[case::explicit_one(Some(1), 1)]
6216 #[case::explicit_max(Some(9), 9)]
6217 fn test_build_index_params_ivf_rq_num_bits(
6218 #[case] requested: Option<i32>,
6219 #[case] expected: u8,
6220 ) {
6221 assert_eq!(build_ivf_rq_num_bits(requested).unwrap(), expected);
6222 }
6223
6224 #[rstest]
6225 #[case::negative(-1)]
6226 #[case::zero(0)]
6227 #[case::above_max(10)]
6228 #[case::conversion_overflow(i32::MAX)]
6229 fn test_build_index_params_rejects_invalid_ivf_rq_num_bits(#[case] requested: i32) {
6230 let error = build_ivf_rq_num_bits(Some(requested))
6231 .expect_err("invalid IVF_RQ num_bits should fail");
6232 let message = error.to_string();
6233
6234 assert_eq!(mutation_error_code(error), ErrorCode::InvalidInput);
6235 assert!(
6236 message.contains(&format!(
6237 "IVF_RQ num_bits must be in 1..=9, got {requested}"
6238 )),
6239 "unexpected error message: {message}"
6240 );
6241 }
6242
6243 #[test]
6244 fn test_build_engine_fts_query_match() {
6245 let mut ns_match = lance_namespace::models::MatchQuery::new("hello world".to_string());
6246 ns_match.column = Some("body".to_string());
6247 ns_match.operator = Some("AND".to_string());
6248 ns_match.fuzziness = Some(1);
6249 ns_match.max_expansions = Some(30);
6250 ns_match.boost = Some(2.0);
6251 ns_match.prefix_length = Some(2);
6252
6253 let mut ns_query = lance_namespace::models::FtsQuery::new();
6254 ns_query.r#match = Some(Box::new(ns_match));
6255
6256 match build_engine_fts_query(&ns_query).unwrap() {
6257 FtsQuery::Match(m) => {
6258 assert_eq!(m.terms, "hello world");
6259 assert_eq!(m.column, Some("body".to_string()));
6260 assert_eq!(m.operator, Operator::And);
6261 assert_eq!(m.fuzziness, Some(1));
6262 assert_eq!(m.max_expansions, 30);
6263 assert_eq!(m.boost, 2.0);
6264 assert_eq!(m.prefix_length, 2);
6265 }
6266 other => panic!("expected Match, got {:?}", other),
6267 }
6268 }
6269
6270 fn ns_match_query(terms: &str, column: &str) -> lance_namespace::models::FtsQuery {
6273 let mut m = lance_namespace::models::MatchQuery::new(terms.to_string());
6274 m.column = Some(column.to_string());
6275 let mut q = lance_namespace::models::FtsQuery::new();
6276 q.r#match = Some(Box::new(m));
6277 q
6278 }
6279
6280 #[test]
6281 fn test_build_engine_fts_query_phrase() {
6282 let mut ns_phrase = lance_namespace::models::PhraseQuery::new("hello world".to_string());
6283 ns_phrase.column = Some("body".to_string());
6284 ns_phrase.slop = Some(2);
6285
6286 let mut ns_query = lance_namespace::models::FtsQuery::new();
6287 ns_query.phrase = Some(Box::new(ns_phrase));
6288
6289 match build_engine_fts_query(&ns_query).unwrap() {
6290 FtsQuery::Phrase(p) => {
6291 assert_eq!(p.terms, "hello world");
6292 assert_eq!(p.column, Some("body".to_string()));
6293 assert_eq!(p.slop, 2);
6294 }
6295 other => panic!("expected Phrase, got {:?}", other),
6296 }
6297 }
6298
6299 #[test]
6300 fn test_build_engine_fts_query_multi_match() {
6301 let mut m1 = lance_namespace::models::MatchQuery::new("hello".to_string());
6302 m1.column = Some("title".to_string());
6303 let mut m2 = lance_namespace::models::MatchQuery::new("hello".to_string());
6304 m2.column = Some("body".to_string());
6305 m2.boost = Some(2.0);
6306
6307 let ns_multi = lance_namespace::models::MultiMatchQuery::new(vec![m1, m2]);
6308 let mut ns_query = lance_namespace::models::FtsQuery::new();
6309 ns_query.multi_match = Some(Box::new(ns_multi));
6310
6311 match build_engine_fts_query(&ns_query).unwrap() {
6312 FtsQuery::MultiMatch(mm) => {
6313 assert_eq!(mm.match_queries.len(), 2);
6314 assert_eq!(mm.match_queries[0].terms, "hello");
6315 assert_eq!(mm.match_queries[0].column, Some("title".to_string()));
6316 assert_eq!(mm.match_queries[1].column, Some("body".to_string()));
6317 assert_eq!(mm.match_queries[1].boost, 2.0);
6318 }
6319 other => panic!("expected MultiMatch, got {:?}", other),
6320 }
6321 }
6322
6323 #[test]
6324 fn test_build_engine_fts_query_boolean() {
6325 let ns_boolean = lance_namespace::models::BooleanQuery::new(
6327 vec![ns_match_query("must-term", "body")],
6328 vec![ns_match_query("must-not-term", "body")],
6329 vec![ns_match_query("should-term", "body")],
6330 );
6331 let mut ns_query = lance_namespace::models::FtsQuery::new();
6332 ns_query.boolean = Some(Box::new(ns_boolean));
6333
6334 match build_engine_fts_query(&ns_query).unwrap() {
6335 FtsQuery::Boolean(b) => {
6336 assert!(matches!(&b.must[..], [FtsQuery::Match(m)] if m.terms == "must-term"));
6337 assert!(
6338 matches!(&b.must_not[..], [FtsQuery::Match(m)] if m.terms == "must-not-term")
6339 );
6340 assert!(matches!(&b.should[..], [FtsQuery::Match(m)] if m.terms == "should-term"));
6341 }
6342 other => panic!("expected Boolean, got {:?}", other),
6343 }
6344 }
6345
6346 #[test]
6347 fn test_build_engine_fts_query_boost() {
6348 let mut ns_boost = lance_namespace::models::BoostQuery::new(
6349 ns_match_query("positive-term", "body"),
6350 ns_match_query("negative-term", "body"),
6351 );
6352 ns_boost.negative_boost = Some(0.25);
6353
6354 let mut ns_query = lance_namespace::models::FtsQuery::new();
6355 ns_query.boost = Some(Box::new(ns_boost));
6356
6357 match build_engine_fts_query(&ns_query).unwrap() {
6358 FtsQuery::Boost(b) => {
6359 assert!(
6360 matches!(b.positive.as_ref(), FtsQuery::Match(m) if m.terms == "positive-term")
6361 );
6362 assert!(
6363 matches!(b.negative.as_ref(), FtsQuery::Match(m) if m.terms == "negative-term")
6364 );
6365 assert_eq!(b.negative_boost, 0.25);
6366 }
6367 other => panic!("expected Boost, got {:?}", other),
6368 }
6369 }
6370
6371 #[test]
6372 fn test_build_engine_fts_query_requires_a_variant() {
6373 let empty = lance_namespace::models::FtsQuery::new();
6375 assert!(build_engine_fts_query(&empty).is_err());
6376 }
6377 use lance::dataset::Dataset;
6378 use lance::index::DatasetIndexExt;
6379 use lance_core::utils::tempfile::{TempStdDir, TempStrDir};
6380 use lance_core::utils::testing::CountingObjectStore;
6381 use lance_io::object_store::{providers::local::FileStoreProvider, uri_to_url};
6382 use lance_namespace::error::ErrorCode;
6383 use lance_namespace::models::{
6384 CreateTableRequest, JsonArrowDataType, JsonArrowField, JsonArrowSchema, ListTablesRequest,
6385 QueryTableRequestColumns,
6386 };
6387 use lance_namespace::schema::convert_json_arrow_schema;
6388 use std::io::Cursor;
6389 use std::sync::{
6390 Arc,
6391 atomic::{AtomicUsize, Ordering},
6392 };
6393 use url::Url;
6394
6395 fn assert_plan_contains_all(plan: &str, expected_fragments: &[&str], context: &str) {
6396 for expected_fragment in expected_fragments {
6397 assert!(
6398 plan.contains(expected_fragment),
6399 "{}. Missing fragment: '{}'. Plan:\n{}",
6400 context,
6401 expected_fragment,
6402 plan
6403 );
6404 }
6405 }
6406
6407 fn mutation_error_code(err: lance_core::Error) -> ErrorCode {
6408 match err {
6409 lance_core::Error::Namespace { source, .. } => source
6410 .downcast_ref::<NamespaceError>()
6411 .expect("mutation error should wrap a NamespaceError")
6412 .code(),
6413 other => panic!("expected Namespace error, got: {other:?}"),
6414 }
6415 }
6416
6417 #[test]
6422 fn test_map_mutation_error_commit_conflict_alignment() {
6423 let boxed = || -> Box<dyn std::error::Error + Send + Sync + 'static> {
6424 Box::<dyn std::error::Error + Send + Sync>::from("inner conflict")
6425 };
6426
6427 let throttling_cases = vec![lance_core::Error::commit_conflict_source(1, boxed())];
6428 for err in throttling_cases {
6429 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6430 err,
6431 "update",
6432 "memory://t",
6433 ));
6434 assert_eq!(code, ErrorCode::Throttling);
6435 }
6436
6437 let concurrent_cases = vec![
6438 lance_core::Error::too_much_write_contention("contention"),
6439 lance_core::Error::retryable_commit_conflict_source(1, boxed()),
6440 lance_core::Error::incompatible_transaction_source(boxed()),
6441 lance_core::Error::version_conflict("conflict", 0, 3),
6442 ];
6443 for err in concurrent_cases {
6444 let code = mutation_error_code(DirectoryNamespace::map_mutation_error(
6445 err,
6446 "update",
6447 "memory://t",
6448 ));
6449 assert_eq!(code, ErrorCode::ConcurrentModification);
6450 }
6451 }
6452
6453 async fn create_test_namespace() -> (DirectoryNamespace, TempStdDir) {
6455 let temp_dir = TempStdDir::default();
6456
6457 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
6458 .build()
6459 .await
6460 .unwrap();
6461 (namespace, temp_dir)
6462 }
6463
6464 #[tokio::test]
6467 async fn test_list_versions_under_ordering_and_limit() {
6468 use lance_table::io::commit::ManifestNamingScheme;
6469
6470 async fn seed_and_check(ns: &DirectoryNamespace) {
6471 let table_path = ns.base_path.clone().join("lv_test.lance");
6472 for v in 1..=7u64 {
6473 let p = ManifestNamingScheme::V2.manifest_path(&table_path, v);
6474 ns.object_store.put(&p, b"m".as_slice()).await.unwrap();
6475 }
6476 let staging = Path::parse(format!(
6480 "{}-cee4fbbb-eb19-4ea3-8ca7-54f5ec33dedc",
6481 ManifestNamingScheme::V2.manifest_path(&table_path, 8)
6482 ))
6483 .unwrap();
6484 ns.object_store
6485 .put(&staging, b"s".as_slice())
6486 .await
6487 .unwrap();
6488 let detached = table_path.clone().join(VERSIONS_DIR).join("d123.manifest");
6489 ns.object_store
6490 .put(&detached, b"d".as_slice())
6491 .await
6492 .unwrap();
6493 fn versions(r: &[TableVersion]) -> Vec<i64> {
6494 r.iter().map(|t| t.version).collect()
6495 }
6496
6497 let got = ns
6498 .list_versions_under(&table_path, true, Some(1))
6499 .await
6500 .unwrap();
6501 assert_eq!(versions(&got), vec![7]);
6502 let got = ns
6503 .list_versions_under(&table_path, true, Some(3))
6504 .await
6505 .unwrap();
6506 assert_eq!(versions(&got), vec![7, 6, 5]);
6507
6508 let got = ns
6509 .list_versions_under(&table_path, false, Some(2))
6510 .await
6511 .unwrap();
6512 assert_eq!(versions(&got), vec![1, 2]);
6513
6514 let got = ns
6515 .list_versions_under(&table_path, true, None)
6516 .await
6517 .unwrap();
6518 assert_eq!(versions(&got), vec![7, 6, 5, 4, 3, 2, 1]);
6519 let got = ns
6520 .list_versions_under(&table_path, false, None)
6521 .await
6522 .unwrap();
6523 assert_eq!(versions(&got), vec![1, 2, 3, 4, 5, 6, 7]);
6524
6525 let got = ns
6526 .list_versions_under(&table_path, true, Some(0))
6527 .await
6528 .unwrap();
6529 assert!(got.is_empty());
6530 let got = ns
6531 .list_versions_under(&table_path, true, Some(100))
6532 .await
6533 .unwrap();
6534 assert_eq!(versions(&got), vec![7, 6, 5, 4, 3, 2, 1]);
6535
6536 let got = ns
6538 .list_versions_under(&table_path, true, Some(-1))
6539 .await
6540 .unwrap();
6541 assert_eq!(versions(&got), vec![7, 6, 5, 4, 3, 2, 1]);
6542 }
6543
6544 let ns_mem = DirectoryNamespaceBuilder::new("memory://lv-test")
6545 .build()
6546 .await
6547 .unwrap();
6548 assert!(ns_mem.object_store.list_is_lexically_ordered);
6549 seed_and_check(&ns_mem).await;
6550
6551 let (ns_fs, _tmp) = create_test_namespace().await;
6552 assert!(!ns_fs.object_store.list_is_lexically_ordered);
6553 seed_and_check(&ns_fs).await;
6554 }
6555
6556 #[tokio::test]
6560 async fn test_list_versions_under_early_stop_bounded_consumption() {
6561 use lance_io::object_store::providers::memory::MemoryStoreProvider;
6562 use lance_table::io::commit::ManifestNamingScheme;
6563
6564 #[derive(Debug)]
6565 struct EntryCountingStore {
6566 target: Arc<dyn OSObjectStore>,
6567 entries_listed: Arc<AtomicUsize>,
6568 }
6569
6570 impl std::fmt::Display for EntryCountingStore {
6571 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6572 write!(f, "EntryCountingStore({})", self.target)
6573 }
6574 }
6575
6576 #[async_trait]
6577 impl OSObjectStore for EntryCountingStore {
6578 async fn put_opts(
6579 &self,
6580 location: &Path,
6581 bytes: PutPayload,
6582 opts: PutOptions,
6583 ) -> OSResult<PutResult> {
6584 self.target.put_opts(location, bytes, opts).await
6585 }
6586
6587 async fn put_multipart_opts(
6588 &self,
6589 location: &Path,
6590 opts: PutMultipartOptions,
6591 ) -> OSResult<Box<dyn MultipartUpload>> {
6592 self.target.put_multipart_opts(location, opts).await
6593 }
6594
6595 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
6596 self.target.get_opts(location, options).await
6597 }
6598
6599 fn delete_stream(
6600 &self,
6601 locations: BoxStream<'static, OSResult<Path>>,
6602 ) -> BoxStream<'static, OSResult<Path>> {
6603 self.target.delete_stream(locations)
6604 }
6605
6606 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
6607 let entries_listed = self.entries_listed.clone();
6608 self.target
6609 .list(prefix)
6610 .inspect(move |_| {
6611 entries_listed.fetch_add(1, Ordering::SeqCst);
6612 })
6613 .boxed()
6614 }
6615
6616 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
6617 self.target.list_with_delimiter(prefix).await
6618 }
6619
6620 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
6621 self.target.copy_opts(from, to, opts).await
6622 }
6623 }
6624
6625 #[derive(Debug)]
6626 struct EntryCountingMemoryProvider {
6627 entries_listed: Arc<AtomicUsize>,
6628 }
6629
6630 #[async_trait]
6631 impl lance_io::object_store::ObjectStoreProvider for EntryCountingMemoryProvider {
6632 async fn new_store(
6633 &self,
6634 base_path: Url,
6635 params: &ObjectStoreParams,
6636 ) -> Result<ObjectStore> {
6637 let mut store = MemoryStoreProvider.new_store(base_path, params).await?;
6638 store.inner = Arc::new(EntryCountingStore {
6639 target: store.inner.clone(),
6640 entries_listed: self.entries_listed.clone(),
6641 });
6642 Ok(store)
6643 }
6644
6645 fn extract_path(&self, url: &Url) -> Result<Path> {
6646 MemoryStoreProvider.extract_path(url)
6647 }
6648
6649 fn calculate_object_store_prefix(
6650 &self,
6651 url: &Url,
6652 storage_options: Option<&HashMap<String, String>>,
6653 ) -> Result<String> {
6654 MemoryStoreProvider.calculate_object_store_prefix(url, storage_options)
6655 }
6656 }
6657
6658 let entries_listed = Arc::new(AtomicUsize::new(0));
6659 let registry = Arc::new(ObjectStoreRegistry::default());
6660 registry.insert(
6661 "memory-object-store",
6662 Arc::new(EntryCountingMemoryProvider {
6663 entries_listed: entries_listed.clone(),
6664 }),
6665 );
6666 let session = Arc::new(Session::new(0, 0, registry));
6667 let ns = DirectoryNamespaceBuilder::new("memory-object-store://lv-count")
6668 .session(session)
6669 .build()
6670 .await
6671 .unwrap();
6672 assert!(ns.object_store.list_is_lexically_ordered);
6673
6674 let table_path = ns.base_path.clone().join("lv_count.lance");
6675 for v in 1..=100u64 {
6676 let p = ManifestNamingScheme::V2.manifest_path(&table_path, v);
6677 ns.object_store.put(&p, b"m".as_slice()).await.unwrap();
6678 }
6679 let staging = Path::parse(format!(
6681 "{}-cee4fbbb-eb19-4ea3-8ca7-54f5ec33dedc",
6682 ManifestNamingScheme::V2.manifest_path(&table_path, 101)
6683 ))
6684 .unwrap();
6685 ns.object_store
6686 .put(&staging, b"s".as_slice())
6687 .await
6688 .unwrap();
6689
6690 let consumed_before = entries_listed.load(Ordering::SeqCst);
6691 let got = ns
6692 .list_versions_under(&table_path, true, Some(1))
6693 .await
6694 .unwrap();
6695 let consumed = entries_listed.load(Ordering::SeqCst) - consumed_before;
6696
6697 assert_eq!(got.len(), 1);
6698 assert_eq!(got[0].version, 100);
6699 assert_eq!(
6700 consumed, 2,
6701 "latest-version query must consume only the staging entry plus the \
6702 first committed manifest, not the whole directory (consumed {} of \
6703 101 entries)",
6704 consumed
6705 );
6706 }
6707
6708 #[derive(Debug)]
6709 #[allow(dead_code)]
6710 struct CountingFileStoreProvider {
6711 listing_count: Arc<AtomicUsize>,
6712 }
6713
6714 #[async_trait]
6715 impl lance_io::object_store::ObjectStoreProvider for CountingFileStoreProvider {
6716 async fn new_store(
6717 &self,
6718 base_path: Url,
6719 params: &ObjectStoreParams,
6720 ) -> Result<ObjectStore> {
6721 let provider = FileStoreProvider;
6722 let mut store = provider.new_store(base_path, params).await?;
6723 store.inner = Arc::new(CountingObjectStore::new(
6724 store.inner.clone(),
6725 self.listing_count.clone(),
6726 ));
6727 Ok(store)
6728 }
6729
6730 fn extract_path(&self, url: &Url) -> Result<Path> {
6731 let provider = FileStoreProvider;
6732 provider.extract_path(url)
6733 }
6734
6735 fn calculate_object_store_prefix(
6736 &self,
6737 url: &Url,
6738 storage_options: Option<&HashMap<String, String>>,
6739 ) -> Result<String> {
6740 let provider = FileStoreProvider;
6741 provider.calculate_object_store_prefix(url, storage_options)
6742 }
6743 }
6744
6745 #[allow(dead_code)]
6746 fn file_object_store_uri(path: &str) -> String {
6747 let file_url = uri_to_url(path).unwrap();
6748 let mut url = Url::parse("file-object-store:///").unwrap();
6749 url.set_path(file_url.path());
6750 url.to_string()
6751 }
6752
6753 #[allow(dead_code)]
6754 fn build_listing_counting_session(listing_count: Arc<AtomicUsize>) -> Arc<Session> {
6755 let registry = Arc::new(ObjectStoreRegistry::default());
6756 registry.insert(
6757 "file-object-store",
6758 Arc::new(CountingFileStoreProvider { listing_count }),
6759 );
6760 Arc::new(Session::new(0, 0, registry))
6761 }
6762
6763 use futures::stream::BoxStream;
6767 use object_store::{
6768 CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
6769 PutMultipartOptions, PutPayload, PutResult, Result as OSResult,
6770 };
6771 use std::ops::Range;
6772
6773 #[derive(Debug, Clone, Copy)]
6774 enum ListBehavior {
6775 Throttle,
6776 ServiceUnavailable,
6777 Internal,
6778 NotFound,
6779 EmptyListing,
6780 }
6781
6782 #[derive(Debug)]
6783 struct FailingListStore {
6784 target: Arc<dyn OSObjectStore>,
6785 behavior: Arc<Mutex<Option<ListBehavior>>>,
6786 }
6787
6788 impl std::fmt::Display for FailingListStore {
6789 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
6790 write!(f, "FailingListStore({})", self.target)
6791 }
6792 }
6793
6794 #[async_trait]
6795 impl OSObjectStore for FailingListStore {
6796 async fn put_opts(
6797 &self,
6798 location: &Path,
6799 bytes: PutPayload,
6800 opts: PutOptions,
6801 ) -> OSResult<PutResult> {
6802 self.target.put_opts(location, bytes, opts).await
6803 }
6804
6805 async fn put_multipart_opts(
6806 &self,
6807 location: &Path,
6808 opts: PutMultipartOptions,
6809 ) -> OSResult<Box<dyn MultipartUpload>> {
6810 self.target.put_multipart_opts(location, opts).await
6811 }
6812
6813 async fn get_opts(&self, location: &Path, options: GetOptions) -> OSResult<GetResult> {
6814 self.target.get_opts(location, options).await
6815 }
6816
6817 async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> OSResult<Vec<Bytes>> {
6818 self.target.get_ranges(location, ranges).await
6819 }
6820
6821 fn delete_stream(
6822 &self,
6823 locations: BoxStream<'static, OSResult<Path>>,
6824 ) -> BoxStream<'static, OSResult<Path>> {
6825 self.target.delete_stream(locations)
6826 }
6827
6828 fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, OSResult<ObjectMeta>> {
6829 self.target.list(prefix)
6830 }
6831
6832 async fn list_with_delimiter(&self, prefix: Option<&Path>) -> OSResult<ListResult> {
6833 let behavior = *self.behavior.lock().unwrap();
6834 match behavior {
6835 None => self.target.list_with_delimiter(prefix).await,
6836 Some(ListBehavior::EmptyListing) => Ok(ListResult {
6837 common_prefixes: Vec::new(),
6838 objects: Vec::new(),
6839 extensions: Default::default(),
6840 }),
6841 Some(ListBehavior::Throttle) => Err(ObjectStoreError::Generic {
6844 store: "test",
6845 source: "Error performing list request: response error, after 3 retries, \
6846 max_retries: 3, retry_timeout: 180s - HTTP status server error \
6847 (503 Service Unavailable): ServerBusy: The server is busy"
6848 .into(),
6849 }),
6850 Some(ListBehavior::ServiceUnavailable) => Err(ObjectStoreError::Generic {
6851 store: "test",
6852 source: "Error performing list request: 503 Service Unavailable".into(),
6853 }),
6854 Some(ListBehavior::Internal) => Err(ObjectStoreError::Generic {
6855 store: "test",
6856 source: "Error performing list request: catastrophic unclassified failure"
6857 .into(),
6858 }),
6859 Some(ListBehavior::NotFound) => Err(ObjectStoreError::NotFound {
6860 path: "test_table.lance".to_string(),
6861 source: "entity not found".into(),
6862 }),
6863 }
6864 }
6865
6866 async fn copy_opts(&self, from: &Path, to: &Path, opts: CopyOptions) -> OSResult<()> {
6867 self.target.copy_opts(from, to, opts).await
6868 }
6869 }
6870
6871 #[derive(Debug)]
6872 struct FailingListStoreProvider {
6873 behavior: Arc<Mutex<Option<ListBehavior>>>,
6874 }
6875
6876 #[async_trait]
6877 impl lance_io::object_store::ObjectStoreProvider for FailingListStoreProvider {
6878 async fn new_store(
6879 &self,
6880 base_path: Url,
6881 params: &ObjectStoreParams,
6882 ) -> Result<ObjectStore> {
6883 let mut store = FileStoreProvider.new_store(base_path, params).await?;
6884 store.inner = Arc::new(FailingListStore {
6885 target: store.inner.clone(),
6886 behavior: self.behavior.clone(),
6887 });
6888 Ok(store)
6889 }
6890
6891 fn extract_path(&self, url: &Url) -> Result<Path> {
6892 FileStoreProvider.extract_path(url)
6893 }
6894
6895 fn calculate_object_store_prefix(
6896 &self,
6897 url: &Url,
6898 storage_options: Option<&HashMap<String, String>>,
6899 ) -> Result<String> {
6900 FileStoreProvider.calculate_object_store_prefix(url, storage_options)
6901 }
6902 }
6903
6904 fn build_failing_list_session(behavior: Arc<Mutex<Option<ListBehavior>>>) -> Arc<Session> {
6905 let registry = Arc::new(ObjectStoreRegistry::default());
6906 registry.insert(
6907 "file-object-store",
6908 Arc::new(FailingListStoreProvider { behavior }),
6909 );
6910 Arc::new(Session::new(0, 0, registry))
6911 }
6912
6913 async fn failing_list_namespace() -> (
6917 DirectoryNamespace,
6918 TempStdDir,
6919 Arc<Mutex<Option<ListBehavior>>>,
6920 ) {
6921 let temp_dir = TempStdDir::default();
6922 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
6923 let behavior = Arc::new(Mutex::new(None));
6924 let session = build_failing_list_session(behavior.clone());
6925 let namespace = DirectoryNamespaceBuilder::new(root_uri)
6926 .session(session)
6927 .manifest_enabled(false)
6928 .dir_listing_enabled(true)
6929 .build()
6930 .await
6931 .unwrap();
6932 (namespace, temp_dir, behavior)
6933 }
6934
6935 async fn create_named_dir_table(namespace: &DirectoryNamespace, name: &str) {
6936 let schema = create_test_schema();
6937 let ipc_data = create_test_ipc_data(&schema);
6938 let mut create_req = CreateTableRequest::new();
6939 create_req.id = Some(vec![name.to_string()]);
6940 namespace
6941 .create_table(create_req, Bytes::from(ipc_data))
6942 .await
6943 .unwrap();
6944 }
6945
6946 #[tokio::test]
6951 async fn test_table_resolution_propagates_storage_errors_not_table_not_found() {
6952 for (behavior, expected_code, evidence) in [
6953 (ListBehavior::Throttle, ErrorCode::Throttling, "serverbusy"),
6954 (
6955 ListBehavior::ServiceUnavailable,
6956 ErrorCode::ServiceUnavailable,
6957 "503 service unavailable",
6958 ),
6959 (ListBehavior::Internal, ErrorCode::Internal, "catastrophic"),
6960 ] {
6961 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
6962 create_named_dir_table(&namespace, "checkpoint").await;
6963 *toggle.lock().unwrap() = Some(behavior);
6964
6965 let mut describe_req = DescribeTableRequest::new();
6966 describe_req.id = Some(vec!["checkpoint".to_string()]);
6967 let err = namespace.describe_table(describe_req).await.unwrap_err();
6968 let msg = err.to_string();
6969 assert_eq!(
6970 mutation_error_code(err),
6971 expected_code,
6972 "describe_table under {behavior:?}; msg: {msg}"
6973 );
6974 assert!(
6975 msg.to_ascii_lowercase().contains(evidence),
6976 "describe_table message must carry storage evidence '{evidence}', got: {msg}"
6977 );
6978
6979 let mut exists_req = TableExistsRequest::new();
6980 exists_req.id = Some(vec!["checkpoint".to_string()]);
6981 let err = namespace.table_exists(exists_req).await.unwrap_err();
6982 let msg = err.to_string();
6983 assert_eq!(
6984 mutation_error_code(err),
6985 expected_code,
6986 "table_exists under {behavior:?}; msg: {msg}"
6987 );
6988 assert!(
6989 msg.to_ascii_lowercase().contains(evidence),
6990 "table_exists message must carry storage evidence '{evidence}', got: {msg}"
6991 );
6992 }
6993 }
6994
6995 #[tokio::test]
6998 async fn test_table_resolution_missing_table_yields_table_not_found() {
6999 for behavior in [ListBehavior::NotFound, ListBehavior::EmptyListing] {
7000 let (namespace, _temp_dir, toggle) = failing_list_namespace().await;
7001 *toggle.lock().unwrap() = Some(behavior);
7002
7003 let mut describe_req = DescribeTableRequest::new();
7004 describe_req.id = Some(vec!["missing".to_string()]);
7005 let err = namespace.describe_table(describe_req).await.unwrap_err();
7006 assert_eq!(
7007 mutation_error_code(err),
7008 ErrorCode::TableNotFound,
7009 "describe_table under {behavior:?} should be TableNotFound"
7010 );
7011
7012 let mut exists_req = TableExistsRequest::new();
7013 exists_req.id = Some(vec!["missing".to_string()]);
7014 let err = namespace.table_exists(exists_req).await.unwrap_err();
7015 assert_eq!(
7016 mutation_error_code(err),
7017 ErrorCode::TableNotFound,
7018 "table_exists under {behavior:?} should be TableNotFound"
7019 );
7020 }
7021 }
7022
7023 #[tokio::test]
7033 async fn test_hybrid_resolution_falls_through_and_does_not_mask_throttle() {
7034 let temp_dir = TempStdDir::default();
7035 let root_uri = file_object_store_uri(temp_dir.to_str().unwrap());
7036 let behavior = Arc::new(Mutex::new(None));
7037 let session = build_failing_list_session(behavior.clone());
7038
7039 let manifest_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
7041 .session(session.clone())
7042 .manifest_enabled(true)
7043 .dir_listing_enabled(true)
7044 .build()
7045 .await
7046 .unwrap();
7047 create_named_dir_table(&manifest_ns, "seed").await;
7048
7049 let dir_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
7051 .session(session.clone())
7052 .manifest_enabled(false)
7053 .dir_listing_enabled(true)
7054 .build()
7055 .await
7056 .unwrap();
7057 create_named_dir_table(&dir_ns, "checkpoint").await;
7058
7059 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
7062 .session(session)
7063 .manifest_enabled(true)
7064 .dir_listing_enabled(true)
7065 .dir_listing_to_manifest_migration_enabled(true)
7066 .build()
7067 .await
7068 .unwrap();
7069
7070 let mut describe_req = DescribeTableRequest::new();
7073 describe_req.id = Some(vec!["checkpoint".to_string()]);
7074 hybrid_ns
7075 .describe_table(describe_req)
7076 .await
7077 .expect("unregistered on-disk table should resolve via the manifest fall-through");
7078
7079 *behavior.lock().unwrap() = Some(ListBehavior::Throttle);
7083 let mut describe_req = DescribeTableRequest::new();
7084 describe_req.id = Some(vec!["checkpoint".to_string()]);
7085 let err = hybrid_ns.describe_table(describe_req).await.unwrap_err();
7086 let code = mutation_error_code(err);
7087 assert_ne!(
7088 code,
7089 ErrorCode::TableNotFound,
7090 "hybrid resolution masked a throttle as TableNotFound"
7091 );
7092 assert!(
7093 matches!(
7094 code,
7095 ErrorCode::Throttling | ErrorCode::ServiceUnavailable | ErrorCode::Internal
7096 ),
7097 "hybrid resolution should surface a storage error, got {code:?}"
7098 );
7099 }
7100
7101 #[test]
7102 fn test_classify_storage_error_maps_variants_and_preserves_evidence() {
7103 let throttle: Error = ObjectStoreError::Generic {
7104 store: "test",
7105 source: "list request failed, after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
7106 }
7107 .into();
7108 assert!(matches!(throttle, Error::IO { .. }));
7109 let classified = DirectoryNamespace::classify_storage_error(throttle);
7110 let msg = classified.to_string();
7111 assert_eq!(mutation_error_code(classified), ErrorCode::Throttling);
7112 assert!(
7113 msg.to_ascii_lowercase().contains("serverbusy"),
7114 "throttle evidence lost: {msg}"
7115 );
7116
7117 let service: Error = ObjectStoreError::Generic {
7118 store: "test",
7119 source: "504 Gateway Timeout".into(),
7120 }
7121 .into();
7122 assert_eq!(
7123 mutation_error_code(DirectoryNamespace::classify_storage_error(service)),
7124 ErrorCode::ServiceUnavailable
7125 );
7126
7127 let internal: Error = ObjectStoreError::Generic {
7128 store: "test",
7129 source: "disk caught fire".into(),
7130 }
7131 .into();
7132 assert_eq!(
7133 mutation_error_code(DirectoryNamespace::classify_storage_error(internal)),
7134 ErrorCode::Internal
7135 );
7136
7137 let preexisting: Error = NamespaceError::TableAlreadyExists {
7139 message: "t".to_string(),
7140 }
7141 .into();
7142 assert_eq!(
7143 mutation_error_code(DirectoryNamespace::classify_storage_error(preexisting)),
7144 ErrorCode::TableAlreadyExists
7145 );
7146 }
7147
7148 #[test]
7149 fn test_is_manifest_table_absent_error() {
7150 let table_not_found: Error = NamespaceError::TableNotFound {
7151 message: "t".to_string(),
7152 }
7153 .into();
7154 assert!(DirectoryNamespace::is_manifest_table_absent_error(
7155 &table_not_found
7156 ));
7157 let raw_not_found: Error = ObjectStoreError::NotFound {
7158 path: "t".to_string(),
7159 source: "x".into(),
7160 }
7161 .into();
7162 assert!(DirectoryNamespace::is_manifest_table_absent_error(
7163 &raw_not_found
7164 ));
7165
7166 let throttle: Error = ObjectStoreError::Generic {
7167 store: "test",
7168 source: "after 3 retries, max_retries: 3 ServerBusy".into(),
7169 }
7170 .into();
7171 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
7172 &throttle
7173 ));
7174 let internal: Error = NamespaceError::Internal {
7175 message: "boom".to_string(),
7176 }
7177 .into();
7178 assert!(!DirectoryNamespace::is_manifest_table_absent_error(
7179 &internal
7180 ));
7181 }
7182
7183 #[test]
7184 fn test_map_open_error() {
7185 let not_found = || NamespaceError::TableNotFound {
7186 message: "table at 'x' not found: ...".to_string(),
7187 };
7188
7189 let throttle: Error = ObjectStoreError::Generic {
7190 store: "test",
7191 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
7192 }
7193 .into();
7194 assert_eq!(
7195 mutation_error_code(DirectoryNamespace::map_open_error(throttle, not_found())),
7196 ErrorCode::Throttling
7197 );
7198
7199 let generic_io: Error = ObjectStoreError::Generic {
7200 store: "test",
7201 source: "connection reset".into(),
7202 }
7203 .into();
7204 assert_eq!(
7205 mutation_error_code(DirectoryNamespace::map_open_error(generic_io, not_found())),
7206 ErrorCode::Internal
7207 );
7208
7209 let io_not_found: Error = ObjectStoreError::NotFound {
7210 path: "x".to_string(),
7211 source: "missing".into(),
7212 }
7213 .into();
7214 assert_eq!(
7215 mutation_error_code(DirectoryNamespace::map_open_error(
7216 io_not_found,
7217 not_found()
7218 )),
7219 ErrorCode::TableNotFound
7220 );
7221
7222 let dataset_not_found = Error::dataset_not_found("x".to_string(), "missing".into());
7223 assert_eq!(
7224 mutation_error_code(DirectoryNamespace::map_open_error(
7225 dataset_not_found,
7226 not_found()
7227 )),
7228 ErrorCode::TableNotFound
7229 );
7230
7231 let ref_not_found = Error::RefNotFound {
7233 message: "branch 'b' does not exist".to_string(),
7234 };
7235 assert_eq!(
7236 mutation_error_code(DirectoryNamespace::map_open_error(
7237 ref_not_found,
7238 not_found()
7239 )),
7240 ErrorCode::TableNotFound
7241 );
7242
7243 let version_miss = Error::RefNotFound {
7245 message: "version 5 does not exist".to_string(),
7246 };
7247 assert_eq!(
7248 mutation_error_code(DirectoryNamespace::map_open_error(
7249 version_miss,
7250 NamespaceError::TableVersionNotFound {
7251 message: "version 5 not found".to_string(),
7252 },
7253 )),
7254 ErrorCode::TableVersionNotFound
7255 );
7256 let version_throttle: Error = ObjectStoreError::Generic {
7257 store: "test",
7258 source: "after 3 retries, max_retries: 3 - 503 ServerBusy".into(),
7259 }
7260 .into();
7261 assert_eq!(
7262 mutation_error_code(DirectoryNamespace::map_open_error(
7263 version_throttle,
7264 NamespaceError::TableVersionNotFound {
7265 message: "version 5 not found".to_string(),
7266 },
7267 )),
7268 ErrorCode::Throttling
7269 );
7270 }
7271
7272 fn create_test_ipc_data(schema: &JsonArrowSchema) -> Vec<u8> {
7274 use arrow::ipc::writer::StreamWriter;
7275
7276 let arrow_schema = convert_json_arrow_schema(schema).unwrap();
7277 let arrow_schema = Arc::new(arrow_schema);
7278 let batch = arrow::record_batch::RecordBatch::new_empty(arrow_schema.clone());
7279 let mut buffer = Vec::new();
7280 {
7281 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
7282 writer.write(&batch).unwrap();
7283 writer.finish().unwrap();
7284 }
7285 buffer
7286 }
7287
7288 fn create_ipc_data_from_batches(
7289 schema: Arc<arrow_schema::Schema>,
7290 batches: Vec<arrow::record_batch::RecordBatch>,
7291 ) -> Vec<u8> {
7292 use arrow::ipc::writer::StreamWriter;
7293
7294 let mut buffer = Vec::new();
7295 {
7296 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
7297 for batch in &batches {
7298 writer.write(batch).unwrap();
7299 }
7300 writer.finish().unwrap();
7301 }
7302 buffer
7303 }
7304
7305 fn create_non_empty_test_ipc_data() -> Vec<u8> {
7306 use arrow::array::{Int32Array, StringArray};
7307 use arrow::record_batch::RecordBatch;
7308
7309 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
7310 let batch = RecordBatch::try_new(
7311 schema.clone(),
7312 vec![
7313 Arc::new(Int32Array::from(vec![1, 2])),
7314 Arc::new(StringArray::from(vec![Some("alice"), Some("bob")])),
7315 ],
7316 )
7317 .unwrap();
7318 create_ipc_data_from_batches(schema, vec![batch])
7319 }
7320
7321 fn create_single_row_test_ipc_data() -> Vec<u8> {
7322 use arrow::array::{Int32Array, StringArray};
7323 use arrow::record_batch::RecordBatch;
7324
7325 let schema = Arc::new(convert_json_arrow_schema(&create_test_schema()).unwrap());
7326 let batch = RecordBatch::try_new(
7327 schema.clone(),
7328 vec![
7329 Arc::new(Int32Array::from(vec![10])),
7330 Arc::new(StringArray::from(vec![Some("carol")])),
7331 ],
7332 )
7333 .unwrap();
7334 create_ipc_data_from_batches(schema, vec![batch])
7335 }
7336
7337 fn create_test_schema() -> JsonArrowSchema {
7339 let int_type = JsonArrowDataType::new("int32".to_string());
7340 let string_type = JsonArrowDataType::new("utf8".to_string());
7341
7342 let id_field = JsonArrowField {
7343 name: "id".to_string(),
7344 r#type: Box::new(int_type),
7345 nullable: false,
7346 metadata: None,
7347 };
7348
7349 let name_field = JsonArrowField {
7350 name: "name".to_string(),
7351 r#type: Box::new(string_type),
7352 nullable: true,
7353 metadata: None,
7354 };
7355
7356 JsonArrowSchema {
7357 fields: vec![id_field, name_field],
7358 metadata: None,
7359 }
7360 }
7361
7362 fn create_scalar_table_ipc_data() -> Vec<u8> {
7363 use arrow::array::{Int32Array, StringArray};
7364 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7365
7366 let schema = Arc::new(ArrowSchema::new(vec![
7367 Field::new("id", DataType::Int32, false),
7368 Field::new("name", DataType::Utf8, true),
7369 ]));
7370 let batch = arrow::record_batch::RecordBatch::try_new(
7371 schema.clone(),
7372 vec![
7373 Arc::new(Int32Array::from(vec![1, 2, 3])),
7374 Arc::new(StringArray::from(vec!["alice", "bob", "cory"])),
7375 ],
7376 )
7377 .unwrap();
7378 create_ipc_data_from_batches(schema, vec![batch])
7379 }
7380
7381 async fn create_legacy_manifest_without_primary_key_metadata(root: &str) {
7382 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7383 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
7384
7385 let schema = Arc::new(ArrowSchema::new(vec![
7386 Field::new("object_id", DataType::Utf8, false),
7387 Field::new("object_type", DataType::Utf8, false),
7388 Field::new("location", DataType::Utf8, true),
7389 Field::new("metadata", DataType::Utf8, true),
7390 Field::new(
7391 "base_objects",
7392 DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))),
7393 true,
7394 ),
7395 ]));
7396 let batch = RecordBatch::new_empty(schema.clone());
7397 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema);
7398 Dataset::write(Box::new(reader), &format!("{}/__manifest", root), None)
7399 .await
7400 .unwrap();
7401 }
7402
7403 async fn manifest_has_primary_key_metadata(root: &str) -> bool {
7404 let dataset = Dataset::open(&format!("{}/__manifest", root))
7405 .await
7406 .unwrap();
7407 dataset
7408 .schema()
7409 .field("object_id")
7410 .map(|field| {
7411 field
7412 .metadata
7413 .contains_key(lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
7414 })
7415 .unwrap_or(false)
7416 }
7417
7418 fn create_vector_table_ipc_data() -> Vec<u8> {
7419 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
7420 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7421
7422 let schema = Arc::new(ArrowSchema::new(vec![
7423 Field::new("id", DataType::Int32, false),
7424 Field::new(
7425 "vector",
7426 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2),
7427 true,
7428 ),
7429 ]));
7430 let vector_field = Arc::new(Field::new("item", DataType::Float32, true));
7431 let vectors = FixedSizeListArray::try_new(
7432 vector_field,
7433 2,
7434 Arc::new(Float32Array::from(vec![0.1, 0.2, 0.3, 0.4, 0.5, 0.6])),
7435 None,
7436 )
7437 .unwrap();
7438 let batch = arrow::record_batch::RecordBatch::try_new(
7439 schema.clone(),
7440 vec![Arc::new(Int32Array::from(vec![1, 2, 3])), Arc::new(vectors)],
7441 )
7442 .unwrap();
7443 create_ipc_data_from_batches(schema, vec![batch])
7444 }
7445
7446 async fn create_scalar_table(namespace: &DirectoryNamespace, table_name: &str) {
7447 let mut create_table_request = CreateTableRequest::new();
7448 create_table_request.id = Some(vec![table_name.to_string()]);
7449 namespace
7450 .create_table(
7451 create_table_request,
7452 Bytes::from(create_scalar_table_ipc_data()),
7453 )
7454 .await
7455 .unwrap();
7456 }
7457
7458 async fn create_vector_table(namespace: &DirectoryNamespace, table_name: &str) {
7459 let mut create_table_request = CreateTableRequest::new();
7460 create_table_request.id = Some(vec![table_name.to_string()]);
7461 namespace
7462 .create_table(
7463 create_table_request,
7464 Bytes::from(create_vector_table_ipc_data()),
7465 )
7466 .await
7467 .unwrap();
7468 }
7469
7470 async fn open_dataset(namespace: &DirectoryNamespace, table_name: &str) -> Dataset {
7471 let mut describe_request = DescribeTableRequest::new();
7472 describe_request.id = Some(vec![table_name.to_string()]);
7473 let table_uri = namespace
7474 .describe_table(describe_request)
7475 .await
7476 .unwrap()
7477 .location
7478 .expect("table location should exist");
7479 Dataset::open(&table_uri).await.unwrap()
7480 }
7481
7482 async fn create_scalar_index(
7483 namespace: &DirectoryNamespace,
7484 table_name: &str,
7485 index_name: &str,
7486 ) -> Option<String> {
7487 use lance_namespace::models::CreateTableIndexRequest;
7488
7489 let mut create_index_request =
7490 CreateTableIndexRequest::new("id".to_string(), "BTREE".to_string());
7491 create_index_request.id = Some(vec![table_name.to_string()]);
7492 create_index_request.name = Some(index_name.to_string());
7493 namespace
7494 .create_table_scalar_index(create_index_request)
7495 .await
7496 .unwrap()
7497 .transaction_id
7498 }
7499
7500 async fn create_branch_with_commits(
7505 namespace: &DirectoryNamespace,
7506 table_name: &str,
7507 branch_name: &str,
7508 extra_versions: usize,
7509 ) -> String {
7510 let mut main = open_dataset(namespace, table_name).await;
7511 let fork_version = main.version().version;
7512 let branch = main
7513 .create_branch(branch_name, fork_version, None)
7514 .await
7515 .unwrap();
7516 let branch_uri = branch.uri().to_string();
7517 for i in 0..extra_versions {
7518 append_scalar_version(&branch_uri, (i as i32 + 1) * 100).await;
7519 }
7520 branch_uri
7521 }
7522
7523 async fn append_scalar_version(uri: &str, seed: i32) {
7526 use arrow::array::{Int32Array, StringArray};
7527 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
7528 let schema = Arc::new(ArrowSchema::new(vec![
7529 Field::new("id", DataType::Int32, false),
7530 Field::new("name", DataType::Utf8, true),
7531 ]));
7532 let batch = arrow::record_batch::RecordBatch::try_new(
7533 schema.clone(),
7534 vec![
7535 Arc::new(Int32Array::from(vec![seed, seed + 1])),
7536 Arc::new(StringArray::from(vec![Some("x"), Some("y")])),
7537 ],
7538 )
7539 .unwrap();
7540 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
7541 Dataset::write(
7542 reader,
7543 uri,
7544 Some(WriteParams {
7545 mode: WriteMode::Append,
7546 ..Default::default()
7547 }),
7548 )
7549 .await
7550 .unwrap();
7551 }
7552
7553 async fn list_versions(
7555 namespace: &DirectoryNamespace,
7556 table_name: &str,
7557 branch: Option<&str>,
7558 ) -> Result<Vec<TableVersion>> {
7559 let req = ListTableVersionsRequest {
7560 id: Some(vec![table_name.to_string()]),
7561 branch: branch.map(|b| b.to_string()),
7562 ..Default::default()
7563 };
7564 namespace.list_table_versions(req).await.map(|r| r.versions)
7565 }
7566
7567 #[tokio::test]
7568 async fn test_list_table_versions_on_branch() {
7569 let (namespace, _temp_dir) = create_test_namespace().await;
7570 create_scalar_table(&namespace, "users").await;
7571 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7572
7573 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7576 .await
7577 .unwrap();
7578 assert!(branch_versions.len() >= 2);
7579 assert!(
7580 branch_versions
7581 .iter()
7582 .all(|v| v.manifest_path.contains("tree/exp")),
7583 "branch versions must resolve to branch manifests: {:?}",
7584 branch_versions
7585 );
7586
7587 let main_versions = list_versions(&namespace, "users", None).await.unwrap();
7589 let main_explicit = list_versions(&namespace, "users", Some("main"))
7590 .await
7591 .unwrap();
7592 assert_eq!(main_versions.len(), main_explicit.len());
7593 assert!(
7594 main_versions
7595 .iter()
7596 .all(|v| !v.manifest_path.contains("tree/"))
7597 );
7598
7599 let missing = list_versions(&namespace, "users", Some("does-not-exist")).await;
7601 assert!(missing.is_err());
7602 assert!(missing.unwrap_err().to_string().contains("not found"));
7603 }
7604
7605 #[tokio::test]
7606 async fn test_describe_table_version_on_branch() {
7607 let (namespace, _temp_dir) = create_test_namespace().await;
7608 create_scalar_table(&namespace, "users").await;
7609 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7610
7611 let branch_versions = list_versions(&namespace, "users", Some("exp"))
7612 .await
7613 .unwrap();
7614 let latest = branch_versions.iter().map(|v| v.version).max().unwrap();
7615
7616 let req = DescribeTableVersionRequest {
7618 id: Some(vec!["users".to_string()]),
7619 branch: Some("exp".to_string()),
7620 ..Default::default()
7621 };
7622 let resp = namespace.describe_table_version(req).await.unwrap();
7623 assert_eq!(resp.version.version, latest);
7624 assert!(resp.version.manifest_path.contains("tree/exp"));
7625
7626 let req = DescribeTableVersionRequest {
7628 id: Some(vec!["users".to_string()]),
7629 version: Some(latest),
7630 branch: Some("exp".to_string()),
7631 ..Default::default()
7632 };
7633 assert!(namespace.describe_table_version(req).await.is_ok());
7634
7635 let req = DescribeTableVersionRequest {
7637 id: Some(vec!["users".to_string()]),
7638 version: Some(999_999),
7639 branch: Some("exp".to_string()),
7640 ..Default::default()
7641 };
7642 assert!(namespace.describe_table_version(req).await.is_err());
7643
7644 let req = DescribeTableVersionRequest {
7646 id: Some(vec!["users".to_string()]),
7647 branch: Some("nope".to_string()),
7648 ..Default::default()
7649 };
7650 let err = namespace.describe_table_version(req).await;
7651 assert!(err.is_err() && err.unwrap_err().to_string().contains("not found"));
7652 }
7653
7654 #[tokio::test]
7655 async fn test_restore_table_on_branch() {
7656 use lance_namespace::models::RestoreTableRequest;
7657
7658 let (namespace, _temp_dir) = create_test_namespace().await;
7659 create_scalar_table(&namespace, "users").await;
7660 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7661
7662 let before = list_versions(&namespace, "users", Some("exp"))
7663 .await
7664 .unwrap();
7665 let branch_latest = before.iter().map(|v| v.version).max().unwrap();
7666 let earliest = before.iter().map(|v| v.version).min().unwrap();
7667 let main_before = list_versions(&namespace, "users", None)
7668 .await
7669 .unwrap()
7670 .len();
7671
7672 let req = RestoreTableRequest {
7675 id: Some(vec!["users".to_string()]),
7676 version: earliest,
7677 branch: Some("exp".to_string()),
7678 ..Default::default()
7679 };
7680 let resp = namespace.restore_table(req).await.unwrap();
7681 assert!(resp.transaction_id.is_some());
7682
7683 let after = list_versions(&namespace, "users", Some("exp"))
7684 .await
7685 .unwrap();
7686 let new_latest = after.iter().map(|v| v.version).max().unwrap();
7687 assert!(
7688 new_latest > branch_latest,
7689 "restore should add a branch version"
7690 );
7691
7692 let main_after = list_versions(&namespace, "users", None)
7693 .await
7694 .unwrap()
7695 .len();
7696 assert_eq!(main_after, main_before, "main must be unaffected");
7697 }
7698
7699 #[tokio::test]
7700 async fn test_batch_delete_table_versions_on_branch() {
7701 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
7702
7703 let (namespace, _temp_dir) = create_test_namespace().await;
7704 create_scalar_table(&namespace, "users").await;
7705 create_branch_with_commits(&namespace, "users", "exp", 2).await;
7706
7707 let before = list_versions(&namespace, "users", Some("exp"))
7708 .await
7709 .unwrap();
7710 let main_before = list_versions(&namespace, "users", None).await.unwrap();
7711
7712 let req = BatchDeleteTableVersionsRequest {
7717 id: Some(vec!["users".to_string()]),
7718 branch: Some("exp".to_string()),
7719 ranges: vec![VersionRange::new(0, -1)],
7720 ..Default::default()
7721 };
7722 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
7723 assert_eq!(
7724 resp.deleted_count,
7725 Some(before.len() as i64),
7726 "every branch manifest should be physically deleted"
7727 );
7728
7729 assert!(
7731 list_versions(&namespace, "users", Some("exp"))
7732 .await
7733 .is_err()
7734 );
7735 let main_after = list_versions(&namespace, "users", None).await.unwrap();
7736 assert_eq!(
7737 main_after.len(),
7738 main_before.len(),
7739 "main must be untouched"
7740 );
7741 }
7742
7743 #[tokio::test]
7744 async fn test_create_table_version_on_branch() {
7745 use futures::TryStreamExt;
7746 use lance_namespace::models::CreateTableVersionRequest;
7747
7748 let (namespace, _temp_dir) = create_test_namespace().await;
7749 create_scalar_table(&namespace, "users").await;
7750 let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 1).await;
7751
7752 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7754 let versions_dir = branch_ds.versions_dir();
7755 let store = branch_ds.object_store(None).await.unwrap();
7756 let existing = store
7757 .inner
7758 .list(Some(&versions_dir))
7759 .try_collect::<Vec<_>>()
7760 .await
7761 .unwrap()
7762 .into_iter()
7763 .find(|m| {
7764 m.location
7765 .filename()
7766 .map(|f| f.ends_with(".manifest"))
7767 .unwrap_or(false)
7768 })
7769 .expect("a branch manifest");
7770 let bytes = store
7771 .inner
7772 .get(&existing.location)
7773 .await
7774 .unwrap()
7775 .bytes()
7776 .await
7777 .unwrap();
7778 let staging = versions_dir.join("staging_manifest");
7779 store.inner.put(&staging, bytes.into()).await.unwrap();
7780
7781 let main_before = list_versions(&namespace, "users", None)
7782 .await
7783 .unwrap()
7784 .len();
7785 let new_version = list_versions(&namespace, "users", Some("exp"))
7786 .await
7787 .unwrap()
7788 .iter()
7789 .map(|v| v.version)
7790 .max()
7791 .unwrap()
7792 + 1;
7793
7794 let req = CreateTableVersionRequest {
7795 id: Some(vec!["users".to_string()]),
7796 version: new_version,
7797 manifest_path: staging.to_string(),
7798 naming_scheme: Some("V2".to_string()),
7799 branch: Some("exp".to_string()),
7800 ..Default::default()
7801 };
7802 let resp = namespace.create_table_version(req).await.unwrap();
7803 let info = resp.version.expect("version info");
7804 assert!(
7806 info.manifest_path.contains("tree/exp"),
7807 "got {}",
7808 info.manifest_path
7809 );
7810
7811 let after = list_versions(&namespace, "users", Some("exp"))
7813 .await
7814 .unwrap();
7815 assert!(after.iter().any(|v| v.version == new_version));
7816 let main_after = list_versions(&namespace, "users", None)
7817 .await
7818 .unwrap()
7819 .len();
7820 assert_eq!(main_after, main_before, "main must be unaffected");
7821 }
7822
7823 #[tokio::test]
7828 async fn test_external_manifest_store_resolves_branch_from_base_path() {
7829 use futures::TryStreamExt;
7830 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
7831 use lance_table::io::commit::external_manifest::ExternalManifestStore;
7832
7833 let (namespace, _temp_dir) = create_test_namespace().await;
7834 create_scalar_table(&namespace, "users").await; let branch_uri = create_branch_with_commits(&namespace, "users", "exp", 2).await;
7836
7837 let namespace = Arc::new(namespace);
7838 let table_id = vec!["users".to_string()];
7839 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
7840 let branch_base = branch_ds.branch_location().path;
7841 let root_base = branch_ds.branch_location().find_main().unwrap().path;
7842 let store = LanceNamespaceExternalManifestStore::new(
7843 namespace.clone(),
7844 table_id.clone(),
7845 root_base.clone(),
7846 );
7847
7848 let (branch_latest, branch_path) = store
7851 .get_latest_version(branch_base.as_ref())
7852 .await
7853 .unwrap()
7854 .expect("branch has versions");
7855 let (_main_latest, main_path) = store
7856 .get_latest_version(root_base.as_ref())
7857 .await
7858 .unwrap()
7859 .expect("main has versions");
7860 assert!(
7861 branch_path.contains("tree/exp"),
7862 "branch latest must resolve to the branch tree: {}",
7863 branch_path
7864 );
7865 assert!(
7866 !main_path.contains("tree/exp"),
7867 "main latest must not resolve to a branch tree: {}",
7868 main_path
7869 );
7870
7871 let described = store
7873 .get(branch_base.as_ref(), branch_latest)
7874 .await
7875 .unwrap();
7876 assert!(
7877 described.contains("tree/exp"),
7878 "describe on the branch must resolve to the branch tree: {}",
7879 described
7880 );
7881
7882 assert!(store.get_latest_version("somewhere/else").await.is_err());
7884
7885 let versions_dir = branch_ds.versions_dir();
7888 let obj = branch_ds.object_store(None).await.unwrap();
7889 let existing = obj
7890 .inner
7891 .list(Some(&versions_dir))
7892 .try_collect::<Vec<_>>()
7893 .await
7894 .unwrap()
7895 .into_iter()
7896 .find(|m| {
7897 m.location
7898 .filename()
7899 .map(|f| f.ends_with(".manifest"))
7900 .unwrap_or(false)
7901 })
7902 .expect("a branch manifest");
7903 let bytes = obj
7904 .inner
7905 .get(&existing.location)
7906 .await
7907 .unwrap()
7908 .bytes()
7909 .await
7910 .unwrap();
7911 let size = bytes.len() as u64;
7912 let staging = versions_dir.clone().join("staging_manifest");
7913 obj.inner.put(&staging, bytes.into()).await.unwrap();
7914
7915 let committed = store
7916 .put(
7917 &branch_base,
7918 branch_latest + 1,
7919 &staging,
7920 size,
7921 None,
7922 obj.inner.as_ref(),
7923 ManifestNamingScheme::V2,
7924 )
7925 .await
7926 .unwrap();
7927 assert!(
7928 committed.path.to_string().contains("tree/exp"),
7929 "a commit through a branch-qualified base must land on the branch tree: {}",
7930 committed.path
7931 );
7932 }
7933
7934 #[tokio::test]
7942 async fn test_write_into_namespace_on_branch_appends_to_branch() {
7943 use lance::dataset::builder::DatasetBuilder;
7944 use lance_namespace::models::CreateTableBranchRequest;
7945
7946 let temp = TempStdDir::default();
7947 let namespace = Arc::new(
7948 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
7949 .manifest_enabled(true)
7950 .table_version_tracking_enabled(true)
7951 .ops_metrics_enabled(true)
7952 .build()
7953 .await
7954 .unwrap(),
7955 );
7956 let ns: Arc<dyn LanceNamespace> = namespace.clone();
7957 let table_id = vec!["t".to_string()];
7958 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
7960 id: Some(table_id.clone()),
7961 name: "exp".to_string(),
7962 ..Default::default()
7963 })
7964 .await
7965 .unwrap();
7966
7967 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
7968 ns.list_table_versions(ListTableVersionsRequest {
7969 id: Some(table_id),
7970 ..Default::default()
7971 })
7972 .await
7973 .unwrap()
7974 .versions
7975 .len()
7976 };
7977 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
7978 let commits_before = namespace
7979 .retrieve_ops_metrics()
7980 .get("create_table_version")
7981 .copied()
7982 .unwrap_or(0);
7983
7984 let branch_ds = Dataset::write_into_namespace_on_branch(
7985 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
7986 ns.clone(),
7987 table_id.clone(),
7988 "exp",
7989 Some(WriteParams {
7990 mode: WriteMode::Append,
7991 ..Default::default()
7992 }),
7993 )
7994 .await
7995 .unwrap();
7996 assert_eq!(branch_ds.manifest.branch.as_deref(), Some("exp"));
7997 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
7998
7999 let commits_after = namespace
8002 .retrieve_ops_metrics()
8003 .get("create_table_version")
8004 .copied()
8005 .unwrap_or(0);
8006 assert_eq!(
8007 commits_after,
8008 commits_before + 1,
8009 "the branch append must register its version via create_table_version"
8010 );
8011 let exp_versions = ns
8012 .list_table_versions(ListTableVersionsRequest {
8013 id: Some(table_id.clone()),
8014 branch: Some("exp".to_string()),
8015 ..Default::default()
8016 })
8017 .await
8018 .unwrap()
8019 .versions;
8020 assert!(
8021 exp_versions
8022 .iter()
8023 .all(|v| v.manifest_path.contains("tree/exp")),
8024 "branch versions must resolve to the branch tree: {:?}",
8025 exp_versions
8026 );
8027 assert_eq!(
8028 main_chain_len(ns.clone(), table_id.clone()).await,
8029 main_before,
8030 "main's catalog must be untouched by the branch append"
8031 );
8032
8033 Dataset::write_into_namespace(
8036 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
8037 ns.clone(),
8038 table_id.clone(),
8039 Some(WriteParams {
8040 mode: WriteMode::Append,
8041 ..Default::default()
8042 }),
8043 )
8044 .await
8045 .unwrap();
8046 assert_eq!(
8047 main_chain_len(ns.clone(), table_id.clone()).await,
8048 main_before + 1,
8049 "a managed main append must register its version in the catalog"
8050 );
8051 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8052 .await
8053 .unwrap()
8054 .load()
8055 .await
8056 .unwrap();
8057 assert_eq!(
8058 scan_id_column(&fresh).await,
8059 vec![1, 2, 100],
8060 "a fresh managed open must resolve the appended version, not a stale latest"
8061 );
8062 }
8063
8064 #[tokio::test]
8066 async fn test_write_into_namespace_on_branch_rejects_create() {
8067 use arrow::array::{Int32Array, StringArray};
8068 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
8069
8070 let (namespace, _temp_dir) = create_test_namespace().await;
8071 let namespace = Arc::new(namespace);
8072
8073 let schema = Arc::new(ArrowSchema::new(vec![
8074 Field::new("id", DataType::Int32, false),
8075 Field::new("name", DataType::Utf8, true),
8076 ]));
8077 let batch = arrow::record_batch::RecordBatch::try_new(
8078 schema.clone(),
8079 vec![
8080 Arc::new(Int32Array::from(vec![1])),
8081 Arc::new(StringArray::from(vec![Some("a")])),
8082 ],
8083 )
8084 .unwrap();
8085 let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
8086
8087 let result = Dataset::write_into_namespace_on_branch(
8088 reader,
8089 namespace.clone(),
8090 vec!["new_table".to_string()],
8091 "exp",
8092 Some(WriteParams {
8093 mode: WriteMode::Create,
8094 ..Default::default()
8095 }),
8096 )
8097 .await;
8098 assert!(result.is_err(), "create on a branch must be rejected");
8099 assert!(
8100 result.unwrap_err().to_string().contains("branch"),
8101 "error should mention the branch restriction"
8102 );
8103 }
8104
8105 #[tokio::test]
8106 async fn test_branch_name_validation_rejects_traversal() {
8107 let (namespace, _temp_dir) = create_test_namespace().await;
8108 create_scalar_table(&namespace, "users").await;
8109
8110 let err = list_versions(&namespace, "users", Some("../evil")).await;
8113 assert!(err.is_err());
8114 assert!(err.unwrap_err().to_string().contains("invalid branch name"));
8115 }
8116
8117 #[tokio::test]
8118 async fn test_branch_ops_reject_zombie_branch() {
8119 use futures::TryStreamExt;
8120 use lance_namespace::models::{
8121 BatchDeleteTableVersionsRequest, CreateTableVersionRequest, RestoreTableRequest,
8122 VersionRange,
8123 };
8124
8125 let (namespace, _temp_dir) = create_test_namespace().await;
8126 create_scalar_table(&namespace, "users").await;
8127
8128 let dataset = open_dataset(&namespace, "users").await;
8129 let store = dataset.object_store(None).await.unwrap();
8130 let manifest = store
8131 .inner
8132 .list(Some(&dataset.versions_dir()))
8133 .try_collect::<Vec<_>>()
8134 .await
8135 .unwrap()
8136 .into_iter()
8137 .find(|m| {
8138 m.location
8139 .filename()
8140 .map(|f| f.ends_with(".manifest"))
8141 .unwrap_or(false)
8142 })
8143 .expect("a manifest");
8144 let bytes = store
8145 .inner
8146 .get(&manifest.location)
8147 .await
8148 .unwrap()
8149 .bytes()
8150 .await
8151 .unwrap();
8152 let zombie = dataset
8153 .branch_location()
8154 .find_branch(Some("ghost"))
8155 .unwrap()
8156 .path
8157 .join(VERSIONS_DIR)
8158 .join(manifest.location.filename().unwrap());
8159 store.inner.put(&zombie, bytes.into()).await.unwrap();
8160
8161 assert!(dataset.branches().get("ghost").await.is_err());
8162
8163 fn rejected<T: std::fmt::Debug>(label: &str, r: Result<T>) {
8164 match r {
8165 Ok(v) => panic!("{label} must reject the zombie branch, got Ok({v:?})"),
8166 Err(e) => assert!(e.to_string().contains("not found"), "{label}: {e}"),
8167 }
8168 }
8169
8170 rejected(
8171 "list",
8172 list_versions(&namespace, "users", Some("ghost")).await,
8173 );
8174 rejected(
8175 "describe",
8176 namespace
8177 .describe_table_version(DescribeTableVersionRequest {
8178 id: Some(vec!["users".to_string()]),
8179 branch: Some("ghost".to_string()),
8180 ..Default::default()
8181 })
8182 .await,
8183 );
8184 rejected(
8185 "create",
8186 namespace
8187 .create_table_version(CreateTableVersionRequest {
8188 id: Some(vec!["users".to_string()]),
8189 version: 2,
8190 manifest_path: zombie.to_string(),
8191 branch: Some("ghost".to_string()),
8192 ..Default::default()
8193 })
8194 .await,
8195 );
8196 rejected(
8197 "restore",
8198 namespace
8199 .restore_table(RestoreTableRequest {
8200 id: Some(vec!["users".to_string()]),
8201 version: 1,
8202 branch: Some("ghost".to_string()),
8203 ..Default::default()
8204 })
8205 .await,
8206 );
8207 rejected(
8208 "batch_delete",
8209 namespace
8210 .batch_delete_table_versions(BatchDeleteTableVersionsRequest {
8211 id: Some(vec!["users".to_string()]),
8212 branch: Some("ghost".to_string()),
8213 ranges: vec![VersionRange::new(1, 1)],
8214 ..Default::default()
8215 })
8216 .await,
8217 );
8218 }
8219
8220 #[tokio::test]
8225 async fn test_batch_delete_table_versions_main_v2() {
8226 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
8227
8228 let (namespace, _temp_dir) = create_test_namespace().await;
8229 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
8231 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
8235 assert!(before.len() >= 3);
8236 assert!(
8239 before
8240 .iter()
8241 .all(|v| v.manifest_path.rsplit('/').next().unwrap().len() == 29),
8242 "expected V2-named manifests: {:?}",
8243 before
8244 );
8245 let min_v = before.iter().map(|v| v.version).min().unwrap();
8246 let max_v = before.iter().map(|v| v.version).max().unwrap();
8247
8248 let req = BatchDeleteTableVersionsRequest {
8251 id: Some(vec!["users".to_string()]),
8252 ranges: vec![VersionRange::new(min_v, max_v)],
8253 ..Default::default()
8254 };
8255 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
8256 assert_eq!(
8257 resp.deleted_count,
8258 Some((before.len() - 1) as i64),
8259 "V2 manifests must actually be deleted (was 0 before the fix)"
8260 );
8261
8262 let after = list_versions(&namespace, "users", None).await.unwrap();
8263 assert_eq!(after.len(), 1);
8264 assert_eq!(after[0].version, max_v);
8265 }
8266
8267 #[tokio::test]
8269 async fn test_batch_delete_end_is_exclusive() {
8270 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
8271
8272 let (namespace, _temp_dir) = create_test_namespace().await;
8273 create_scalar_table(&namespace, "users").await; let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
8275 append_scalar_version(&main_uri, 100).await; append_scalar_version(&main_uri, 200).await; let before = list_versions(&namespace, "users", None).await.unwrap();
8279 let min_v = before.iter().map(|v| v.version).min().unwrap();
8280
8281 let req = BatchDeleteTableVersionsRequest {
8282 id: Some(vec!["users".to_string()]),
8283 ranges: vec![VersionRange::new(min_v, min_v + 1)],
8284 ..Default::default()
8285 };
8286 let resp = namespace.batch_delete_table_versions(req).await.unwrap();
8287 assert_eq!(
8288 resp.deleted_count,
8289 Some(1),
8290 "only min_v is in [min_v, min_v+1)"
8291 );
8292
8293 let after = list_versions(&namespace, "users", None).await.unwrap();
8294 assert!(
8295 !after.iter().any(|v| v.version == min_v),
8296 "min_v must be deleted"
8297 );
8298 assert_eq!(after.len(), before.len() - 1, "exactly one version removed");
8299 }
8300
8301 #[tokio::test]
8302 async fn test_batch_delete_rejects_unbounded_range() {
8303 use lance_namespace::models::{BatchDeleteTableVersionsRequest, VersionRange};
8304
8305 let (namespace, _temp_dir) = create_test_namespace().await;
8306 create_scalar_table(&namespace, "users").await;
8307
8308 let req = BatchDeleteTableVersionsRequest {
8311 id: Some(vec!["users".to_string()]),
8312 ranges: vec![VersionRange::new(0, i64::MAX)],
8313 ..Default::default()
8314 };
8315 let err = namespace.batch_delete_table_versions(req).await;
8316 assert!(err.is_err());
8317 assert!(
8318 err.unwrap_err().to_string().contains("limit"),
8319 "expected a range-too-large error"
8320 );
8321 }
8322
8323 async fn create_managed_namespace(path: &str) -> Arc<dyn LanceNamespace> {
8325 Arc::new(
8326 DirectoryNamespaceBuilder::new(path)
8327 .manifest_enabled(true)
8328 .table_version_tracking_enabled(true)
8329 .build()
8330 .await
8331 .unwrap(),
8332 )
8333 }
8334
8335 fn single_int_schema() -> Arc<arrow::datatypes::Schema> {
8336 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
8337 Arc::new(ArrowSchema::new(vec![Field::new(
8338 "id",
8339 DataType::Int32,
8340 false,
8341 )]))
8342 }
8343
8344 fn single_int_batch(seed: i32) -> arrow::record_batch::RecordBatch {
8345 use arrow::array::Int32Array;
8346 arrow::record_batch::RecordBatch::try_new(
8347 single_int_schema(),
8348 vec![Arc::new(Int32Array::from(vec![seed]))],
8349 )
8350 .unwrap()
8351 }
8352
8353 async fn create_managed_table(ns: &Arc<dyn LanceNamespace>, table_id: &[String]) -> Dataset {
8356 let mut ds = Dataset::write_into_namespace(
8357 RecordBatchIterator::new(vec![Ok(single_int_batch(1))], single_int_schema()),
8358 ns.clone(),
8359 table_id.to_vec(),
8360 Some(WriteParams {
8361 mode: WriteMode::Create,
8362 ..Default::default()
8363 }),
8364 )
8365 .await
8366 .unwrap();
8367 ds.append(
8368 RecordBatchIterator::new(vec![Ok(single_int_batch(2))], single_int_schema()),
8369 None,
8370 )
8371 .await
8372 .unwrap();
8373 ds
8374 }
8375
8376 async fn scan_id_column(ds: &Dataset) -> Vec<i32> {
8378 use arrow::array::Int32Array;
8379 use futures::TryStreamExt;
8380 let batches: Vec<arrow::record_batch::RecordBatch> = ds
8381 .scan()
8382 .try_into_stream()
8383 .await
8384 .unwrap()
8385 .try_collect()
8386 .await
8387 .unwrap();
8388 let mut ids: Vec<i32> = batches
8389 .iter()
8390 .flat_map(|b| {
8391 b.column(0)
8392 .as_any()
8393 .downcast_ref::<Int32Array>()
8394 .unwrap()
8395 .values()
8396 .to_vec()
8397 })
8398 .collect();
8399 ids.sort();
8400 ids
8401 }
8402
8403 #[tokio::test]
8408 async fn test_managed_branch_open_and_commit() {
8409 use futures::TryStreamExt;
8410 use lance::dataset::builder::DatasetBuilder;
8411 use lance_namespace::models::CreateTableBranchRequest;
8412
8413 let temp = TempStdDir::default();
8414 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8415 let table_id = vec!["t".to_string()];
8416 create_managed_table(&ns, &table_id).await;
8417 let main_before = ns
8418 .list_table_versions(ListTableVersionsRequest {
8419 id: Some(table_id.clone()),
8420 ..Default::default()
8421 })
8422 .await
8423 .unwrap()
8424 .versions
8425 .len();
8426
8427 ns.create_table_branch(CreateTableBranchRequest {
8430 id: Some(table_id.clone()),
8431 name: "exp".to_string(),
8432 ..Default::default()
8433 })
8434 .await
8435 .unwrap();
8436
8437 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8440 .await
8441 .unwrap()
8442 .with_branch("exp", None)
8443 .load()
8444 .await
8445 .unwrap();
8446 assert_eq!(
8447 branch_ds.manifest.branch.as_deref(),
8448 Some("exp"),
8449 "with_branch on a managed table must open the branch chain"
8450 );
8451 let branch_base = branch_ds.branch_location().path;
8452 assert!(
8453 branch_base.as_ref().ends_with("tree/exp"),
8454 "the branch dataset must be rooted at the branch chain: {}",
8455 branch_base
8456 );
8457 let branch_v_before = branch_ds.version().version;
8458
8459 branch_ds
8461 .append(
8462 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8463 None,
8464 )
8465 .await
8466 .unwrap();
8467 assert_eq!(
8468 branch_ds.manifest.branch.as_deref(),
8469 Some("exp"),
8470 "the commit must stay on the branch"
8471 );
8472 assert!(
8473 branch_ds.version().version > branch_v_before,
8474 "the branch version must advance after the commit"
8475 );
8476 assert_eq!(scan_id_column(&branch_ds).await, vec![1, 2, 3]);
8477
8478 let store = branch_ds.object_store(None).await.unwrap();
8482 let branch_data = branch_base.clone().join("data");
8483 let branch_files = store
8484 .inner
8485 .list(Some(&branch_data))
8486 .try_collect::<Vec<_>>()
8487 .await
8488 .unwrap();
8489 assert!(
8490 !branch_files.is_empty(),
8491 "the branch commit must place data files under the branch chain"
8492 );
8493
8494 let table_uri = ns
8496 .describe_table(DescribeTableRequest {
8497 id: Some(table_id.clone()),
8498 ..Default::default()
8499 })
8500 .await
8501 .unwrap()
8502 .location
8503 .unwrap();
8504 let fs_branch_ds = DatasetBuilder::from_uri(&table_uri)
8505 .with_branch("exp", None)
8506 .load()
8507 .await
8508 .unwrap();
8509 assert_eq!(fs_branch_ds.manifest.branch.as_deref(), Some("exp"));
8510 assert_eq!(scan_id_column(&fs_branch_ds).await, vec![1, 2, 3]);
8511
8512 let main_after = ns
8515 .list_table_versions(ListTableVersionsRequest {
8516 id: Some(table_id.clone()),
8517 ..Default::default()
8518 })
8519 .await
8520 .unwrap()
8521 .versions
8522 .len();
8523 assert_eq!(
8524 main_after, main_before,
8525 "committing on the branch must not change main's chain"
8526 );
8527 let main_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8528 .await
8529 .unwrap()
8530 .load()
8531 .await
8532 .unwrap();
8533 assert_eq!(main_ds.manifest.branch, None);
8534 assert_eq!(scan_id_column(&main_ds).await, vec![1, 2]);
8535 }
8536
8537 #[tokio::test]
8542 async fn test_managed_branch_tags() {
8543 use lance::dataset::builder::DatasetBuilder;
8544 use lance::dataset::refs::Ref;
8545 use lance_namespace::models::CreateTableBranchRequest;
8546
8547 let temp = TempStdDir::default();
8548 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8549 let table_id = vec!["t".to_string()];
8550 let main_ds = create_managed_table(&ns, &table_id).await;
8551 ns.create_table_branch(CreateTableBranchRequest {
8552 id: Some(table_id.clone()),
8553 name: "exp".to_string(),
8554 ..Default::default()
8555 })
8556 .await
8557 .unwrap();
8558 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8559 .await
8560 .unwrap()
8561 .with_branch("exp", None)
8562 .load()
8563 .await
8564 .unwrap();
8565 branch_ds
8566 .append(
8567 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8568 None,
8569 )
8570 .await
8571 .unwrap();
8572 let branch_version = branch_ds.version().version;
8573
8574 main_ds
8577 .tags()
8578 .create("exp-tag", ("exp", Some(branch_version)))
8579 .await
8580 .unwrap();
8581 let tag = main_ds.tags().get("exp-tag").await.unwrap();
8582 assert_eq!(tag.branch.as_deref(), Some("exp"));
8583 assert_eq!(tag.version, branch_version);
8584
8585 branch_ds
8587 .tags()
8588 .create("exp-tag2", branch_version)
8589 .await
8590 .unwrap();
8591 let tag2 = branch_ds.tags().get("exp-tag2").await.unwrap();
8592 assert_eq!(tag2.branch.as_deref(), Some("exp"));
8593
8594 let tag_open = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8597 .await
8598 .unwrap()
8599 .with_tag("exp-tag")
8600 .load()
8601 .await
8602 .unwrap();
8603 assert_eq!(tag_open.manifest.branch.as_deref(), Some("exp"));
8604 assert_eq!(tag_open.version().version, branch_version);
8605 assert_eq!(scan_id_column(&tag_open).await, vec![1, 2, 3]);
8606
8607 let tag_checkout = main_ds
8609 .checkout_version(Ref::Tag("exp-tag".to_string()))
8610 .await
8611 .unwrap();
8612 assert_eq!(tag_checkout.manifest.branch.as_deref(), Some("exp"));
8613 assert_eq!(scan_id_column(&tag_checkout).await, vec![1, 2, 3]);
8614
8615 let err = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8617 .await
8618 .unwrap()
8619 .with_tag("no-such-tag")
8620 .load()
8621 .await;
8622 assert!(err.is_err(), "a missing tag must error");
8623 }
8624
8625 #[tokio::test]
8630 async fn test_managed_cross_branch_checkout() {
8631 use lance::dataset::builder::DatasetBuilder;
8632 use lance::dataset::refs::Ref;
8633 use lance_namespace::models::CreateTableBranchRequest;
8634
8635 let temp = TempStdDir::default();
8636 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8637 let table_id = vec!["t".to_string()];
8638 let mut main_ds = create_managed_table(&ns, &table_id).await;
8639 ns.create_table_branch(CreateTableBranchRequest {
8640 id: Some(table_id.clone()),
8641 name: "exp".to_string(),
8642 ..Default::default()
8643 })
8644 .await
8645 .unwrap();
8646
8647 let mut branch_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8649 .await
8650 .unwrap()
8651 .with_branch("exp", None)
8652 .load()
8653 .await
8654 .unwrap();
8655 branch_ds
8656 .append(
8657 RecordBatchIterator::new(vec![Ok(single_int_batch(3))], single_int_schema()),
8658 None,
8659 )
8660 .await
8661 .unwrap();
8662 let overlap_version = branch_ds.version().version;
8663 while main_ds.version().version < overlap_version {
8664 main_ds
8665 .append(
8666 RecordBatchIterator::new(vec![Ok(single_int_batch(100))], single_int_schema()),
8667 None,
8668 )
8669 .await
8670 .unwrap();
8671 }
8672
8673 let on_branch = main_ds
8676 .checkout_version(Ref::Version(Some("exp".to_string()), Some(overlap_version)))
8677 .await
8678 .unwrap();
8679 assert_eq!(on_branch.manifest.branch.as_deref(), Some("exp"));
8680 assert_eq!(scan_id_column(&on_branch).await, vec![1, 2, 3]);
8681
8682 let mut on_branch_latest = main_ds.checkout_branch("exp").await.unwrap();
8684 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8685 assert_eq!(on_branch_latest.version().version, overlap_version);
8686
8687 let main_chain_len = |ns: Arc<dyn LanceNamespace>, table_id: Vec<String>| async move {
8690 ns.list_table_versions(ListTableVersionsRequest {
8691 id: Some(table_id),
8692 ..Default::default()
8693 })
8694 .await
8695 .unwrap()
8696 .versions
8697 .len()
8698 };
8699 let main_before = main_chain_len(ns.clone(), table_id.clone()).await;
8700 on_branch_latest
8701 .append(
8702 RecordBatchIterator::new(vec![Ok(single_int_batch(4))], single_int_schema()),
8703 None,
8704 )
8705 .await
8706 .unwrap();
8707 assert_eq!(on_branch_latest.manifest.branch.as_deref(), Some("exp"));
8708 assert_eq!(scan_id_column(&on_branch_latest).await, vec![1, 2, 3, 4]);
8709 assert_eq!(
8710 main_chain_len(ns.clone(), table_id.clone()).await,
8711 main_before,
8712 "a commit on the checked-out branch must not advance main's chain"
8713 );
8714
8715 let on_main = branch_ds
8717 .checkout_version(Ref::Version(None, Some(1)))
8718 .await
8719 .unwrap();
8720 assert_eq!(on_main.manifest.branch, None);
8721 assert_eq!(scan_id_column(&on_main).await, vec![1]);
8722
8723 ns.create_table_branch(CreateTableBranchRequest {
8725 id: Some(table_id.clone()),
8726 name: "exp2".to_string(),
8727 ..Default::default()
8728 })
8729 .await
8730 .unwrap();
8731 let on_branch2 = branch_ds.checkout_branch("exp2").await.unwrap();
8732 assert_eq!(on_branch2.manifest.branch.as_deref(), Some("exp2"));
8733
8734 let err = main_ds
8736 .checkout_version(Ref::Version(Some("exp".to_string()), Some(999)))
8737 .await;
8738 assert!(err.is_err(), "a version missing from the branch must error");
8739 }
8740
8741 #[tokio::test]
8747 async fn test_commit_builder_honors_explicit_handler_for_dataset_dest() {
8748 use lance::dataset::write::{CommitBuilder, InsertBuilder};
8749 use lance::dataset::{WriteDestination, builder::DatasetBuilder};
8750 use lance::io::commit::namespace_manifest::LanceNamespaceExternalManifestStore;
8751 use lance_table::io::commit::external_manifest::ExternalManifestCommitHandler;
8752
8753 let temp = TempStdDir::default();
8754 let namespace = Arc::new(
8755 DirectoryNamespaceBuilder::new(temp.to_str().unwrap())
8756 .manifest_enabled(true)
8757 .table_version_tracking_enabled(true)
8758 .ops_metrics_enabled(true)
8759 .build()
8760 .await
8761 .unwrap(),
8762 );
8763 let ns: Arc<dyn LanceNamespace> = namespace.clone();
8764 let table_id = vec!["t".to_string()];
8765 create_managed_table(&ns, &table_id).await; let table_uri = ns
8769 .describe_table(DescribeTableRequest {
8770 id: Some(table_id.clone()),
8771 ..Default::default()
8772 })
8773 .await
8774 .unwrap()
8775 .location
8776 .unwrap();
8777 let plain_ds = Arc::new(Dataset::open(&table_uri).await.unwrap());
8778
8779 let transaction = InsertBuilder::new(WriteDestination::Dataset(plain_ds.clone()))
8780 .with_params(&WriteParams {
8781 mode: WriteMode::Append,
8782 ..Default::default()
8783 })
8784 .execute_uncommitted(vec![single_int_batch(3)])
8785 .await
8786 .unwrap();
8787
8788 let handler = Arc::new(ExternalManifestCommitHandler {
8789 external_manifest_store: Arc::new(
8790 LanceNamespaceExternalManifestStore::for_table_uri(
8791 ns.clone(),
8792 table_id.clone(),
8793 &table_uri,
8794 )
8795 .unwrap(),
8796 ),
8797 });
8798 let commits_before = namespace
8799 .retrieve_ops_metrics()
8800 .get("create_table_version")
8801 .copied()
8802 .unwrap_or(0);
8803 let committed = CommitBuilder::new(WriteDestination::Dataset(plain_ds))
8804 .with_commit_handler(handler)
8805 .execute(transaction)
8806 .await
8807 .unwrap();
8808 assert_eq!(scan_id_column(&committed).await, vec![1, 2, 3]);
8809
8810 let commits_after = namespace
8811 .retrieve_ops_metrics()
8812 .get("create_table_version")
8813 .copied()
8814 .unwrap_or(0);
8815 assert_eq!(
8816 commits_after,
8817 commits_before + 1,
8818 "the explicit handler must route the commit through create_table_version"
8819 );
8820 let fresh = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8821 .await
8822 .unwrap()
8823 .load()
8824 .await
8825 .unwrap();
8826 assert_eq!(
8827 scan_id_column(&fresh).await,
8828 vec![1, 2, 3],
8829 "a fresh managed open must resolve the committed version"
8830 );
8831 }
8832
8833 #[tokio::test]
8835 async fn test_managed_branch_from_non_latest_fork() {
8836 use lance::dataset::builder::DatasetBuilder;
8837 use lance_namespace::models::CreateTableBranchRequest;
8838
8839 let temp = TempStdDir::default();
8840 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
8841 let table_id = vec!["t".to_string()];
8842 create_managed_table(&ns, &table_id).await; ns.create_table_branch(CreateTableBranchRequest {
8845 id: Some(table_id.clone()),
8846 name: "old".to_string(),
8847 from_version: Some(1),
8848 ..Default::default()
8849 })
8850 .await
8851 .unwrap();
8852
8853 let old_ds = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
8854 .await
8855 .unwrap()
8856 .with_branch("old", None)
8857 .load()
8858 .await
8859 .unwrap();
8860 assert_eq!(old_ds.manifest.branch.as_deref(), Some("old"));
8861 assert_eq!(
8862 scan_id_column(&old_ds).await,
8863 vec![1],
8864 "the fork must contain only the fork-point data"
8865 );
8866 }
8867
8868 #[test]
8871 fn test_manifest_version_from_filename() {
8872 assert_eq!(
8874 DirectoryNamespace::manifest_version_from_filename("5.manifest"),
8875 Some(5)
8876 );
8877 assert_eq!(
8878 DirectoryNamespace::manifest_version_from_filename("0.manifest"),
8879 Some(0)
8880 );
8881 let v2_five = format!("{:020}.manifest", u64::MAX - 5);
8883 assert_eq!(
8884 DirectoryNamespace::manifest_version_from_filename(&v2_five),
8885 Some(5)
8886 );
8887 let v2_zero = format!("{:020}.manifest", u64::MAX);
8888 assert_eq!(
8889 DirectoryNamespace::manifest_version_from_filename(&v2_zero),
8890 Some(0)
8891 );
8892 assert_eq!(
8894 DirectoryNamespace::manifest_version_from_filename("data.lance"),
8895 None
8896 );
8897 assert_eq!(
8898 DirectoryNamespace::manifest_version_from_filename("d5.manifest"),
8899 None
8900 );
8901 }
8902
8903 #[tokio::test]
8904 async fn test_create_table() {
8905 let (namespace, _temp_dir) = create_test_namespace().await;
8906
8907 let schema = create_test_schema();
8909 let ipc_data = create_test_ipc_data(&schema);
8910
8911 let mut request = CreateTableRequest::new();
8912 request.id = Some(vec!["test_table".to_string()]);
8913
8914 let response = namespace
8915 .create_table(request, bytes::Bytes::from(ipc_data))
8916 .await
8917 .unwrap();
8918
8919 assert!(response.location.is_some());
8920 assert!(response.location.unwrap().ends_with("test_table.lance"));
8921 assert_eq!(response.version, Some(1));
8922 }
8923
8924 #[tokio::test]
8925 async fn test_create_table_without_data() {
8926 let (namespace, _temp_dir) = create_test_namespace().await;
8927
8928 let mut request = CreateTableRequest::new();
8929 request.id = Some(vec!["test_table".to_string()]);
8930
8931 let result = namespace.create_table(request, bytes::Bytes::new()).await;
8932 assert!(result.is_err());
8933 assert!(
8934 result
8935 .unwrap_err()
8936 .to_string()
8937 .contains("Arrow IPC stream) is required")
8938 );
8939 }
8940
8941 #[tokio::test]
8942 async fn test_create_table_with_invalid_id() {
8943 let (namespace, _temp_dir) = create_test_namespace().await;
8944
8945 let schema = create_test_schema();
8947 let ipc_data = create_test_ipc_data(&schema);
8948
8949 let mut request = CreateTableRequest::new();
8951 request.id = Some(vec![]);
8952
8953 let result = namespace
8954 .create_table(request, bytes::Bytes::from(ipc_data.clone()))
8955 .await;
8956 assert!(result.is_err());
8957
8958 let mut create_ns_req = CreateNamespaceRequest::new();
8961 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
8962 namespace.create_namespace(create_ns_req).await.unwrap();
8963
8964 let mut request = CreateTableRequest::new();
8966 request.id = Some(vec!["test_namespace".to_string(), "table".to_string()]);
8967
8968 let result = namespace
8969 .create_table(request, bytes::Bytes::from(ipc_data))
8970 .await;
8971 assert!(
8973 result.is_ok(),
8974 "Multi-level table IDs should work with manifest enabled"
8975 );
8976 }
8977
8978 #[tokio::test]
8979 async fn test_list_tables() {
8980 let (namespace, _temp_dir) = create_test_namespace().await;
8981
8982 let mut request = ListTablesRequest::new();
8984 request.id = Some(vec![]);
8985 let response = namespace.list_tables(request).await.unwrap();
8986 assert_eq!(response.tables.len(), 0);
8987
8988 let schema = create_test_schema();
8990 let ipc_data = create_test_ipc_data(&schema);
8991
8992 let mut create_request = CreateTableRequest::new();
8994 create_request.id = Some(vec!["table1".to_string()]);
8995 namespace
8996 .create_table(create_request, bytes::Bytes::from(ipc_data.clone()))
8997 .await
8998 .unwrap();
8999
9000 let mut create_request = CreateTableRequest::new();
9002 create_request.id = Some(vec!["table2".to_string()]);
9003 namespace
9004 .create_table(create_request, bytes::Bytes::from(ipc_data))
9005 .await
9006 .unwrap();
9007
9008 let mut request = ListTablesRequest::new();
9010 request.id = Some(vec![]);
9011 let response = namespace.list_tables(request).await.unwrap();
9012 let tables = response.tables;
9013 assert_eq!(tables.len(), 2);
9014 assert!(tables.contains(&"table1".to_string()));
9015 assert!(tables.contains(&"table2".to_string()));
9016 }
9017
9018 #[tokio::test]
9019 async fn test_list_tables_pagination() {
9020 let (namespace, _temp_dir) = create_test_namespace().await;
9021
9022 let schema = create_test_schema();
9023 let ipc_data = create_test_ipc_data(&schema);
9024
9025 for name in ["alpha", "bravo", "charlie"] {
9026 let mut req = CreateTableRequest::new();
9027 req.id = Some(vec![name.to_string()]);
9028 namespace
9029 .create_table(req, bytes::Bytes::from(ipc_data.clone()))
9030 .await
9031 .unwrap();
9032 }
9033
9034 let first_page = namespace
9036 .list_tables(ListTablesRequest {
9037 id: Some(vec![]),
9038 limit: Some(2),
9039 ..Default::default()
9040 })
9041 .await
9042 .unwrap();
9043
9044 assert_eq!(first_page.tables, vec!["alpha", "bravo"]);
9045 assert_eq!(first_page.page_token.as_deref(), Some("bravo"));
9046
9047 let second_page = namespace
9049 .list_tables(ListTablesRequest {
9050 id: Some(vec![]),
9051 limit: Some(2),
9052 page_token: first_page.page_token.clone(),
9053 ..Default::default()
9054 })
9055 .await
9056 .unwrap();
9057
9058 assert_eq!(second_page.tables, vec!["charlie"]);
9059 assert!(second_page.page_token.is_none());
9060 }
9061
9062 #[tokio::test]
9063 async fn test_list_tables_pagination_limit_zero() {
9064 let (namespace, _temp_dir) = create_test_namespace().await;
9065
9066 let schema = create_test_schema();
9067 let ipc_data = create_test_ipc_data(&schema);
9068
9069 let mut req = CreateTableRequest::new();
9070 req.id = Some(vec!["alpha".to_string()]);
9071 namespace
9072 .create_table(req, bytes::Bytes::from(ipc_data))
9073 .await
9074 .unwrap();
9075
9076 let response = namespace
9077 .list_tables(ListTablesRequest {
9078 id: Some(vec![]),
9079 limit: Some(0),
9080 ..Default::default()
9081 })
9082 .await
9083 .unwrap();
9084
9085 assert!(response.tables.is_empty());
9086 assert!(response.page_token.is_none());
9087 }
9088
9089 #[tokio::test]
9090 async fn test_list_tables_with_namespace_id() {
9091 let (namespace, _temp_dir) = create_test_namespace().await;
9092
9093 let mut create_ns_req = CreateNamespaceRequest::new();
9095 create_ns_req.id = Some(vec!["test_namespace".to_string()]);
9096 namespace.create_namespace(create_ns_req).await.unwrap();
9097
9098 let mut request = ListTablesRequest::new();
9100 request.id = Some(vec!["test_namespace".to_string()]);
9101
9102 let result = namespace.list_tables(request).await;
9103 assert!(
9105 result.is_ok(),
9106 "list_tables should work with child namespace when manifest is enabled"
9107 );
9108 let response = result.unwrap();
9109 assert_eq!(
9110 response.tables.len(),
9111 0,
9112 "Namespace should have no tables yet"
9113 );
9114 }
9115
9116 #[tokio::test]
9117 async fn test_create_scalar_index() {
9118 let (namespace, _temp_dir) = create_test_namespace().await;
9119 create_scalar_table(&namespace, "users").await;
9120
9121 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9122 let dataset = open_dataset(&namespace, "users").await;
9123 let expected_transaction_id = dataset
9124 .read_transaction()
9125 .await
9126 .unwrap()
9127 .map(|transaction| transaction.uuid);
9128 assert_eq!(transaction_id, expected_transaction_id);
9129 let indices = dataset.load_indices().await.unwrap();
9130 assert!(indices.iter().any(|index| index.name == "users_id_idx"));
9131 }
9132
9133 #[tokio::test]
9134 async fn test_create_vector_index() {
9135 use lance_namespace::models::CreateTableIndexRequest;
9136
9137 let (namespace, _temp_dir) = create_test_namespace().await;
9138 create_vector_table(&namespace, "vectors").await;
9139
9140 let mut create_index_request =
9141 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
9142 create_index_request.id = Some(vec!["vectors".to_string()]);
9143 create_index_request.name = Some("vector_idx".to_string());
9144 create_index_request.distance_type = Some("l2".to_string());
9145 let transaction_id = namespace
9146 .create_table_index(create_index_request)
9147 .await
9148 .unwrap()
9149 .transaction_id;
9150
9151 let dataset = open_dataset(&namespace, "vectors").await;
9152 let expected_transaction_id = dataset
9153 .read_transaction()
9154 .await
9155 .unwrap()
9156 .map(|transaction| transaction.uuid);
9157 assert_eq!(transaction_id, expected_transaction_id);
9158 let indices = dataset.load_indices().await.unwrap();
9159 assert!(indices.iter().any(|index| index.name == "vector_idx"));
9160 }
9161
9162 #[tokio::test]
9163 async fn test_list_table_indices() {
9164 use lance_namespace::models::{CreateTableIndexRequest, ListTableIndicesRequest};
9165
9166 let (namespace, _temp_dir) = create_test_namespace().await;
9167 create_scalar_table(&namespace, "users").await;
9168 create_scalar_index(&namespace, "users", "a_idx").await;
9169 create_scalar_index(&namespace, "users", "b_idx").await;
9170 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9171
9172 let response = namespace
9173 .list_table_indices(ListTableIndicesRequest {
9174 id: Some(vec!["users".to_string()]),
9175 ..Default::default()
9176 })
9177 .await
9178 .unwrap();
9179
9180 assert_eq!(response.indexes.len(), 3);
9181 assert_eq!(response.indexes[0].index_name, "a_idx");
9182 assert_eq!(response.indexes[1].index_name, "b_idx");
9183 assert_eq!(response.indexes[2].index_name, "users_id_idx");
9184 assert!(response.page_token.is_none());
9185 let users_id_idx = response
9186 .indexes
9187 .iter()
9188 .find(|index| index.index_name == "users_id_idx")
9189 .unwrap();
9190 assert_eq!(users_id_idx.columns, vec!["id"]);
9191 assert_eq!(users_id_idx.status, "SUCCEEDED");
9192
9193 assert_eq!(users_id_idx.index_type.as_deref(), Some("BTree"));
9195 assert!(
9196 users_id_idx
9197 .type_url
9198 .as_deref()
9199 .is_some_and(|s| !s.is_empty())
9200 );
9201 assert_eq!(users_id_idx.num_indexed_rows, Some(3));
9202 assert_eq!(users_id_idx.num_unindexed_rows, Some(0));
9203 assert_eq!(users_id_idx.num_segments, Some(1));
9204 assert!(users_id_idx.size_bytes.is_some_and(|size| size > 0));
9205 assert!(users_id_idx.created_at.is_some());
9206 assert!(users_id_idx.index_version.is_some());
9207 assert!(users_id_idx.index_details.is_some());
9208
9209 let dataset = open_dataset(&namespace, "users").await;
9210 let expected_transaction_id = dataset
9211 .read_transaction()
9212 .await
9213 .unwrap()
9214 .map(|transaction| transaction.uuid);
9215 assert_eq!(transaction_id, expected_transaction_id);
9216 let indices = dataset.load_indices().await.unwrap();
9217 assert_eq!(
9218 indices
9219 .iter()
9220 .filter(|index| index.name == "users_id_idx")
9221 .count(),
9222 1
9223 );
9224
9225 let first_page = namespace
9226 .list_table_indices(ListTableIndicesRequest {
9227 id: Some(vec!["users".to_string()]),
9228 limit: Some(2),
9229 ..Default::default()
9230 })
9231 .await
9232 .unwrap();
9233
9234 assert_eq!(first_page.indexes.len(), 2);
9235 assert_eq!(first_page.indexes[0].index_name, "a_idx");
9236 assert_eq!(first_page.indexes[1].index_name, "b_idx");
9237 assert_eq!(first_page.page_token.as_deref(), Some("b_idx"));
9238
9239 let second_page = namespace
9240 .list_table_indices(ListTableIndicesRequest {
9241 id: Some(vec!["users".to_string()]),
9242 page_token: first_page.page_token.clone(),
9243 limit: Some(2),
9244 ..Default::default()
9245 })
9246 .await
9247 .unwrap();
9248
9249 assert_eq!(second_page.indexes.len(), 1);
9250 assert_eq!(second_page.indexes[0].index_name, "users_id_idx");
9251 assert!(second_page.page_token.is_none());
9252
9253 create_vector_table(&namespace, "vectors").await;
9255 let mut create_index_request =
9256 CreateTableIndexRequest::new("vector".to_string(), "IVF_FLAT".to_string());
9257 create_index_request.id = Some(vec!["vectors".to_string()]);
9258 create_index_request.name = Some("vector_idx".to_string());
9259 create_index_request.distance_type = Some("l2".to_string());
9260 namespace
9261 .create_table_index(create_index_request)
9262 .await
9263 .unwrap();
9264
9265 let vector_response = namespace
9266 .list_table_indices(ListTableIndicesRequest {
9267 id: Some(vec!["vectors".to_string()]),
9268 ..Default::default()
9269 })
9270 .await
9271 .unwrap();
9272
9273 assert_eq!(vector_response.indexes.len(), 1);
9274 let vector_idx = &vector_response.indexes[0];
9275 assert_eq!(vector_idx.index_name, "vector_idx");
9276 assert_eq!(vector_idx.columns, vec!["vector"]);
9277 assert_eq!(vector_idx.index_type.as_deref(), Some("IVF_FLAT"));
9278 assert!(
9279 vector_idx
9280 .type_url
9281 .as_deref()
9282 .is_some_and(|s| !s.is_empty())
9283 );
9284 assert!(vector_idx.num_indexed_rows.is_some());
9285 assert!(vector_idx.num_unindexed_rows.is_some());
9286 assert_eq!(vector_idx.num_segments, Some(1));
9287 assert!(vector_idx.created_at.is_some());
9288 assert!(vector_idx.index_version.is_some());
9289 assert!(vector_idx.index_details.is_some());
9290 }
9291
9292 #[tokio::test]
9293 async fn test_describe_table_index_stats() {
9294 use lance_namespace::models::DescribeTableIndexStatsRequest;
9295
9296 let (namespace, _temp_dir) = create_test_namespace().await;
9297 create_scalar_table(&namespace, "users").await;
9298 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9299
9300 let response = namespace
9301 .describe_table_index_stats(DescribeTableIndexStatsRequest {
9302 id: Some(vec!["users".to_string()]),
9303 index_name: Some("users_id_idx".to_string()),
9304 ..Default::default()
9305 })
9306 .await
9307 .unwrap();
9308 assert_eq!(response.index_type, Some("BTree".to_string()));
9309 assert_eq!(response.num_indices, Some(1));
9310 assert_eq!(response.num_indexed_rows, Some(3));
9311 assert_eq!(response.num_unindexed_rows, Some(0));
9312
9313 let dataset = open_dataset(&namespace, "users").await;
9314 let expected_transaction_id = dataset
9315 .read_transaction()
9316 .await
9317 .unwrap()
9318 .map(|transaction| transaction.uuid);
9319 assert_eq!(transaction_id, expected_transaction_id);
9320 let stats: serde_json::Value =
9321 serde_json::from_str(&dataset.index_statistics("users_id_idx").await.unwrap()).unwrap();
9322 assert_eq!(stats["index_type"], "BTree");
9323 assert_eq!(stats["num_indices"], 1);
9324 assert_eq!(stats["num_indexed_rows"], 3);
9325 assert_eq!(stats["num_unindexed_rows"], 0);
9326 }
9327
9328 #[tokio::test]
9329 async fn test_describe_transaction() {
9330 use lance_namespace::models::DescribeTransactionRequest;
9331
9332 let (namespace, _temp_dir) = create_test_namespace().await;
9333 create_scalar_table(&namespace, "users").await;
9334 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9335 let dataset = open_dataset(&namespace, "users").await;
9336 let latest_transaction = dataset.read_transaction().await.unwrap();
9337 assert_eq!(
9338 transaction_id,
9339 latest_transaction
9340 .as_ref()
9341 .map(|transaction| transaction.uuid.clone())
9342 );
9343
9344 if let Some(transaction_id) = transaction_id {
9345 let response = namespace
9346 .describe_transaction(DescribeTransactionRequest {
9347 id: Some(vec!["users".to_string(), transaction_id.clone()]),
9348 ..Default::default()
9349 })
9350 .await
9351 .unwrap();
9352 assert_eq!(response.status, "SUCCEEDED");
9353 assert_eq!(
9354 response
9355 .properties
9356 .as_ref()
9357 .and_then(|props| props.get("operation")),
9358 Some(&"CreateIndex".to_string())
9359 );
9360 assert_eq!(
9361 response
9362 .properties
9363 .as_ref()
9364 .and_then(|props| props.get("uuid")),
9365 Some(&transaction_id)
9366 );
9367 } else {
9368 assert!(latest_transaction.is_none());
9369 }
9370 }
9371
9372 #[tokio::test]
9373 async fn test_drop_table_index() {
9374 use lance_namespace::models::{DropTableIndexRequest, ListTableIndicesRequest};
9375
9376 let (namespace, _temp_dir) = create_test_namespace().await;
9377 create_scalar_table(&namespace, "users").await;
9378 let create_transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
9379
9380 let drop_transaction_id = namespace
9381 .drop_table_index(DropTableIndexRequest {
9382 id: Some(vec!["users".to_string()]),
9383 index_name: Some("users_id_idx".to_string()),
9384 ..Default::default()
9385 })
9386 .await
9387 .unwrap()
9388 .transaction_id;
9389
9390 let dataset = open_dataset(&namespace, "users").await;
9391 let previous_dataset = dataset
9392 .checkout_version(dataset.version().version - 1)
9393 .await
9394 .unwrap();
9395 let previous_transaction_id = previous_dataset
9396 .read_transaction()
9397 .await
9398 .unwrap()
9399 .map(|transaction| transaction.uuid);
9400 assert_eq!(create_transaction_id, previous_transaction_id);
9401 let expected_drop_transaction_id = dataset
9402 .read_transaction()
9403 .await
9404 .unwrap()
9405 .map(|transaction| transaction.uuid);
9406 assert_eq!(drop_transaction_id, expected_drop_transaction_id);
9407 let indices = dataset.load_indices().await.unwrap();
9408 assert!(!indices.iter().any(|index| index.name == "users_id_idx"));
9409
9410 let list_response = namespace
9411 .list_table_indices(ListTableIndicesRequest {
9412 id: Some(vec!["users".to_string()]),
9413 ..Default::default()
9414 })
9415 .await
9416 .unwrap();
9417 assert!(list_response.indexes.is_empty());
9418 }
9419
9420 #[tokio::test]
9421 async fn test_describe_table() {
9422 let (namespace, _temp_dir) = create_test_namespace().await;
9423
9424 let schema = create_test_schema();
9426 let ipc_data = create_test_ipc_data(&schema);
9427
9428 let mut create_request = CreateTableRequest::new();
9429 create_request.id = Some(vec!["test_table".to_string()]);
9430 namespace
9431 .create_table(create_request, bytes::Bytes::from(ipc_data))
9432 .await
9433 .unwrap();
9434
9435 let mut request = DescribeTableRequest::new();
9437 request.id = Some(vec!["test_table".to_string()]);
9438 let response = namespace.describe_table(request).await.unwrap();
9439
9440 assert!(response.location.is_some());
9441 assert!(response.location.unwrap().ends_with("test_table.lance"));
9442 }
9443
9444 #[tokio::test]
9445 async fn test_describe_nonexistent_table() {
9446 let (namespace, _temp_dir) = create_test_namespace().await;
9447
9448 let mut request = DescribeTableRequest::new();
9449 request.id = Some(vec!["nonexistent".to_string()]);
9450
9451 let result = namespace.describe_table(request).await;
9452 assert!(result.is_err());
9453 assert!(result.unwrap_err().to_string().contains("Table not found"));
9454 }
9455
9456 #[tokio::test]
9457 async fn test_table_exists() {
9458 let (namespace, _temp_dir) = create_test_namespace().await;
9459
9460 let schema = create_test_schema();
9462 let ipc_data = create_test_ipc_data(&schema);
9463
9464 let mut create_request = CreateTableRequest::new();
9465 create_request.id = Some(vec!["existing_table".to_string()]);
9466 namespace
9467 .create_table(create_request, bytes::Bytes::from(ipc_data))
9468 .await
9469 .unwrap();
9470
9471 let mut request = TableExistsRequest::new();
9473 request.id = Some(vec!["existing_table".to_string()]);
9474 let result = namespace.table_exists(request).await;
9475 assert!(result.is_ok());
9476
9477 let mut request = TableExistsRequest::new();
9479 request.id = Some(vec!["nonexistent".to_string()]);
9480 let result = namespace.table_exists(request).await;
9481 assert!(result.is_err());
9482 assert!(result.unwrap_err().to_string().contains("Table not found"));
9483 }
9484
9485 #[tokio::test]
9486 async fn test_drop_table() {
9487 let (namespace, _temp_dir) = create_test_namespace().await;
9488
9489 let schema = create_test_schema();
9491 let ipc_data = create_test_ipc_data(&schema);
9492
9493 let mut create_request = CreateTableRequest::new();
9494 create_request.id = Some(vec!["table_to_drop".to_string()]);
9495 namespace
9496 .create_table(create_request, bytes::Bytes::from(ipc_data))
9497 .await
9498 .unwrap();
9499
9500 let mut exists_request = TableExistsRequest::new();
9502 exists_request.id = Some(vec!["table_to_drop".to_string()]);
9503 assert!(namespace.table_exists(exists_request.clone()).await.is_ok());
9504
9505 let mut drop_request = DropTableRequest::new();
9507 drop_request.id = Some(vec!["table_to_drop".to_string()]);
9508 let response = namespace.drop_table(drop_request).await.unwrap();
9509 assert!(response.location.is_some());
9510
9511 assert!(namespace.table_exists(exists_request).await.is_err());
9513 }
9514
9515 #[tokio::test]
9516 async fn test_drop_nonexistent_table() {
9517 let (namespace, _temp_dir) = create_test_namespace().await;
9518
9519 let mut request = DropTableRequest::new();
9520 request.id = Some(vec!["nonexistent".to_string()]);
9521
9522 let result = namespace.drop_table(request).await;
9524 let _ = result;
9527 }
9528
9529 #[tokio::test]
9530 async fn test_root_namespace_operations() {
9531 let (namespace, _temp_dir) = create_test_namespace().await;
9532
9533 let mut request = ListNamespacesRequest::new();
9535 request.id = Some(vec![]);
9536 let result = namespace.list_namespaces(request).await;
9537 assert!(result.is_ok());
9538 assert_eq!(result.unwrap().namespaces.len(), 0);
9539
9540 let mut request = DescribeNamespaceRequest::new();
9542 request.id = Some(vec![]);
9543 let result = namespace.describe_namespace(request).await;
9544 assert!(result.is_ok());
9545
9546 let mut request = NamespaceExistsRequest::new();
9548 request.id = Some(vec![]);
9549 let result = namespace.namespace_exists(request).await;
9550 assert!(result.is_ok());
9551
9552 let mut request = CreateNamespaceRequest::new();
9554 request.id = Some(vec![]);
9555 let result = namespace.create_namespace(request).await;
9556 assert!(result.is_err());
9557 assert!(result.unwrap_err().to_string().contains("already exists"));
9558
9559 let mut request = DropNamespaceRequest::new();
9561 request.id = Some(vec![]);
9562 let result = namespace.drop_namespace(request).await;
9563 assert!(result.is_err());
9564 assert!(
9565 result
9566 .unwrap_err()
9567 .to_string()
9568 .contains("cannot be dropped")
9569 );
9570 }
9571
9572 #[tokio::test]
9573 async fn test_non_root_namespace_operations() {
9574 let (namespace, _temp_dir) = create_test_namespace().await;
9575
9576 let mut request = CreateNamespaceRequest::new();
9579 request.id = Some(vec!["child".to_string()]);
9580 let result = namespace.create_namespace(request).await;
9581 assert!(
9582 result.is_ok(),
9583 "Child namespace creation should succeed with manifest enabled"
9584 );
9585
9586 let mut request = NamespaceExistsRequest::new();
9588 request.id = Some(vec!["child".to_string()]);
9589 let result = namespace.namespace_exists(request).await;
9590 assert!(
9591 result.is_ok(),
9592 "Child namespace should exist after creation"
9593 );
9594
9595 let mut request = DropNamespaceRequest::new();
9597 request.id = Some(vec!["child".to_string()]);
9598 let result = namespace.drop_namespace(request).await;
9599 assert!(
9600 result.is_ok(),
9601 "Child namespace drop should succeed with manifest enabled"
9602 );
9603
9604 let mut request = NamespaceExistsRequest::new();
9606 request.id = Some(vec!["child".to_string()]);
9607 let result = namespace.namespace_exists(request).await;
9608 assert!(
9609 result.is_err(),
9610 "Child namespace should not exist after drop"
9611 );
9612 }
9613
9614 #[tokio::test]
9615 async fn test_config_custom_root() {
9616 let temp_dir = TempStdDir::default();
9617 let custom_path = temp_dir.join("custom");
9618 std::fs::create_dir(&custom_path).unwrap();
9619
9620 let namespace = DirectoryNamespaceBuilder::new(custom_path.to_string_lossy().to_string())
9621 .build()
9622 .await
9623 .unwrap();
9624
9625 let schema = create_test_schema();
9627 let ipc_data = create_test_ipc_data(&schema);
9628
9629 let mut request = CreateTableRequest::new();
9631 request.id = Some(vec!["test_table".to_string()]);
9632
9633 let response = namespace
9634 .create_table(request, bytes::Bytes::from(ipc_data))
9635 .await
9636 .unwrap();
9637
9638 assert!(response.location.unwrap().contains("custom"));
9639 }
9640
9641 #[tokio::test]
9642 async fn test_config_storage_options() {
9643 let temp_dir = TempStdDir::default();
9644
9645 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9646 .storage_option("option1", "value1")
9647 .storage_option("option2", "value2")
9648 .build()
9649 .await
9650 .unwrap();
9651
9652 let schema = create_test_schema();
9654 let ipc_data = create_test_ipc_data(&schema);
9655
9656 let mut request = CreateTableRequest::new();
9658 request.id = Some(vec!["test_table".to_string()]);
9659
9660 let response = namespace
9661 .create_table(request, bytes::Bytes::from(ipc_data))
9662 .await
9663 .unwrap();
9664
9665 let storage_options = response.storage_options.unwrap();
9666 assert_eq!(storage_options.get("option1"), Some(&"value1".to_string()));
9667 assert_eq!(storage_options.get("option2"), Some(&"value2".to_string()));
9668 }
9669
9670 #[tokio::test]
9674 async fn test_no_storage_options_without_vendor() {
9675 use lance_namespace::models::DeclareTableRequest;
9676
9677 let temp_dir = TempStdDir::default();
9678
9679 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9681 .manifest_enabled(false)
9682 .storage_option("aws_access_key_id", "AKID")
9683 .storage_option("aws_secret_access_key", "SECRET")
9684 .storage_option("region", "us-east-1")
9685 .build()
9686 .await
9687 .unwrap();
9688
9689 let schema = create_test_schema();
9690 let ipc_data = create_test_ipc_data(&schema);
9691
9692 let mut create_req = CreateTableRequest::new();
9694 create_req.id = Some(vec!["t1".to_string()]);
9695 namespace
9696 .create_table(create_req, bytes::Bytes::from(ipc_data))
9697 .await
9698 .unwrap();
9699
9700 let mut desc_req = DescribeTableRequest::new();
9702 desc_req.id = Some(vec!["t1".to_string()]);
9703 let resp = namespace.describe_table(desc_req).await.unwrap();
9704 assert!(resp.storage_options.is_none());
9705
9706 let mut decl_req = DeclareTableRequest::new();
9708 decl_req.id = Some(vec!["t2".to_string()]);
9709 let resp = namespace.declare_table(decl_req).await.unwrap();
9710 assert!(resp.storage_options.is_none());
9711 }
9712
9713 #[tokio::test]
9715 async fn test_no_storage_options_without_vendor_manifest() {
9716 let temp_dir = TempStdDir::default();
9717
9718 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9719 .storage_option("aws_access_key_id", "AKID")
9720 .storage_option("aws_secret_access_key", "SECRET")
9721 .storage_option("region", "us-east-1")
9722 .build()
9723 .await
9724 .unwrap();
9725
9726 let schema = create_test_schema();
9727 let ipc_data = create_test_ipc_data(&schema);
9728
9729 let mut create_req = CreateTableRequest::new();
9730 create_req.id = Some(vec!["t1".to_string()]);
9731 namespace
9732 .create_table(create_req, bytes::Bytes::from(ipc_data))
9733 .await
9734 .unwrap();
9735
9736 let mut desc_req = DescribeTableRequest::new();
9738 desc_req.id = Some(vec!["t1".to_string()]);
9739 let resp = namespace.describe_table(desc_req).await.unwrap();
9740 assert!(resp.storage_options.is_none());
9741 }
9742
9743 #[tokio::test]
9744 async fn test_from_properties_manifest_enabled() {
9745 let temp_dir = TempStdDir::default();
9746
9747 let mut properties = HashMap::new();
9748 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9749 properties.insert("manifest_enabled".to_string(), "true".to_string());
9750 properties.insert("dir_listing_enabled".to_string(), "false".to_string());
9751
9752 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9753 assert!(builder.manifest_enabled);
9754 assert!(!builder.dir_listing_enabled);
9755
9756 let namespace = builder.build().await.unwrap();
9757
9758 let schema = create_test_schema();
9760 let ipc_data = create_test_ipc_data(&schema);
9761
9762 let mut request = CreateTableRequest::new();
9764 request.id = Some(vec!["test_table".to_string()]);
9765
9766 let response = namespace
9767 .create_table(request, bytes::Bytes::from(ipc_data))
9768 .await
9769 .unwrap();
9770
9771 assert!(response.location.is_some());
9772 }
9773
9774 #[tokio::test]
9775 async fn test_from_properties_dir_listing_enabled() {
9776 let temp_dir = TempStdDir::default();
9777
9778 let mut properties = HashMap::new();
9779 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9780 properties.insert("manifest_enabled".to_string(), "false".to_string());
9781 properties.insert("dir_listing_enabled".to_string(), "true".to_string());
9782
9783 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9784 assert!(!builder.manifest_enabled);
9785 assert!(builder.dir_listing_enabled);
9786
9787 let namespace = builder.build().await.unwrap();
9788
9789 let schema = create_test_schema();
9791 let ipc_data = create_test_ipc_data(&schema);
9792
9793 let mut request = CreateTableRequest::new();
9795 request.id = Some(vec!["test_table".to_string()]);
9796
9797 let response = namespace
9798 .create_table(request, bytes::Bytes::from(ipc_data))
9799 .await
9800 .unwrap();
9801
9802 assert!(response.location.is_some());
9803 }
9804
9805 #[tokio::test]
9806 async fn test_from_properties_defaults() {
9807 let temp_dir = TempStdDir::default();
9808
9809 let mut properties = HashMap::new();
9810 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9811
9812 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9813 assert!(builder.manifest_enabled);
9814 assert!(builder.dir_listing_enabled);
9815 assert!(!builder.inline_optimization_enabled);
9816 }
9817
9818 #[test]
9819 fn test_builder_disables_inline_optimization_by_default() {
9820 let builder = DirectoryNamespaceBuilder::new("memory://");
9821 assert!(!builder.inline_optimization_enabled);
9822 }
9823
9824 #[tokio::test]
9825 async fn test_from_properties_with_storage_options() {
9826 let temp_dir = TempStdDir::default();
9827
9828 let mut properties = HashMap::new();
9829 properties.insert("root".to_string(), temp_dir.to_str().unwrap().to_string());
9830 properties.insert("manifest_enabled".to_string(), "true".to_string());
9831 properties.insert("storage.region".to_string(), "us-west-2".to_string());
9832 properties.insert("storage.bucket".to_string(), "my-bucket".to_string());
9833
9834 let builder = DirectoryNamespaceBuilder::from_properties(properties, None).unwrap();
9835 assert!(builder.manifest_enabled);
9836 assert!(builder.storage_options.is_some());
9837
9838 let storage_options = builder.storage_options.unwrap();
9839 assert_eq!(
9840 storage_options.get("region"),
9841 Some(&"us-west-2".to_string())
9842 );
9843 assert_eq!(
9844 storage_options.get("bucket"),
9845 Some(&"my-bucket".to_string())
9846 );
9847 }
9848
9849 #[tokio::test]
9850 async fn test_various_arrow_types() {
9851 let (namespace, _temp_dir) = create_test_namespace().await;
9852
9853 let fields = vec![
9855 JsonArrowField {
9856 name: "bool_col".to_string(),
9857 r#type: Box::new(JsonArrowDataType::new("bool".to_string())),
9858 nullable: true,
9859 metadata: None,
9860 },
9861 JsonArrowField {
9862 name: "int8_col".to_string(),
9863 r#type: Box::new(JsonArrowDataType::new("int8".to_string())),
9864 nullable: true,
9865 metadata: None,
9866 },
9867 JsonArrowField {
9868 name: "float64_col".to_string(),
9869 r#type: Box::new(JsonArrowDataType::new("float64".to_string())),
9870 nullable: true,
9871 metadata: None,
9872 },
9873 JsonArrowField {
9874 name: "binary_col".to_string(),
9875 r#type: Box::new(JsonArrowDataType::new("binary".to_string())),
9876 nullable: true,
9877 metadata: None,
9878 },
9879 ];
9880
9881 let schema = JsonArrowSchema {
9882 fields,
9883 metadata: None,
9884 };
9885
9886 let ipc_data = create_test_ipc_data(&schema);
9888
9889 let mut request = CreateTableRequest::new();
9890 request.id = Some(vec!["complex_table".to_string()]);
9891
9892 let response = namespace
9893 .create_table(request, bytes::Bytes::from(ipc_data))
9894 .await
9895 .unwrap();
9896
9897 assert!(response.location.is_some());
9898 }
9899
9900 #[tokio::test]
9901 async fn test_connect_dir() {
9902 let temp_dir = TempStdDir::default();
9903
9904 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
9905 .build()
9906 .await
9907 .unwrap();
9908
9909 let mut request = ListTablesRequest::new();
9911 request.id = Some(vec![]);
9912 let response = namespace.list_tables(request).await.unwrap();
9913 assert_eq!(response.tables.len(), 0);
9914 }
9915
9916 #[tokio::test]
9917 async fn test_create_table_with_ipc_data() {
9918 use arrow::array::{Int32Array, StringArray};
9919 use arrow::ipc::writer::StreamWriter;
9920
9921 let (namespace, _temp_dir) = create_test_namespace().await;
9922
9923 let schema = create_test_schema();
9925
9926 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
9928 let arrow_schema = Arc::new(arrow_schema);
9929
9930 let id_array = Int32Array::from(vec![1, 2, 3]);
9932 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
9933 let batch = arrow::record_batch::RecordBatch::try_new(
9934 arrow_schema.clone(),
9935 vec![Arc::new(id_array), Arc::new(name_array)],
9936 )
9937 .unwrap();
9938
9939 let mut buffer = Vec::new();
9941 {
9942 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
9943 writer.write(&batch).unwrap();
9944 writer.finish().unwrap();
9945 }
9946
9947 let mut request = CreateTableRequest::new();
9949 request.id = Some(vec!["test_table_with_data".to_string()]);
9950
9951 let response = namespace
9952 .create_table(request, Bytes::from(buffer))
9953 .await
9954 .unwrap();
9955
9956 assert_eq!(response.version, Some(1));
9957 assert!(
9958 response
9959 .location
9960 .unwrap()
9961 .contains("test_table_with_data.lance")
9962 );
9963
9964 let mut exists_request = TableExistsRequest::new();
9966 exists_request.id = Some(vec!["test_table_with_data".to_string()]);
9967 namespace.table_exists(exists_request).await.unwrap();
9968 }
9969
9970 #[tokio::test]
9971 async fn test_child_namespace_create_and_list() {
9972 let (namespace, _temp_dir) = create_test_namespace().await;
9973
9974 for i in 1..=3 {
9976 let mut create_req = CreateNamespaceRequest::new();
9977 create_req.id = Some(vec![format!("ns{}", i)]);
9978 let result = namespace.create_namespace(create_req).await;
9979 assert!(result.is_ok(), "Failed to create child namespace ns{}", i);
9980 }
9981
9982 let list_req = ListNamespacesRequest {
9984 id: Some(vec![]),
9985 ..Default::default()
9986 };
9987 let result = namespace.list_namespaces(list_req).await;
9988 assert!(result.is_ok());
9989 let namespaces = result.unwrap().namespaces;
9990 assert_eq!(namespaces.len(), 3);
9991 assert!(namespaces.contains(&"ns1".to_string()));
9992 assert!(namespaces.contains(&"ns2".to_string()));
9993 assert!(namespaces.contains(&"ns3".to_string()));
9994 }
9995
9996 #[tokio::test]
9997 async fn test_nested_namespace_hierarchy() {
9998 let (namespace, _temp_dir) = create_test_namespace().await;
9999
10000 let mut create_req = CreateNamespaceRequest::new();
10002 create_req.id = Some(vec!["parent".to_string()]);
10003 namespace.create_namespace(create_req).await.unwrap();
10004
10005 let mut create_req = CreateNamespaceRequest::new();
10007 create_req.id = Some(vec!["parent".to_string(), "child1".to_string()]);
10008 namespace.create_namespace(create_req).await.unwrap();
10009
10010 let mut create_req = CreateNamespaceRequest::new();
10011 create_req.id = Some(vec!["parent".to_string(), "child2".to_string()]);
10012 namespace.create_namespace(create_req).await.unwrap();
10013
10014 let list_req = ListNamespacesRequest {
10016 id: Some(vec!["parent".to_string()]),
10017 ..Default::default()
10018 };
10019 let result = namespace.list_namespaces(list_req).await;
10020 assert!(result.is_ok());
10021 let children = result.unwrap().namespaces;
10022 assert_eq!(children.len(), 2);
10023 assert!(children.contains(&"child1".to_string()));
10024 assert!(children.contains(&"child2".to_string()));
10025
10026 let list_req = ListNamespacesRequest {
10028 id: Some(vec![]),
10029 ..Default::default()
10030 };
10031 let result = namespace.list_namespaces(list_req).await;
10032 assert!(result.is_ok());
10033 let root_namespaces = result.unwrap().namespaces;
10034 assert_eq!(root_namespaces.len(), 1);
10035 assert_eq!(root_namespaces[0], "parent");
10036 }
10037
10038 #[tokio::test]
10039 async fn test_table_in_child_namespace() {
10040 let (namespace, _temp_dir) = create_test_namespace().await;
10041
10042 let mut create_ns_req = CreateNamespaceRequest::new();
10044 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10045 namespace.create_namespace(create_ns_req).await.unwrap();
10046
10047 let schema = create_test_schema();
10049 let ipc_data = create_test_ipc_data(&schema);
10050 let mut create_table_req = CreateTableRequest::new();
10051 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10052 let result = namespace
10053 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10054 .await;
10055 assert!(result.is_ok(), "Failed to create table in child namespace");
10056
10057 let list_req = ListTablesRequest {
10059 id: Some(vec!["test_ns".to_string()]),
10060 ..Default::default()
10061 };
10062 let result = namespace.list_tables(list_req).await;
10063 assert!(result.is_ok());
10064 let tables = result.unwrap().tables;
10065 assert_eq!(tables.len(), 1);
10066 assert_eq!(tables[0], "table1");
10067
10068 let mut exists_req = TableExistsRequest::new();
10070 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10071 let result = namespace.table_exists(exists_req).await;
10072 assert!(result.is_ok());
10073
10074 let mut describe_req = DescribeTableRequest::new();
10076 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10077 let result = namespace.describe_table(describe_req).await;
10078 assert!(result.is_ok());
10079 let response = result.unwrap();
10080 assert!(response.location.is_some());
10081 }
10082
10083 #[tokio::test]
10092 async fn test_read_self_heals_after_manifest_created_by_other_connection() {
10093 let temp_dir = TempStdDir::default();
10094 let root = temp_dir.to_str().unwrap();
10095
10096 let reader = DirectoryNamespaceBuilder::new(root).build().await.unwrap();
10099
10100 let writer = DirectoryNamespaceBuilder::new(root).build().await.unwrap();
10103 let mut create_ns_req = CreateNamespaceRequest::new();
10104 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10105 writer.create_namespace(create_ns_req).await.unwrap();
10106 let ipc_data = create_test_ipc_data(&create_test_schema());
10107 let mut create_table_req = CreateTableRequest::new();
10108 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10109 writer
10110 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10111 .await
10112 .unwrap();
10113
10114 let mut describe_req = DescribeTableRequest::new();
10117 describe_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10118 let resp = reader
10119 .describe_table(describe_req)
10120 .await
10121 .expect("describe_table must resolve a table registered after build");
10122 assert!(resp.location.is_some());
10123
10124 let mut exists_req = TableExistsRequest::new();
10125 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10126 reader
10127 .table_exists(exists_req)
10128 .await
10129 .expect("table_exists must resolve a table registered after build");
10130
10131 let list_req = ListTablesRequest {
10132 id: Some(vec!["test_ns".to_string()]),
10133 ..Default::default()
10134 };
10135 let tables = reader.list_tables(list_req).await.unwrap().tables;
10136 assert_eq!(tables, vec!["table1".to_string()]);
10137 }
10138
10139 #[tokio::test]
10147 async fn test_migration_root_read_self_heals_registered_alias() {
10148 use lance_namespace::models::RegisterTableRequest;
10149
10150 let temp_dir = TempStdDir::default();
10151 let temp_path = temp_dir.to_str().unwrap();
10152
10153 let reader = DirectoryNamespaceBuilder::new(temp_path)
10156 .dir_listing_enabled(true)
10157 .dir_listing_to_manifest_migration_enabled(true)
10158 .build()
10159 .await
10160 .unwrap();
10161
10162 let writer = DirectoryNamespaceBuilder::new(temp_path)
10166 .dir_listing_enabled(true)
10167 .dir_listing_to_manifest_migration_enabled(true)
10168 .build()
10169 .await
10170 .unwrap();
10171 let ipc_data = create_test_ipc_data(&create_test_schema());
10172 let table_uri = format!("{}/external_table.lance", temp_path);
10173 let cursor = Cursor::new(ipc_data);
10174 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
10175 let batches: Vec<_> = stream_reader
10176 .collect::<std::result::Result<Vec<_>, _>>()
10177 .unwrap();
10178 let schema = batches[0].schema();
10179 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
10180 let batch_reader = RecordBatchIterator::new(batch_results, schema);
10181 Dataset::write(Box::new(batch_reader), &table_uri, None)
10182 .await
10183 .unwrap();
10184 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
10185 register_req.id = Some(vec!["registered_table".to_string()]);
10186 writer.register_table(register_req).await.unwrap();
10187
10188 let mut describe_req = DescribeTableRequest::new();
10191 describe_req.id = Some(vec!["registered_table".to_string()]);
10192 reader
10193 .describe_table(describe_req)
10194 .await
10195 .expect("describe_table must resolve a manifest alias registered after build");
10196
10197 let mut exists_req = TableExistsRequest::new();
10198 exists_req.id = Some(vec!["registered_table".to_string()]);
10199 reader
10200 .table_exists(exists_req)
10201 .await
10202 .expect("table_exists must resolve a manifest alias registered after build");
10203
10204 let list_req = ListTablesRequest {
10205 id: Some(vec![]),
10206 ..Default::default()
10207 };
10208 let tables = reader.list_tables(list_req).await.unwrap().tables;
10209 assert!(
10210 tables.contains(&"registered_table".to_string()),
10211 "list_tables must include the manifest alias registered after build, got {:?}",
10212 tables
10213 );
10214 }
10215
10216 #[tokio::test]
10217 async fn test_multiple_tables_in_child_namespace() {
10218 let (namespace, _temp_dir) = create_test_namespace().await;
10219
10220 let mut create_ns_req = CreateNamespaceRequest::new();
10222 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10223 namespace.create_namespace(create_ns_req).await.unwrap();
10224
10225 let schema = create_test_schema();
10227 let ipc_data = create_test_ipc_data(&schema);
10228 for i in 1..=3 {
10229 let mut create_table_req = CreateTableRequest::new();
10230 create_table_req.id = Some(vec!["test_ns".to_string(), format!("table{}", i)]);
10231 namespace
10232 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
10233 .await
10234 .unwrap();
10235 }
10236
10237 let list_req = ListTablesRequest {
10239 id: Some(vec!["test_ns".to_string()]),
10240 ..Default::default()
10241 };
10242 let result = namespace.list_tables(list_req).await;
10243 assert!(result.is_ok());
10244 let tables = result.unwrap().tables;
10245 assert_eq!(tables.len(), 3);
10246 assert!(tables.contains(&"table1".to_string()));
10247 assert!(tables.contains(&"table2".to_string()));
10248 assert!(tables.contains(&"table3".to_string()));
10249 }
10250
10251 #[tokio::test]
10252 async fn test_drop_table_in_child_namespace() {
10253 let (namespace, _temp_dir) = create_test_namespace().await;
10254
10255 let mut create_ns_req = CreateNamespaceRequest::new();
10257 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10258 namespace.create_namespace(create_ns_req).await.unwrap();
10259
10260 let schema = create_test_schema();
10262 let ipc_data = create_test_ipc_data(&schema);
10263 let mut create_table_req = CreateTableRequest::new();
10264 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10265 namespace
10266 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10267 .await
10268 .unwrap();
10269
10270 let mut drop_req = DropTableRequest::new();
10272 drop_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10273 let result = namespace.drop_table(drop_req).await;
10274 assert!(result.is_ok(), "Failed to drop table in child namespace");
10275
10276 let mut exists_req = TableExistsRequest::new();
10278 exists_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10279 let result = namespace.table_exists(exists_req).await;
10280 assert!(result.is_err());
10281 }
10282
10283 #[tokio::test]
10284 async fn test_deeply_nested_namespace() {
10285 let (namespace, _temp_dir) = create_test_namespace().await;
10286
10287 let mut create_req = CreateNamespaceRequest::new();
10289 create_req.id = Some(vec!["level1".to_string()]);
10290 namespace.create_namespace(create_req).await.unwrap();
10291
10292 let mut create_req = CreateNamespaceRequest::new();
10293 create_req.id = Some(vec!["level1".to_string(), "level2".to_string()]);
10294 namespace.create_namespace(create_req).await.unwrap();
10295
10296 let mut create_req = CreateNamespaceRequest::new();
10297 create_req.id = Some(vec![
10298 "level1".to_string(),
10299 "level2".to_string(),
10300 "level3".to_string(),
10301 ]);
10302 namespace.create_namespace(create_req).await.unwrap();
10303
10304 let schema = create_test_schema();
10306 let ipc_data = create_test_ipc_data(&schema);
10307 let mut create_table_req = CreateTableRequest::new();
10308 create_table_req.id = Some(vec![
10309 "level1".to_string(),
10310 "level2".to_string(),
10311 "level3".to_string(),
10312 "table1".to_string(),
10313 ]);
10314 let result = namespace
10315 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10316 .await;
10317 assert!(
10318 result.is_ok(),
10319 "Failed to create table in deeply nested namespace"
10320 );
10321
10322 let mut exists_req = TableExistsRequest::new();
10324 exists_req.id = Some(vec![
10325 "level1".to_string(),
10326 "level2".to_string(),
10327 "level3".to_string(),
10328 "table1".to_string(),
10329 ]);
10330 let result = namespace.table_exists(exists_req).await;
10331 assert!(result.is_ok());
10332 }
10333
10334 #[tokio::test]
10335 async fn test_namespace_with_properties() {
10336 let (namespace, _temp_dir) = create_test_namespace().await;
10337
10338 let mut properties = HashMap::new();
10340 properties.insert("owner".to_string(), "test_user".to_string());
10341 properties.insert("description".to_string(), "Test namespace".to_string());
10342
10343 let mut create_req = CreateNamespaceRequest::new();
10344 create_req.id = Some(vec!["test_ns".to_string()]);
10345 create_req.properties = Some(properties.clone());
10346 namespace.create_namespace(create_req).await.unwrap();
10347
10348 let describe_req = DescribeNamespaceRequest {
10350 id: Some(vec!["test_ns".to_string()]),
10351 ..Default::default()
10352 };
10353 let result = namespace.describe_namespace(describe_req).await;
10354 assert!(result.is_ok());
10355 let response = result.unwrap();
10356 assert!(response.properties.is_some());
10357 let props = response.properties.unwrap();
10358 assert_eq!(props.get("owner"), Some(&"test_user".to_string()));
10359 assert_eq!(
10360 props.get("description"),
10361 Some(&"Test namespace".to_string())
10362 );
10363 }
10364
10365 #[tokio::test]
10366 async fn test_cannot_drop_namespace_with_tables() {
10367 let (namespace, _temp_dir) = create_test_namespace().await;
10368
10369 let mut create_ns_req = CreateNamespaceRequest::new();
10371 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10372 namespace.create_namespace(create_ns_req).await.unwrap();
10373
10374 let schema = create_test_schema();
10376 let ipc_data = create_test_ipc_data(&schema);
10377 let mut create_table_req = CreateTableRequest::new();
10378 create_table_req.id = Some(vec!["test_ns".to_string(), "table1".to_string()]);
10379 namespace
10380 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10381 .await
10382 .unwrap();
10383
10384 let mut drop_req = DropNamespaceRequest::new();
10386 drop_req.id = Some(vec!["test_ns".to_string()]);
10387 let result = namespace.drop_namespace(drop_req).await;
10388 assert!(
10389 result.is_err(),
10390 "Should not be able to drop namespace with tables"
10391 );
10392 }
10393
10394 #[tokio::test]
10395 async fn test_isolation_between_namespaces() {
10396 let (namespace, _temp_dir) = create_test_namespace().await;
10397
10398 let mut create_req = CreateNamespaceRequest::new();
10400 create_req.id = Some(vec!["ns1".to_string()]);
10401 namespace.create_namespace(create_req).await.unwrap();
10402
10403 let mut create_req = CreateNamespaceRequest::new();
10404 create_req.id = Some(vec!["ns2".to_string()]);
10405 namespace.create_namespace(create_req).await.unwrap();
10406
10407 let schema = create_test_schema();
10409 let ipc_data = create_test_ipc_data(&schema);
10410
10411 let mut create_table_req = CreateTableRequest::new();
10412 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10413 namespace
10414 .create_table(create_table_req, bytes::Bytes::from(ipc_data.clone()))
10415 .await
10416 .unwrap();
10417
10418 let mut create_table_req = CreateTableRequest::new();
10419 create_table_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
10420 namespace
10421 .create_table(create_table_req, bytes::Bytes::from(ipc_data))
10422 .await
10423 .unwrap();
10424
10425 let list_req = ListTablesRequest {
10427 id: Some(vec!["ns1".to_string()]),
10428 page_token: None,
10429 limit: None,
10430 ..Default::default()
10431 };
10432 let result = namespace.list_tables(list_req).await.unwrap();
10433 assert_eq!(result.tables.len(), 1);
10434 assert_eq!(result.tables[0], "table1");
10435
10436 let list_req = ListTablesRequest {
10437 id: Some(vec!["ns2".to_string()]),
10438 page_token: None,
10439 limit: None,
10440 ..Default::default()
10441 };
10442 let result = namespace.list_tables(list_req).await.unwrap();
10443 assert_eq!(result.tables.len(), 1);
10444 assert_eq!(result.tables[0], "table1");
10445
10446 let mut drop_req = DropTableRequest::new();
10448 drop_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10449 namespace.drop_table(drop_req).await.unwrap();
10450
10451 let mut exists_req = TableExistsRequest::new();
10453 exists_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
10454 assert!(namespace.table_exists(exists_req).await.is_err());
10455
10456 let mut exists_req = TableExistsRequest::new();
10457 exists_req.id = Some(vec!["ns2".to_string(), "table1".to_string()]);
10458 assert!(namespace.table_exists(exists_req).await.is_ok());
10459 }
10460
10461 #[tokio::test]
10462 async fn test_migrate_directory_tables() {
10463 let temp_dir = TempStdDir::default();
10464 let temp_path = temp_dir.to_str().unwrap();
10465
10466 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
10468 .manifest_enabled(false)
10469 .dir_listing_enabled(true)
10470 .build()
10471 .await
10472 .unwrap();
10473
10474 let schema = create_test_schema();
10476 let ipc_data = create_test_ipc_data(&schema);
10477
10478 for i in 1..=3 {
10479 let mut create_req = CreateTableRequest::new();
10480 create_req.id = Some(vec![format!("table{}", i)]);
10481 dir_only_ns
10482 .create_table(create_req, bytes::Bytes::from(ipc_data.clone()))
10483 .await
10484 .unwrap();
10485 }
10486
10487 drop(dir_only_ns);
10488
10489 let dual_mode_ns = DirectoryNamespaceBuilder::new(temp_path)
10491 .manifest_enabled(true)
10492 .dir_listing_enabled(true)
10493 .build()
10494 .await
10495 .unwrap();
10496
10497 let mut list_req = ListTablesRequest::new();
10499 list_req.id = Some(vec![]);
10500 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
10501 assert_eq!(tables.len(), 3);
10502
10503 let migrated_count = dual_mode_ns.migrate().await.unwrap();
10505 assert_eq!(migrated_count, 3, "Should migrate all 3 tables");
10506
10507 let mut list_req = ListTablesRequest::new();
10509 list_req.id = Some(vec![]);
10510 let tables = dual_mode_ns.list_tables(list_req).await.unwrap().tables;
10511 assert_eq!(tables.len(), 3);
10512
10513 let migrated_count = dual_mode_ns.migrate().await.unwrap();
10515 assert_eq!(
10516 migrated_count, 0,
10517 "Should not migrate already-migrated tables"
10518 );
10519
10520 drop(dual_mode_ns);
10521
10522 let manifest_only_ns = DirectoryNamespaceBuilder::new(temp_path)
10524 .manifest_enabled(true)
10525 .dir_listing_enabled(false)
10526 .build()
10527 .await
10528 .unwrap();
10529
10530 let mut list_req = ListTablesRequest::new();
10532 list_req.id = Some(vec![]);
10533 let tables = manifest_only_ns.list_tables(list_req).await.unwrap().tables;
10534 assert_eq!(tables.len(), 3);
10535 assert!(tables.contains(&"table1".to_string()));
10536 assert!(tables.contains(&"table2".to_string()));
10537 assert!(tables.contains(&"table3".to_string()));
10538 }
10539
10540 #[tokio::test]
10541 async fn test_migrate_without_manifest() {
10542 let temp_dir = TempStdDir::default();
10543 let temp_path = temp_dir.to_str().unwrap();
10544
10545 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10547 .manifest_enabled(false)
10548 .dir_listing_enabled(true)
10549 .build()
10550 .await
10551 .unwrap();
10552
10553 let migrated_count = namespace.migrate().await.unwrap();
10555 assert_eq!(migrated_count, 0);
10556 }
10557
10558 #[tokio::test]
10559 async fn test_register_table() {
10560 use lance_namespace::models::{RegisterTableRequest, TableExistsRequest};
10561
10562 let temp_dir = TempStdDir::default();
10563 let temp_path = temp_dir.to_str().unwrap();
10564
10565 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10566 .dir_listing_to_manifest_migration_enabled(true)
10567 .build()
10568 .await
10569 .unwrap();
10570
10571 let schema = create_test_schema();
10573 let ipc_data = create_test_ipc_data(&schema);
10574
10575 let table_uri = format!("{}/external_table.lance", temp_path);
10576 let cursor = Cursor::new(ipc_data);
10577 let stream_reader = StreamReader::try_new(cursor, None).unwrap();
10578 let batches: Vec<_> = stream_reader
10579 .collect::<std::result::Result<Vec<_>, _>>()
10580 .unwrap();
10581 let schema = batches[0].schema();
10582 let batch_results: Vec<_> = batches.into_iter().map(Ok).collect();
10583 let reader = RecordBatchIterator::new(batch_results, schema);
10584 Dataset::write(Box::new(reader), &table_uri, None)
10585 .await
10586 .unwrap();
10587
10588 let mut register_req = RegisterTableRequest::new("external_table.lance".to_string());
10590 register_req.id = Some(vec!["registered_table".to_string()]);
10591
10592 let response = namespace.register_table(register_req).await.unwrap();
10593 assert_eq!(response.location, Some("external_table.lance".to_string()));
10594
10595 let mut exists_req = TableExistsRequest::new();
10597 exists_req.id = Some(vec!["registered_table".to_string()]);
10598 assert!(namespace.table_exists(exists_req).await.is_ok());
10599
10600 let mut list_req = ListTablesRequest::new();
10602 list_req.id = Some(vec![]);
10603 let tables = namespace.list_tables(list_req).await.unwrap();
10604 assert!(tables.tables.contains(&"registered_table".to_string()));
10605 }
10606
10607 #[tokio::test]
10608 async fn test_register_table_duplicate_fails() {
10609 use lance_namespace::models::RegisterTableRequest;
10610
10611 let temp_dir = TempStdDir::default();
10612 let temp_path = temp_dir.to_str().unwrap();
10613
10614 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10615 .build()
10616 .await
10617 .unwrap();
10618
10619 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10621 register_req.id = Some(vec!["test_table".to_string()]);
10622
10623 namespace
10624 .register_table(register_req.clone())
10625 .await
10626 .unwrap();
10627
10628 let result = namespace.register_table(register_req).await;
10630 assert!(result.is_err());
10631 assert!(result.unwrap_err().to_string().contains("already exists"));
10632 }
10633
10634 #[tokio::test]
10635 async fn test_deregister_table() {
10636 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
10637
10638 let temp_dir = TempStdDir::default();
10639 let temp_path = temp_dir.to_str().unwrap();
10640
10641 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10644 .manifest_enabled(true)
10645 .dir_listing_enabled(false)
10646 .build()
10647 .await
10648 .unwrap();
10649
10650 let schema = create_test_schema();
10652 let ipc_data = create_test_ipc_data(&schema);
10653
10654 let mut create_req = CreateTableRequest::new();
10655 create_req.id = Some(vec!["test_table".to_string()]);
10656 namespace
10657 .create_table(create_req, bytes::Bytes::from(ipc_data))
10658 .await
10659 .unwrap();
10660
10661 let mut exists_req = TableExistsRequest::new();
10663 exists_req.id = Some(vec!["test_table".to_string()]);
10664 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
10665
10666 let mut deregister_req = DeregisterTableRequest::new();
10668 deregister_req.id = Some(vec!["test_table".to_string()]);
10669 let response = namespace.deregister_table(deregister_req).await.unwrap();
10670
10671 assert!(
10673 response.location.is_some(),
10674 "Deregister should return location"
10675 );
10676 let location = response.location.as_ref().unwrap();
10677 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10680 .expect("Failed to convert temp path to URL");
10681 let expected_prefix = expected_url.to_string();
10682 assert!(
10683 location.starts_with(&expected_prefix),
10684 "Location should start with '{}', got: {}",
10685 expected_prefix,
10686 location
10687 );
10688 assert!(
10689 location.contains("test_table"),
10690 "Location should contain table name: {}",
10691 location
10692 );
10693 assert_eq!(response.id, Some(vec!["test_table".to_string()]));
10694
10695 assert!(namespace.table_exists(exists_req).await.is_err());
10697
10698 let dataset = Dataset::open(location).await;
10700 assert!(
10701 dataset.is_ok(),
10702 "Physical table data should still exist at {}",
10703 location
10704 );
10705 }
10706
10707 #[tokio::test]
10708 async fn test_deregister_table_in_child_namespace() {
10709 use lance_namespace::models::{
10710 CreateNamespaceRequest, DeregisterTableRequest, TableExistsRequest,
10711 };
10712
10713 let temp_dir = TempStdDir::default();
10714 let temp_path = temp_dir.to_str().unwrap();
10715
10716 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10717 .build()
10718 .await
10719 .unwrap();
10720
10721 let mut create_ns_req = CreateNamespaceRequest::new();
10723 create_ns_req.id = Some(vec!["test_ns".to_string()]);
10724 namespace.create_namespace(create_ns_req).await.unwrap();
10725
10726 let schema = create_test_schema();
10728 let ipc_data = create_test_ipc_data(&schema);
10729
10730 let mut create_req = CreateTableRequest::new();
10731 create_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10732 namespace
10733 .create_table(create_req, bytes::Bytes::from(ipc_data))
10734 .await
10735 .unwrap();
10736
10737 let mut deregister_req = DeregisterTableRequest::new();
10739 deregister_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10740 let response = namespace.deregister_table(deregister_req).await.unwrap();
10741
10742 assert!(
10744 response.location.is_some(),
10745 "Deregister should return location"
10746 );
10747 let location = response.location.as_ref().unwrap();
10748 let expected_url = lance_io::object_store::uri_to_url(temp_path)
10751 .expect("Failed to convert temp path to URL");
10752 let expected_prefix = expected_url.to_string();
10753 assert!(
10754 location.starts_with(&expected_prefix),
10755 "Location should start with '{}', got: {}",
10756 expected_prefix,
10757 location
10758 );
10759 assert!(
10760 location.contains("test_ns") && location.contains("test_table"),
10761 "Location should contain namespace and table name: {}",
10762 location
10763 );
10764 assert_eq!(
10765 response.id,
10766 Some(vec!["test_ns".to_string(), "test_table".to_string()])
10767 );
10768
10769 let mut exists_req = TableExistsRequest::new();
10771 exists_req.id = Some(vec!["test_ns".to_string(), "test_table".to_string()]);
10772 assert!(namespace.table_exists(exists_req).await.is_err());
10773 }
10774
10775 #[tokio::test]
10776 async fn test_register_without_manifest_fails() {
10777 use lance_namespace::models::RegisterTableRequest;
10778
10779 let temp_dir = TempStdDir::default();
10780 let temp_path = temp_dir.to_str().unwrap();
10781
10782 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10784 .manifest_enabled(false)
10785 .build()
10786 .await
10787 .unwrap();
10788
10789 let mut register_req = RegisterTableRequest::new("test_table.lance".to_string());
10791 register_req.id = Some(vec!["test_table".to_string()]);
10792 let result = namespace.register_table(register_req).await;
10793 assert!(result.is_err());
10794 assert!(
10795 result
10796 .unwrap_err()
10797 .to_string()
10798 .contains("manifest mode is enabled")
10799 );
10800
10801 }
10804
10805 #[tokio::test]
10806 async fn test_register_table_rejects_absolute_uri() {
10807 use lance_namespace::models::RegisterTableRequest;
10808
10809 let temp_dir = TempStdDir::default();
10810 let temp_path = temp_dir.to_str().unwrap();
10811
10812 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10813 .build()
10814 .await
10815 .unwrap();
10816
10817 let mut register_req = RegisterTableRequest::new("s3://bucket/table.lance".to_string());
10819 register_req.id = Some(vec!["test_table".to_string()]);
10820 let result = namespace.register_table(register_req).await;
10821 assert!(result.is_err());
10822 let err_msg = result.unwrap_err().to_string();
10823 assert!(err_msg.contains("Absolute URIs are not allowed"));
10824 }
10825
10826 #[tokio::test]
10827 async fn test_register_table_rejects_absolute_path() {
10828 use lance_namespace::models::RegisterTableRequest;
10829
10830 let temp_dir = TempStdDir::default();
10831 let temp_path = temp_dir.to_str().unwrap();
10832
10833 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10834 .build()
10835 .await
10836 .unwrap();
10837
10838 let mut register_req = RegisterTableRequest::new("/tmp/table.lance".to_string());
10840 register_req.id = Some(vec!["test_table".to_string()]);
10841 let result = namespace.register_table(register_req).await;
10842 assert!(result.is_err());
10843 let err_msg = result.unwrap_err().to_string();
10844 assert!(err_msg.contains("Absolute paths are not allowed"));
10845 }
10846
10847 #[tokio::test]
10848 async fn test_register_table_rejects_path_traversal() {
10849 use lance_namespace::models::RegisterTableRequest;
10850
10851 let temp_dir = TempStdDir::default();
10852 let temp_path = temp_dir.to_str().unwrap();
10853
10854 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10855 .build()
10856 .await
10857 .unwrap();
10858
10859 let mut register_req = RegisterTableRequest::new("../outside/table.lance".to_string());
10861 register_req.id = Some(vec!["test_table".to_string()]);
10862 let result = namespace.register_table(register_req).await;
10863 assert!(result.is_err());
10864 let err_msg = result.unwrap_err().to_string();
10865 assert!(err_msg.contains("Path traversal is not allowed"));
10866 }
10867
10868 #[tokio::test]
10869 async fn test_namespace_write() {
10870 use arrow::array::Int32Array;
10871 use arrow::datatypes::{DataType, Field as ArrowField, Schema as ArrowSchema};
10872 use arrow::record_batch::{RecordBatch, RecordBatchIterator};
10873 use lance::dataset::{Dataset, WriteMode, WriteParams};
10874 use lance_namespace::LanceNamespace;
10875
10876 let (namespace, _temp_dir) = create_test_namespace().await;
10877 let namespace = Arc::new(namespace) as Arc<dyn LanceNamespace>;
10878
10879 let table_id = vec!["test_ns".to_string(), "test_table".to_string()];
10881 let schema = Arc::new(ArrowSchema::new(vec![
10882 ArrowField::new("a", DataType::Int32, false),
10883 ArrowField::new("b", DataType::Int32, false),
10884 ]));
10885
10886 let data1 = RecordBatch::try_new(
10888 schema.clone(),
10889 vec![
10890 Arc::new(Int32Array::from(vec![1, 2, 3])),
10891 Arc::new(Int32Array::from(vec![10, 20, 30])),
10892 ],
10893 )
10894 .unwrap();
10895
10896 let reader1 = RecordBatchIterator::new(vec![data1].into_iter().map(Ok), schema.clone());
10897 let dataset =
10898 Dataset::write_into_namespace(reader1, namespace.clone(), table_id.clone(), None)
10899 .await
10900 .unwrap();
10901
10902 assert_eq!(dataset.count_rows(None).await.unwrap(), 3);
10903 assert_eq!(dataset.version().version, 1);
10904
10905 let data2 = RecordBatch::try_new(
10907 schema.clone(),
10908 vec![
10909 Arc::new(Int32Array::from(vec![4, 5])),
10910 Arc::new(Int32Array::from(vec![40, 50])),
10911 ],
10912 )
10913 .unwrap();
10914
10915 let params_append = WriteParams {
10916 mode: WriteMode::Append,
10917 ..Default::default()
10918 };
10919
10920 let reader2 = RecordBatchIterator::new(vec![data2].into_iter().map(Ok), schema.clone());
10921 let dataset = Dataset::write_into_namespace(
10922 reader2,
10923 namespace.clone(),
10924 table_id.clone(),
10925 Some(params_append),
10926 )
10927 .await
10928 .unwrap();
10929
10930 assert_eq!(dataset.count_rows(None).await.unwrap(), 5);
10931 assert_eq!(dataset.version().version, 2);
10932
10933 let data3 = RecordBatch::try_new(
10935 schema.clone(),
10936 vec![
10937 Arc::new(Int32Array::from(vec![100, 200])),
10938 Arc::new(Int32Array::from(vec![1000, 2000])),
10939 ],
10940 )
10941 .unwrap();
10942
10943 let params_overwrite = WriteParams {
10944 mode: WriteMode::Overwrite,
10945 ..Default::default()
10946 };
10947
10948 let reader3 = RecordBatchIterator::new(vec![data3].into_iter().map(Ok), schema.clone());
10949 let dataset = Dataset::write_into_namespace(
10950 reader3,
10951 namespace.clone(),
10952 table_id.clone(),
10953 Some(params_overwrite),
10954 )
10955 .await
10956 .unwrap();
10957
10958 assert_eq!(dataset.count_rows(None).await.unwrap(), 2);
10959 assert_eq!(dataset.version().version, 3);
10960
10961 let result = dataset.scan().try_into_batch().await.unwrap();
10963 let a_col = result
10964 .column_by_name("a")
10965 .unwrap()
10966 .as_any()
10967 .downcast_ref::<Int32Array>()
10968 .unwrap();
10969 assert_eq!(a_col.values(), &[100, 200]);
10970 }
10971
10972 #[tokio::test]
10977 async fn test_declare_table_v1_mode() {
10978 use lance_namespace::models::{
10979 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
10980 };
10981
10982 let temp_dir = TempStdDir::default();
10983 let temp_path = temp_dir.to_str().unwrap();
10984
10985 let namespace = DirectoryNamespaceBuilder::new(temp_path)
10987 .manifest_enabled(false)
10988 .build()
10989 .await
10990 .unwrap();
10991
10992 let mut declare_req = DeclareTableRequest::new();
10994 declare_req.id = Some(vec!["test_table".to_string()]);
10995 let response = namespace.declare_table(declare_req).await.unwrap();
10996
10997 assert!(response.location.is_some());
10999 let location = response.location.as_ref().unwrap();
11000 assert!(location.ends_with("test_table.lance"));
11001
11002 let mut exists_req = TableExistsRequest::new();
11004 exists_req.id = Some(vec!["test_table".to_string()]);
11005 assert!(namespace.table_exists(exists_req).await.is_ok());
11006
11007 let mut describe_req = DescribeTableRequest::new();
11009 describe_req.id = Some(vec!["test_table".to_string()]);
11010 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11011 assert!(describe_response.location.is_some());
11012 assert!(describe_response.version.is_none()); assert!(describe_response.schema.is_none()); assert_eq!(describe_response.is_only_declared, None);
11015
11016 let mut describe_req = DescribeTableRequest::new();
11017 describe_req.id = Some(vec!["test_table".to_string()]);
11018 describe_req.check_declared = Some(true);
11019 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11020 assert_eq!(describe_response.is_only_declared, Some(true));
11021
11022 let mut list_req = ListTablesRequest::new();
11023 list_req.id = Some(vec![]);
11024 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
11025 assert_eq!(list_response.tables, vec!["test_table".to_string()]);
11026
11027 list_req.include_declared = Some(false);
11028 let list_response = namespace.list_tables(list_req).await.unwrap();
11029 assert!(list_response.tables.is_empty());
11030 }
11031
11032 #[tokio::test]
11033 async fn test_insert_into_declared_table_promotes_it_from_declared_state() {
11034 use lance_namespace::models::{
11035 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest,
11036 };
11037
11038 let temp_dir = TempStdDir::default();
11039 let temp_path = temp_dir.to_str().unwrap();
11040
11041 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11042 .manifest_enabled(false)
11043 .build()
11044 .await
11045 .unwrap();
11046
11047 let mut declare_req = DeclareTableRequest::new();
11048 declare_req.id = Some(vec!["test_table".to_string()]);
11049 namespace.declare_table(declare_req).await.unwrap();
11050
11051 let schema = create_test_schema();
11052 let ipc_data = create_test_ipc_data(&schema);
11053 let mut insert_req = InsertIntoTableRequest::new();
11054 insert_req.id = Some(vec!["test_table".to_string()]);
11055 namespace
11056 .insert_into_table(insert_req, bytes::Bytes::from(ipc_data))
11057 .await
11058 .unwrap();
11059
11060 let mut describe_req = DescribeTableRequest::new();
11061 describe_req.id = Some(vec!["test_table".to_string()]);
11062 describe_req.load_detailed_metadata = Some(true);
11063 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11064
11065 assert_eq!(describe_response.is_only_declared, Some(false));
11066 assert_eq!(describe_response.version, Some(1));
11067 assert!(describe_response.schema.is_some());
11068
11069 let mut list_req = ListTablesRequest::new();
11070 list_req.id = Some(vec![]);
11071 list_req.include_declared = Some(false);
11072 assert_eq!(
11073 namespace.list_tables(list_req).await.unwrap().tables,
11074 vec!["test_table".to_string()]
11075 );
11076 }
11077
11078 #[tokio::test]
11079 async fn test_create_table_after_declare_table_v1_mode_creates_table() {
11080 use lance_namespace::models::{
11081 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11082 };
11083
11084 let temp_dir = TempStdDir::default();
11085 let temp_path = temp_dir.to_str().unwrap();
11086
11087 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11088 .manifest_enabled(false)
11089 .build()
11090 .await
11091 .unwrap();
11092
11093 let mut declare_req = DeclareTableRequest::new();
11094 declare_req.id = Some(vec!["test_table".to_string()]);
11095 namespace.declare_table(declare_req).await.unwrap();
11096
11097 let mut create_req = CreateTableRequest::new();
11098 create_req.id = Some(vec!["test_table".to_string()]);
11099 let response = namespace
11100 .create_table(
11101 create_req,
11102 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11103 )
11104 .await
11105 .unwrap();
11106
11107 assert_eq!(response.version, Some(1));
11108
11109 let mut describe_req = DescribeTableRequest::new();
11110 describe_req.id = Some(vec!["test_table".to_string()]);
11111 describe_req.load_detailed_metadata = Some(true);
11112 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11113 assert_eq!(describe_response.is_only_declared, Some(false));
11114 assert_eq!(describe_response.version, Some(1));
11115
11116 let mut list_req = ListTablesRequest::new();
11117 list_req.id = Some(vec![]);
11118 list_req.include_declared = Some(false);
11119 assert_eq!(
11120 namespace.list_tables(list_req).await.unwrap().tables,
11121 vec!["test_table".to_string()]
11122 );
11123 }
11124
11125 #[tokio::test]
11126 async fn test_insert_into_declared_table_with_manifest_promotes_it() {
11127 use lance_namespace::models::{
11128 DeclareTableRequest, DescribeTableRequest, InsertIntoTableRequest, ListTablesRequest,
11129 };
11130
11131 let temp_dir = TempStdDir::default();
11132 let temp_path = temp_dir.to_str().unwrap();
11133
11134 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11135 .manifest_enabled(true)
11136 .dir_listing_enabled(false)
11137 .build()
11138 .await
11139 .unwrap();
11140
11141 let mut declare_req = DeclareTableRequest::new();
11142 declare_req.id = Some(vec!["test_table".to_string()]);
11143 namespace.declare_table(declare_req).await.unwrap();
11144
11145 let mut insert_req = InsertIntoTableRequest::new();
11146 insert_req.id = Some(vec!["test_table".to_string()]);
11147 namespace
11148 .insert_into_table(
11149 insert_req,
11150 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11151 )
11152 .await
11153 .unwrap();
11154
11155 let mut describe_req = DescribeTableRequest::new();
11156 describe_req.id = Some(vec!["test_table".to_string()]);
11157 describe_req.load_detailed_metadata = Some(true);
11158 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11159 assert_eq!(describe_response.is_only_declared, Some(false));
11160 assert_eq!(describe_response.version, Some(1));
11161
11162 let mut list_req = ListTablesRequest::new();
11163 list_req.id = Some(vec![]);
11164 list_req.include_declared = Some(false);
11165 assert_eq!(
11166 namespace.list_tables(list_req).await.unwrap().tables,
11167 vec!["test_table".to_string()]
11168 );
11169 }
11170
11171 #[tokio::test]
11172 async fn test_create_table_after_declare_table_with_manifest_creates_table() {
11173 use lance_namespace::models::{
11174 CreateTableRequest, DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11175 };
11176
11177 let temp_dir = TempStdDir::default();
11178 let temp_path = temp_dir.to_str().unwrap();
11179
11180 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11181 .manifest_enabled(true)
11182 .dir_listing_enabled(false)
11183 .build()
11184 .await
11185 .unwrap();
11186
11187 let mut declare_req = DeclareTableRequest::new();
11188 declare_req.id = Some(vec!["test_table".to_string()]);
11189 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11190 namespace.declare_table(declare_req).await.unwrap();
11191
11192 let mut create_req = CreateTableRequest::new();
11193 create_req.id = Some(vec!["test_table".to_string()]);
11194 create_req.mode = Some("Overwrite".to_string());
11195 let response = namespace
11196 .create_table(
11197 create_req,
11198 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11199 )
11200 .await
11201 .unwrap();
11202
11203 assert_eq!(response.version, Some(1));
11204 assert_eq!(
11205 response
11206 .properties
11207 .as_ref()
11208 .and_then(|properties| properties.get("owner")),
11209 Some(&"alice".to_string())
11210 );
11211
11212 let mut describe_req = DescribeTableRequest::new();
11213 describe_req.id = Some(vec!["test_table".to_string()]);
11214 describe_req.load_detailed_metadata = Some(true);
11215 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11216 assert_eq!(describe_response.is_only_declared, Some(false));
11217 assert_eq!(describe_response.version, Some(1));
11218 assert_eq!(
11219 describe_response
11220 .properties
11221 .as_ref()
11222 .and_then(|properties| properties.get("owner")),
11223 Some(&"alice".to_string())
11224 );
11225
11226 let mut list_req = ListTablesRequest::new();
11227 list_req.id = Some(vec![]);
11228 list_req.include_declared = Some(false);
11229 assert_eq!(
11230 namespace.list_tables(list_req).await.unwrap().tables,
11231 vec!["test_table".to_string()]
11232 );
11233 }
11234
11235 #[tokio::test]
11236 async fn test_create_table_after_declare_table_with_manifest_rejects_new_properties() {
11237 use lance_namespace::models::{CreateTableRequest, DeclareTableRequest};
11238
11239 let temp_dir = TempStdDir::default();
11240 let temp_path = temp_dir.to_str().unwrap();
11241
11242 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11243 .manifest_enabled(true)
11244 .dir_listing_enabled(false)
11245 .build()
11246 .await
11247 .unwrap();
11248
11249 let mut declare_req = DeclareTableRequest::new();
11250 declare_req.id = Some(vec!["test_table".to_string()]);
11251 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11252 namespace.declare_table(declare_req).await.unwrap();
11253
11254 let mut create_req = CreateTableRequest::new();
11255 create_req.id = Some(vec!["test_table".to_string()]);
11256 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
11257
11258 let result = namespace
11259 .create_table(
11260 create_req,
11261 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11262 )
11263 .await;
11264
11265 assert!(result.is_err());
11266 assert!(
11267 result
11268 .unwrap_err()
11269 .to_string()
11270 .contains("cannot set properties for already declared table")
11271 );
11272 }
11273
11274 #[tokio::test]
11275 async fn test_create_table_with_manifest_exist_ok_keeps_existing_table() {
11276 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
11277
11278 let temp_dir = TempStdDir::default();
11279 let temp_path = temp_dir.to_str().unwrap();
11280
11281 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11282 .manifest_enabled(true)
11283 .dir_listing_enabled(false)
11284 .build()
11285 .await
11286 .unwrap();
11287
11288 let mut create_req = CreateTableRequest::new();
11289 create_req.id = Some(vec!["test_table".to_string()]);
11290 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11291 namespace
11292 .create_table(
11293 create_req,
11294 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11295 )
11296 .await
11297 .unwrap();
11298
11299 let mut create_req = CreateTableRequest::new();
11300 create_req.id = Some(vec!["test_table".to_string()]);
11301 create_req.mode = Some("ExistOk".to_string());
11302 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
11303 let response = namespace
11304 .create_table(
11305 create_req,
11306 bytes::Bytes::from(create_single_row_test_ipc_data()),
11307 )
11308 .await
11309 .unwrap();
11310
11311 assert_eq!(
11312 response
11313 .properties
11314 .as_ref()
11315 .and_then(|properties| properties.get("owner")),
11316 Some(&"alice".to_string())
11317 );
11318 assert_eq!(
11319 open_dataset(&namespace, "test_table")
11320 .await
11321 .count_rows(None)
11322 .await
11323 .unwrap(),
11324 2
11325 );
11326
11327 let mut describe_req = DescribeTableRequest::new();
11328 describe_req.id = Some(vec!["test_table".to_string()]);
11329 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11330 assert_eq!(
11331 describe_response
11332 .properties
11333 .as_ref()
11334 .and_then(|properties| properties.get("owner")),
11335 Some(&"alice".to_string())
11336 );
11337 }
11338
11339 #[tokio::test]
11340 async fn test_create_table_with_manifest_overwrite_replaces_existing_table() {
11341 use lance_namespace::models::{CreateTableRequest, DescribeTableRequest};
11342
11343 let temp_dir = TempStdDir::default();
11344 let temp_path = temp_dir.to_str().unwrap();
11345
11346 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11347 .manifest_enabled(true)
11348 .dir_listing_enabled(false)
11349 .build()
11350 .await
11351 .unwrap();
11352
11353 let mut create_req = CreateTableRequest::new();
11354 create_req.id = Some(vec!["test_table".to_string()]);
11355 create_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11356 namespace
11357 .create_table(
11358 create_req,
11359 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11360 )
11361 .await
11362 .unwrap();
11363
11364 let mut create_req = CreateTableRequest::new();
11365 create_req.id = Some(vec!["test_table".to_string()]);
11366 create_req.mode = Some("overwrite".to_string());
11367 create_req.properties = Some(HashMap::from([("owner".to_string(), "bob".to_string())]));
11368 let response = namespace
11369 .create_table(
11370 create_req,
11371 bytes::Bytes::from(create_single_row_test_ipc_data()),
11372 )
11373 .await
11374 .unwrap();
11375
11376 assert_eq!(response.version, Some(2));
11377 assert_eq!(
11378 response
11379 .properties
11380 .as_ref()
11381 .and_then(|properties| properties.get("owner")),
11382 Some(&"bob".to_string())
11383 );
11384 assert_eq!(
11385 open_dataset(&namespace, "test_table")
11386 .await
11387 .count_rows(None)
11388 .await
11389 .unwrap(),
11390 1
11391 );
11392
11393 let mut describe_req = DescribeTableRequest::new();
11394 describe_req.id = Some(vec!["test_table".to_string()]);
11395 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11396 assert_eq!(
11397 describe_response
11398 .properties
11399 .as_ref()
11400 .and_then(|properties| properties.get("owner")),
11401 Some(&"bob".to_string())
11402 );
11403 }
11404
11405 #[tokio::test]
11406 async fn test_create_table_with_manifest_invalid_mode_rejected() {
11407 use lance_namespace::models::CreateTableRequest;
11408
11409 let temp_dir = TempStdDir::default();
11410 let temp_path = temp_dir.to_str().unwrap();
11411
11412 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11413 .manifest_enabled(true)
11414 .dir_listing_enabled(false)
11415 .build()
11416 .await
11417 .unwrap();
11418
11419 let mut create_req = CreateTableRequest::new();
11420 create_req.id = Some(vec!["test_table".to_string()]);
11421 create_req.mode = Some("append".to_string());
11422 let result = namespace
11423 .create_table(
11424 create_req,
11425 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11426 )
11427 .await;
11428
11429 assert!(result.is_err());
11430 assert!(
11431 result
11432 .unwrap_err()
11433 .to_string()
11434 .contains("Unsupported create_table mode")
11435 );
11436 }
11437
11438 #[tokio::test]
11439 async fn test_merge_insert_into_declared_table_v1_mode_creates_table() {
11440 use lance_namespace::models::{
11441 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11442 MergeInsertIntoTableRequest,
11443 };
11444
11445 let temp_dir = TempStdDir::default();
11446 let temp_path = temp_dir.to_str().unwrap();
11447
11448 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11449 .manifest_enabled(false)
11450 .build()
11451 .await
11452 .unwrap();
11453
11454 let mut declare_req = DeclareTableRequest::new();
11455 declare_req.id = Some(vec!["test_table".to_string()]);
11456 namespace.declare_table(declare_req).await.unwrap();
11457
11458 let mut merge_req = MergeInsertIntoTableRequest::new();
11459 merge_req.id = Some(vec!["test_table".to_string()]);
11460 merge_req.on = Some(vec!["id".to_string()]);
11461 let response = namespace
11462 .merge_insert_into_table(
11463 merge_req,
11464 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11465 )
11466 .await
11467 .unwrap();
11468
11469 assert_eq!(response.num_inserted_rows, Some(2));
11470 assert_eq!(response.num_updated_rows, Some(0));
11471
11472 let mut describe_req = DescribeTableRequest::new();
11473 describe_req.id = Some(vec!["test_table".to_string()]);
11474 describe_req.load_detailed_metadata = Some(true);
11475 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11476 assert_eq!(describe_response.is_only_declared, Some(false));
11477 assert_eq!(describe_response.version, Some(1));
11478
11479 let mut list_req = ListTablesRequest::new();
11480 list_req.id = Some(vec![]);
11481 list_req.include_declared = Some(false);
11482 assert_eq!(
11483 namespace.list_tables(list_req).await.unwrap().tables,
11484 vec!["test_table".to_string()]
11485 );
11486 }
11487
11488 #[tokio::test]
11489 async fn test_merge_insert_into_declared_table_with_manifest_creates_table() {
11490 use lance_namespace::models::{
11491 DeclareTableRequest, DescribeTableRequest, ListTablesRequest,
11492 MergeInsertIntoTableRequest,
11493 };
11494
11495 let temp_dir = TempStdDir::default();
11496 let temp_path = temp_dir.to_str().unwrap();
11497
11498 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11499 .manifest_enabled(true)
11500 .dir_listing_enabled(false)
11501 .build()
11502 .await
11503 .unwrap();
11504
11505 let mut declare_req = DeclareTableRequest::new();
11506 declare_req.id = Some(vec!["test_table".to_string()]);
11507 namespace.declare_table(declare_req).await.unwrap();
11508
11509 let mut merge_req = MergeInsertIntoTableRequest::new();
11510 merge_req.id = Some(vec!["test_table".to_string()]);
11511 merge_req.on = Some(vec!["id".to_string()]);
11512 let response = namespace
11513 .merge_insert_into_table(
11514 merge_req,
11515 bytes::Bytes::from(create_non_empty_test_ipc_data()),
11516 )
11517 .await
11518 .unwrap();
11519
11520 assert_eq!(response.num_inserted_rows, Some(2));
11521 assert_eq!(response.num_updated_rows, Some(0));
11522
11523 let mut describe_req = DescribeTableRequest::new();
11524 describe_req.id = Some(vec!["test_table".to_string()]);
11525 describe_req.load_detailed_metadata = Some(true);
11526 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11527 assert_eq!(describe_response.is_only_declared, Some(false));
11528 assert_eq!(describe_response.version, Some(1));
11529
11530 let mut list_req = ListTablesRequest::new();
11531 list_req.id = Some(vec![]);
11532 list_req.include_declared = Some(false);
11533 assert_eq!(
11534 namespace.list_tables(list_req).await.unwrap().tables,
11535 vec!["test_table".to_string()]
11536 );
11537 }
11538
11539 fn create_composite_key_ipc_data(rows: &[(Option<&str>, i32, &str)]) -> Vec<u8> {
11543 use arrow::array::{Int32Array, StringArray};
11544 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
11545 use arrow::record_batch::RecordBatch;
11546
11547 let schema = Arc::new(ArrowSchema::new(vec![
11548 Field::new("region", DataType::Utf8, true),
11549 Field::new("id", DataType::Int32, false),
11550 Field::new("value", DataType::Utf8, false),
11551 ]));
11552 let batch = RecordBatch::try_new(
11553 schema.clone(),
11554 vec![
11555 Arc::new(StringArray::from_iter(
11556 rows.iter().map(|(region, _, _)| *region),
11557 )),
11558 Arc::new(Int32Array::from_iter_values(
11559 rows.iter().map(|(_, id, _)| *id),
11560 )),
11561 Arc::new(StringArray::from_iter_values(
11562 rows.iter().map(|(_, _, value)| *value),
11563 )),
11564 ],
11565 )
11566 .unwrap();
11567 create_ipc_data_from_batches(schema, vec![batch])
11568 }
11569
11570 async fn read_composite_key_rows(root: &str) -> Vec<(Option<String>, i32, String)> {
11572 use arrow::array::Array;
11573
11574 let dataset = Dataset::open(&format!("{}/test_table.lance", root))
11575 .await
11576 .unwrap();
11577 let batch = dataset.scan().try_into_batch().await.unwrap();
11578 let regions = batch["region"]
11579 .as_any()
11580 .downcast_ref::<arrow::array::StringArray>()
11581 .unwrap();
11582 let ids = batch["id"]
11583 .as_any()
11584 .downcast_ref::<arrow::array::Int32Array>()
11585 .unwrap();
11586 let values = batch["value"]
11587 .as_any()
11588 .downcast_ref::<arrow::array::StringArray>()
11589 .unwrap();
11590
11591 let mut rows: Vec<_> = (0..batch.num_rows())
11592 .map(|row| {
11593 (
11594 regions
11595 .is_valid(row)
11596 .then(|| regions.value(row).to_string()),
11597 ids.value(row),
11598 values.value(row).to_string(),
11599 )
11600 })
11601 .collect();
11602 rows.sort_unstable();
11603 rows
11604 }
11605
11606 #[tokio::test]
11607 async fn test_merge_insert_matches_on_every_column_of_a_composite_key() {
11608 use lance_namespace::models::{DeclareTableRequest, MergeInsertIntoTableRequest};
11609
11610 let temp_dir = TempStdDir::default();
11611 let temp_path = temp_dir.to_str().unwrap();
11612 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11613 .manifest_enabled(false)
11614 .build()
11615 .await
11616 .unwrap();
11617
11618 let mut declare_req = DeclareTableRequest::new();
11619 declare_req.id = Some(vec!["test_table".to_string()]);
11620 namespace.declare_table(declare_req).await.unwrap();
11621
11622 let seed = create_composite_key_ipc_data(&[
11623 (Some("us"), 1, "a"),
11624 (Some("us"), 2, "b"),
11625 (Some("eu"), 1, "c"),
11626 ]);
11627 let mut merge_req = MergeInsertIntoTableRequest::new();
11628 merge_req.id = Some(vec!["test_table".to_string()]);
11629 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11630 namespace
11631 .merge_insert_into_table(merge_req, bytes::Bytes::from(seed))
11632 .await
11633 .unwrap();
11634
11635 let mut merge_req = MergeInsertIntoTableRequest::new();
11638 merge_req.id = Some(vec!["test_table".to_string()]);
11639 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11640 merge_req.when_matched_update_all = Some(true);
11641 let response = namespace
11642 .merge_insert_into_table(
11643 merge_req,
11644 bytes::Bytes::from(create_composite_key_ipc_data(&[
11645 (Some("us"), 1, "updated"),
11646 (Some("eu"), 2, "inserted"),
11647 ])),
11648 )
11649 .await
11650 .unwrap();
11651
11652 assert_eq!(response.num_updated_rows, Some(1));
11653 assert_eq!(response.num_inserted_rows, Some(1));
11654
11655 assert_eq!(
11656 read_composite_key_rows(temp_path).await,
11657 vec![
11658 (Some("eu".into()), 1, "c".into()),
11660 (Some("eu".into()), 2, "inserted".into()),
11661 (Some("us".into()), 1, "updated".into()),
11662 (Some("us".into()), 2, "b".into()),
11663 ]
11664 );
11665 }
11666
11667 #[tokio::test]
11672 async fn test_merge_insert_composite_key_never_matches_a_null_key_column() {
11673 use lance_namespace::models::{DeclareTableRequest, MergeInsertIntoTableRequest};
11674
11675 let temp_dir = TempStdDir::default();
11676 let temp_path = temp_dir.to_str().unwrap();
11677 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11678 .manifest_enabled(false)
11679 .build()
11680 .await
11681 .unwrap();
11682
11683 let mut declare_req = DeclareTableRequest::new();
11684 declare_req.id = Some(vec!["test_table".to_string()]);
11685 namespace.declare_table(declare_req).await.unwrap();
11686
11687 let mut merge_req = MergeInsertIntoTableRequest::new();
11688 merge_req.id = Some(vec!["test_table".to_string()]);
11689 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11690 namespace
11691 .merge_insert_into_table(
11692 merge_req,
11693 bytes::Bytes::from(create_composite_key_ipc_data(&[
11694 (None, 1, "seeded"),
11695 (Some("us"), 1, "us-seeded"),
11696 ])),
11697 )
11698 .await
11699 .unwrap();
11700
11701 let mut merge_req = MergeInsertIntoTableRequest::new();
11702 merge_req.id = Some(vec!["test_table".to_string()]);
11703 merge_req.on = Some(vec!["region".to_string(), "id".to_string()]);
11704 merge_req.when_matched_update_all = Some(true);
11705 let response = namespace
11706 .merge_insert_into_table(
11707 merge_req,
11708 bytes::Bytes::from(create_composite_key_ipc_data(&[(None, 1, "not-a-match")])),
11709 )
11710 .await
11711 .unwrap();
11712
11713 assert_eq!(response.num_updated_rows, Some(0));
11715 assert_eq!(response.num_inserted_rows, Some(1));
11716
11717 assert_eq!(
11718 read_composite_key_rows(temp_path).await,
11719 vec![
11720 (None, 1, "not-a-match".into()),
11721 (None, 1, "seeded".into()),
11722 (Some("us".into()), 1, "us-seeded".into()),
11723 ]
11724 );
11725 }
11726
11727 #[rstest::rstest]
11728 #[case::missing(None, "'on' field is required")]
11729 #[case::empty(Some(vec![]), "must name at least one column")]
11730 #[case::duplicate(
11731 Some(vec!["region".to_string(), "region".to_string()]),
11732 "names column 'region' more than once"
11733 )]
11734 #[tokio::test]
11735 async fn test_merge_insert_rejects_an_invalid_on_key(
11736 #[case] on: Option<Vec<String>>,
11737 #[case] expected_message: &str,
11738 ) {
11739 use lance_namespace::models::{DeclareTableRequest, MergeInsertIntoTableRequest};
11740
11741 let temp_dir = TempStdDir::default();
11742 let temp_path = temp_dir.to_str().unwrap();
11743 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11744 .manifest_enabled(false)
11745 .build()
11746 .await
11747 .unwrap();
11748
11749 let mut declare_req = DeclareTableRequest::new();
11750 declare_req.id = Some(vec!["test_table".to_string()]);
11751 namespace.declare_table(declare_req).await.unwrap();
11752
11753 let mut merge_req = MergeInsertIntoTableRequest::new();
11754 merge_req.id = Some(vec!["test_table".to_string()]);
11755 merge_req.on = on;
11756 let error = namespace
11757 .merge_insert_into_table(
11758 merge_req,
11759 bytes::Bytes::from(create_composite_key_ipc_data(&[(Some("us"), 1, "a")])),
11760 )
11761 .await
11762 .unwrap_err();
11763
11764 let lance_core::Error::Namespace { source, .. } = &error else {
11765 panic!("expected a Namespace error, got: {}", error);
11766 };
11767 let ns_err = source
11768 .downcast_ref::<NamespaceError>()
11769 .expect("expected a NamespaceError source");
11770 assert_eq!(ns_err.code(), lance_namespace::ErrorCode::InvalidInput);
11771 assert!(
11772 error.to_string().contains(expected_message),
11773 "unexpected error message: {error}"
11774 );
11775 }
11776
11777 #[tokio::test]
11778 async fn test_declare_table_with_manifest() {
11779 use lance_namespace::models::{
11780 DeclareTableRequest, DescribeTableRequest, ListTablesRequest, TableExistsRequest,
11781 };
11782
11783 let temp_dir = TempStdDir::default();
11784 let temp_path = temp_dir.to_str().unwrap();
11785
11786 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11788 .manifest_enabled(true)
11789 .dir_listing_enabled(false)
11790 .build()
11791 .await
11792 .unwrap();
11793
11794 let mut declare_req = DeclareTableRequest::new();
11796 declare_req.id = Some(vec!["test_table".to_string()]);
11797 declare_req.properties = Some(HashMap::from([("owner".to_string(), "alice".to_string())]));
11798 let response = namespace.declare_table(declare_req).await.unwrap();
11799
11800 assert!(response.location.is_some());
11802 assert_eq!(
11803 response
11804 .properties
11805 .as_ref()
11806 .and_then(|properties| properties.get("owner")),
11807 Some(&"alice".to_string())
11808 );
11809
11810 let mut exists_req = TableExistsRequest::new();
11812 exists_req.id = Some(vec!["test_table".to_string()]);
11813 assert!(namespace.table_exists(exists_req).await.is_ok());
11814
11815 let mut describe_req = DescribeTableRequest::new();
11816 describe_req.id = Some(vec!["test_table".to_string()]);
11817 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11818 assert_eq!(describe_response.is_only_declared, None);
11819
11820 let mut describe_req = DescribeTableRequest::new();
11821 describe_req.id = Some(vec!["test_table".to_string()]);
11822 describe_req.check_declared = Some(true);
11823 let describe_response = namespace.describe_table(describe_req).await.unwrap();
11824 assert_eq!(describe_response.is_only_declared, Some(true));
11825 assert_eq!(
11826 describe_response
11827 .properties
11828 .as_ref()
11829 .and_then(|properties| properties.get("owner")),
11830 Some(&"alice".to_string())
11831 );
11832
11833 let mut list_req = ListTablesRequest::new();
11834 list_req.id = Some(vec![]);
11835 assert_eq!(
11836 namespace
11837 .list_tables(list_req.clone())
11838 .await
11839 .unwrap()
11840 .tables,
11841 vec!["test_table".to_string()]
11842 );
11843 list_req.include_declared = Some(false);
11844 assert!(
11845 namespace
11846 .list_tables(list_req)
11847 .await
11848 .unwrap()
11849 .tables
11850 .is_empty()
11851 );
11852 }
11853
11854 #[tokio::test]
11855 async fn test_declare_table_with_manifest_marker_already_exists() {
11856 use lance_namespace::error::ErrorCode;
11859 use lance_namespace::models::DeclareTableRequest;
11860
11861 let temp_dir = TempStdDir::default();
11862 let temp_path = temp_dir.to_str().unwrap();
11863
11864 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11865 .manifest_enabled(true)
11866 .dir_listing_enabled(true)
11867 .build()
11868 .await
11869 .unwrap();
11870
11871 let table_dir = temp_dir.join("test_table.lance");
11872 std::fs::create_dir_all(&table_dir).unwrap();
11873 std::fs::write(table_dir.join(".lance-reserved"), b"reserved").unwrap();
11874
11875 let mut declare_req = DeclareTableRequest::new();
11876 declare_req.id = Some(vec!["test_table".to_string()]);
11877 let err = namespace
11878 .declare_table(declare_req)
11879 .await
11880 .expect_err("declare with existing marker must fail");
11881 let msg = err.to_string();
11882 assert!(
11883 msg.contains("already exists") || msg.contains("TableAlreadyExists"),
11884 "expected TableAlreadyExists, got: {msg}"
11885 );
11886 assert_eq!(
11887 mutation_error_code(err),
11888 ErrorCode::TableAlreadyExists,
11889 "expected TableAlreadyExists error code"
11890 );
11891 }
11892
11893 #[tokio::test]
11894 async fn test_declare_table_when_table_exists() {
11895 use lance_namespace::models::DeclareTableRequest;
11896
11897 let temp_dir = TempStdDir::default();
11898 let temp_path = temp_dir.to_str().unwrap();
11899
11900 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11901 .manifest_enabled(false)
11902 .build()
11903 .await
11904 .unwrap();
11905
11906 let schema = create_test_schema();
11908 let ipc_data = create_test_ipc_data(&schema);
11909 let mut create_req = CreateTableRequest::new();
11910 create_req.id = Some(vec!["test_table".to_string()]);
11911 namespace
11912 .create_table(create_req, bytes::Bytes::from(ipc_data))
11913 .await
11914 .unwrap();
11915
11916 let mut declare_req = DeclareTableRequest::new();
11918 declare_req.id = Some(vec!["test_table".to_string()]);
11919 let result = namespace.declare_table(declare_req).await;
11920 assert!(result.is_err());
11921 }
11922
11923 #[tokio::test]
11928 async fn test_deregister_table_v1_mode() {
11929 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
11930
11931 let temp_dir = TempStdDir::default();
11932 let temp_path = temp_dir.to_str().unwrap();
11933
11934 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11936 .manifest_enabled(false)
11937 .dir_listing_enabled(true)
11938 .build()
11939 .await
11940 .unwrap();
11941
11942 let schema = create_test_schema();
11944 let ipc_data = create_test_ipc_data(&schema);
11945 let mut create_req = CreateTableRequest::new();
11946 create_req.id = Some(vec!["test_table".to_string()]);
11947 namespace
11948 .create_table(create_req, bytes::Bytes::from(ipc_data))
11949 .await
11950 .unwrap();
11951
11952 let mut exists_req = TableExistsRequest::new();
11954 exists_req.id = Some(vec!["test_table".to_string()]);
11955 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
11956
11957 let mut deregister_req = DeregisterTableRequest::new();
11959 deregister_req.id = Some(vec!["test_table".to_string()]);
11960 let response = namespace.deregister_table(deregister_req).await.unwrap();
11961
11962 assert!(response.location.is_some());
11964 let location = response.location.as_ref().unwrap();
11965 assert!(location.contains("test_table"));
11966
11967 let result = namespace.table_exists(exists_req).await;
11969 assert!(result.is_err());
11970 assert!(result.unwrap_err().to_string().contains("deregistered"));
11971
11972 let dataset = Dataset::open(location).await;
11974 assert!(dataset.is_ok(), "Physical table data should still exist");
11975 }
11976
11977 #[tokio::test]
11978 async fn test_deregister_table_v1_already_deregistered() {
11979 use lance_namespace::models::DeregisterTableRequest;
11980
11981 let temp_dir = TempStdDir::default();
11982 let temp_path = temp_dir.to_str().unwrap();
11983
11984 let namespace = DirectoryNamespaceBuilder::new(temp_path)
11985 .manifest_enabled(false)
11986 .dir_listing_enabled(true)
11987 .build()
11988 .await
11989 .unwrap();
11990
11991 let schema = create_test_schema();
11993 let ipc_data = create_test_ipc_data(&schema);
11994 let mut create_req = CreateTableRequest::new();
11995 create_req.id = Some(vec!["test_table".to_string()]);
11996 namespace
11997 .create_table(create_req, bytes::Bytes::from(ipc_data))
11998 .await
11999 .unwrap();
12000
12001 let mut deregister_req = DeregisterTableRequest::new();
12003 deregister_req.id = Some(vec!["test_table".to_string()]);
12004 namespace
12005 .deregister_table(deregister_req.clone())
12006 .await
12007 .unwrap();
12008
12009 let result = namespace.deregister_table(deregister_req).await;
12011 assert!(result.is_err());
12012 assert!(
12013 result
12014 .unwrap_err()
12015 .to_string()
12016 .contains("already deregistered")
12017 );
12018 }
12019
12020 #[tokio::test]
12025 async fn test_list_tables_skips_deregistered_v1() {
12026 use lance_namespace::models::DeregisterTableRequest;
12027
12028 let temp_dir = TempStdDir::default();
12029 let temp_path = temp_dir.to_str().unwrap();
12030
12031 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12032 .manifest_enabled(false)
12033 .dir_listing_enabled(true)
12034 .build()
12035 .await
12036 .unwrap();
12037
12038 let schema = create_test_schema();
12040 let ipc_data = create_test_ipc_data(&schema);
12041
12042 let mut create_req1 = CreateTableRequest::new();
12043 create_req1.id = Some(vec!["table1".to_string()]);
12044 namespace
12045 .create_table(create_req1, bytes::Bytes::from(ipc_data.clone()))
12046 .await
12047 .unwrap();
12048
12049 let mut create_req2 = CreateTableRequest::new();
12050 create_req2.id = Some(vec!["table2".to_string()]);
12051 namespace
12052 .create_table(create_req2, bytes::Bytes::from(ipc_data))
12053 .await
12054 .unwrap();
12055
12056 let mut list_req = ListTablesRequest::new();
12058 list_req.id = Some(vec![]);
12059 let list_response = namespace.list_tables(list_req.clone()).await.unwrap();
12060 assert_eq!(list_response.tables.len(), 2);
12061
12062 let mut deregister_req = DeregisterTableRequest::new();
12064 deregister_req.id = Some(vec!["table1".to_string()]);
12065 namespace.deregister_table(deregister_req).await.unwrap();
12066
12067 let list_response = namespace.list_tables(list_req).await.unwrap();
12069 assert_eq!(list_response.tables.len(), 1);
12070 assert!(list_response.tables.contains(&"table2".to_string()));
12071 assert!(!list_response.tables.contains(&"table1".to_string()));
12072 }
12073
12074 #[tokio::test]
12079 async fn test_describe_table_fails_for_deregistered_v1() {
12080 use lance_namespace::models::{DeregisterTableRequest, DescribeTableRequest};
12081
12082 let temp_dir = TempStdDir::default();
12083 let temp_path = temp_dir.to_str().unwrap();
12084
12085 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12086 .manifest_enabled(false)
12087 .dir_listing_enabled(true)
12088 .build()
12089 .await
12090 .unwrap();
12091
12092 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 mut describe_req = DescribeTableRequest::new();
12104 describe_req.id = Some(vec!["test_table".to_string()]);
12105 assert!(namespace.describe_table(describe_req.clone()).await.is_ok());
12106
12107 let mut deregister_req = DeregisterTableRequest::new();
12109 deregister_req.id = Some(vec!["test_table".to_string()]);
12110 namespace.deregister_table(deregister_req).await.unwrap();
12111
12112 let result = namespace.describe_table(describe_req).await;
12114 assert!(result.is_err());
12115 let err = result.unwrap_err();
12116 assert!(matches!(err, Error::Namespace { .. }));
12117 let err_msg = err.to_string();
12118 assert!(err_msg.contains("deregistered"));
12119 assert!(err_msg.contains("table id 'test_table'"));
12120 }
12121
12122 #[tokio::test]
12123 async fn test_table_exists_fails_for_deregistered_v1() {
12124 use lance_namespace::models::{DeregisterTableRequest, TableExistsRequest};
12125
12126 let temp_dir = TempStdDir::default();
12127 let temp_path = temp_dir.to_str().unwrap();
12128
12129 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12130 .manifest_enabled(false)
12131 .dir_listing_enabled(true)
12132 .build()
12133 .await
12134 .unwrap();
12135
12136 let schema = create_test_schema();
12138 let ipc_data = create_test_ipc_data(&schema);
12139 let mut create_req = CreateTableRequest::new();
12140 create_req.id = Some(vec!["test_table".to_string()]);
12141 namespace
12142 .create_table(create_req, bytes::Bytes::from(ipc_data))
12143 .await
12144 .unwrap();
12145
12146 let mut exists_req = TableExistsRequest::new();
12148 exists_req.id = Some(vec!["test_table".to_string()]);
12149 assert!(namespace.table_exists(exists_req.clone()).await.is_ok());
12150
12151 let mut deregister_req = DeregisterTableRequest::new();
12153 deregister_req.id = Some(vec!["test_table".to_string()]);
12154 namespace.deregister_table(deregister_req).await.unwrap();
12155
12156 let result = namespace.table_exists(exists_req).await;
12158 assert!(result.is_err());
12159 let err = result.unwrap_err();
12160 assert!(matches!(err, Error::Namespace { .. }));
12161 let err_msg = err.to_string();
12162 assert!(err_msg.contains("deregistered"));
12163 assert!(err_msg.contains("table id 'test_table'"));
12164 }
12165
12166 #[tokio::test]
12167 async fn test_atomic_table_status_check() {
12168 let temp_dir = TempStdDir::default();
12172 let temp_path = temp_dir.to_str().unwrap();
12173
12174 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12175 .manifest_enabled(false)
12176 .dir_listing_enabled(true)
12177 .build()
12178 .await
12179 .unwrap();
12180
12181 let schema = create_test_schema();
12183 let ipc_data = create_test_ipc_data(&schema);
12184 let mut create_req = CreateTableRequest::new();
12185 create_req.id = Some(vec!["test_table".to_string()]);
12186 namespace
12187 .create_table(create_req, bytes::Bytes::from(ipc_data))
12188 .await
12189 .unwrap();
12190
12191 let status = namespace.check_table_status("test_table").await.unwrap();
12193 assert!(status.exists);
12194 assert!(!status.is_deregistered);
12195 assert!(!status.has_reserved_file);
12196 }
12197
12198 #[tokio::test]
12199 async fn test_table_version_tracking_enabled_managed_versioning() {
12200 use lance_namespace::models::DescribeTableRequest;
12201
12202 let temp_dir = TempStdDir::default();
12203 let temp_path = temp_dir.to_str().unwrap();
12204
12205 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12207 .table_version_tracking_enabled(true)
12208 .build()
12209 .await
12210 .unwrap();
12211
12212 let schema = create_test_schema();
12214 let ipc_data = create_test_ipc_data(&schema);
12215 let mut create_req = CreateTableRequest::new();
12216 create_req.id = Some(vec!["test_table".to_string()]);
12217 namespace
12218 .create_table(create_req, bytes::Bytes::from(ipc_data))
12219 .await
12220 .unwrap();
12221
12222 let mut describe_req = DescribeTableRequest::new();
12224 describe_req.id = Some(vec!["test_table".to_string()]);
12225 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
12226
12227 assert_eq!(
12229 describe_resp.managed_versioning,
12230 Some(true),
12231 "managed_versioning should be true when table_version_tracking_enabled=true"
12232 );
12233 }
12234
12235 #[tokio::test]
12236 async fn test_table_version_tracking_disabled_no_managed_versioning() {
12237 use lance_namespace::models::DescribeTableRequest;
12238
12239 let temp_dir = TempStdDir::default();
12240 let temp_path = temp_dir.to_str().unwrap();
12241
12242 let namespace = DirectoryNamespaceBuilder::new(temp_path)
12244 .table_version_tracking_enabled(false)
12245 .build()
12246 .await
12247 .unwrap();
12248
12249 let schema = create_test_schema();
12251 let ipc_data = create_test_ipc_data(&schema);
12252 let mut create_req = CreateTableRequest::new();
12253 create_req.id = Some(vec!["test_table".to_string()]);
12254 namespace
12255 .create_table(create_req, bytes::Bytes::from(ipc_data))
12256 .await
12257 .unwrap();
12258
12259 let mut describe_req = DescribeTableRequest::new();
12261 describe_req.id = Some(vec!["test_table".to_string()]);
12262 let describe_resp = namespace.describe_table(describe_req).await.unwrap();
12263
12264 assert!(
12266 describe_resp.managed_versioning.is_none(),
12267 "managed_versioning should be None when table_version_tracking_enabled=false, got: {:?}",
12268 describe_resp.managed_versioning
12269 );
12270 }
12271
12272 #[tokio::test]
12273 async fn test_list_table_versions() {
12274 use arrow::array::{Int32Array, RecordBatchIterator};
12275 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12276 use arrow::record_batch::RecordBatch;
12277 use lance::dataset::{Dataset, WriteMode, WriteParams};
12278 use lance_namespace::models::{CreateNamespaceRequest, ListTableVersionsRequest};
12279
12280 let temp_dir = TempStrDir::default();
12281 let temp_path: &str = &temp_dir;
12282
12283 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12284 DirectoryNamespaceBuilder::new(temp_path)
12285 .table_version_tracking_enabled(true)
12286 .build()
12287 .await
12288 .unwrap(),
12289 );
12290
12291 let mut create_ns_req = CreateNamespaceRequest::new();
12293 create_ns_req.id = Some(vec!["workspace".to_string()]);
12294 namespace.create_namespace(create_ns_req).await.unwrap();
12295
12296 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12298 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12299 "id",
12300 DataType::Int32,
12301 false,
12302 )]));
12303 let batch = RecordBatch::try_new(
12304 arrow_schema.clone(),
12305 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12306 )
12307 .unwrap();
12308 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
12309 let write_params = WriteParams {
12310 mode: WriteMode::Create,
12311 ..Default::default()
12312 };
12313 let mut dataset = Dataset::write_into_namespace(
12314 batches,
12315 namespace.clone(),
12316 table_id.clone(),
12317 Some(write_params),
12318 )
12319 .await
12320 .unwrap();
12321
12322 let batch2 = RecordBatch::try_new(
12324 arrow_schema.clone(),
12325 vec![Arc::new(Int32Array::from(vec![100, 200]))],
12326 )
12327 .unwrap();
12328 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
12329 dataset.append(batches, None).await.unwrap();
12330
12331 let batch3 = RecordBatch::try_new(
12333 arrow_schema.clone(),
12334 vec![Arc::new(Int32Array::from(vec![300, 400]))],
12335 )
12336 .unwrap();
12337 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
12338 dataset.append(batches, None).await.unwrap();
12339
12340 let mut list_req = ListTableVersionsRequest::new();
12342 list_req.id = Some(table_id.clone());
12343 let list_resp = namespace.list_table_versions(list_req).await.unwrap();
12344
12345 assert_eq!(
12346 list_resp.versions.len(),
12347 3,
12348 "Should have 3 versions, got: {:?}",
12349 list_resp.versions
12350 );
12351
12352 for expected_version in 1..=3 {
12354 let version = list_resp
12355 .versions
12356 .iter()
12357 .find(|v| v.version == expected_version)
12358 .unwrap_or_else(|| panic!("Expected version {}", expected_version));
12359
12360 assert!(
12361 !version.manifest_path.is_empty(),
12362 "manifest_path should be set for version {}",
12363 expected_version
12364 );
12365 assert!(
12366 version.manifest_path.contains(".manifest"),
12367 "manifest_path should contain .manifest for version {}",
12368 expected_version
12369 );
12370 assert!(
12371 version.manifest_size.is_some(),
12372 "manifest_size should be set for version {}",
12373 expected_version
12374 );
12375 assert!(
12376 version.manifest_size.unwrap() > 0,
12377 "manifest_size should be > 0 for version {}",
12378 expected_version
12379 );
12380 assert!(
12381 version.timestamp_millis.is_some(),
12382 "timestamp_millis should be set for version {}",
12383 expected_version
12384 );
12385 }
12386 }
12387
12388 #[tokio::test]
12389 async fn test_describe_table_version() {
12390 use arrow::array::{Int32Array, RecordBatchIterator};
12391 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12392 use arrow::record_batch::RecordBatch;
12393 use lance::dataset::{Dataset, WriteMode, WriteParams};
12394 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
12395
12396 let temp_dir = TempStrDir::default();
12397 let temp_path: &str = &temp_dir;
12398
12399 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12400 DirectoryNamespaceBuilder::new(temp_path)
12401 .table_version_tracking_enabled(true)
12402 .build()
12403 .await
12404 .unwrap(),
12405 );
12406
12407 let mut create_ns_req = CreateNamespaceRequest::new();
12409 create_ns_req.id = Some(vec!["workspace".to_string()]);
12410 namespace.create_namespace(create_ns_req).await.unwrap();
12411
12412 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12414 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12415 "id",
12416 DataType::Int32,
12417 false,
12418 )]));
12419 let batch = RecordBatch::try_new(
12420 arrow_schema.clone(),
12421 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12422 )
12423 .unwrap();
12424 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12425 let write_params = WriteParams {
12426 mode: WriteMode::Create,
12427 ..Default::default()
12428 };
12429 let mut dataset = Dataset::write_into_namespace(
12430 batches,
12431 namespace.clone(),
12432 table_id.clone(),
12433 Some(write_params),
12434 )
12435 .await
12436 .unwrap();
12437
12438 let batch2 = RecordBatch::try_new(
12440 arrow_schema.clone(),
12441 vec![Arc::new(Int32Array::from(vec![100, 200]))],
12442 )
12443 .unwrap();
12444 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
12445 dataset.append(batches, None).await.unwrap();
12446
12447 let mut describe_req = DescribeTableVersionRequest::new();
12449 describe_req.id = Some(table_id.clone());
12450 describe_req.version = Some(1);
12451 let describe_resp = namespace
12452 .describe_table_version(describe_req)
12453 .await
12454 .unwrap();
12455
12456 let version = &describe_resp.version;
12457 assert_eq!(version.version, 1);
12458 assert!(version.timestamp_millis.is_some());
12459 assert!(
12460 !version.manifest_path.is_empty(),
12461 "manifest_path should be set"
12462 );
12463 assert!(
12464 version.manifest_path.contains(".manifest"),
12465 "manifest_path should contain .manifest"
12466 );
12467 assert!(
12468 version.manifest_size.is_some(),
12469 "manifest_size should be set"
12470 );
12471 assert!(
12472 version.manifest_size.unwrap() > 0,
12473 "manifest_size should be > 0"
12474 );
12475
12476 let mut describe_req = DescribeTableVersionRequest::new();
12478 describe_req.id = Some(table_id.clone());
12479 describe_req.version = Some(2);
12480 let describe_resp = namespace
12481 .describe_table_version(describe_req)
12482 .await
12483 .unwrap();
12484
12485 let version = &describe_resp.version;
12486 assert_eq!(version.version, 2);
12487 assert!(version.timestamp_millis.is_some());
12488 assert!(
12489 !version.manifest_path.is_empty(),
12490 "manifest_path should be set"
12491 );
12492 assert!(
12493 version.manifest_size.is_some(),
12494 "manifest_size should be set"
12495 );
12496 assert!(
12497 version.manifest_size.unwrap() > 0,
12498 "manifest_size should be > 0"
12499 );
12500 }
12501
12502 #[tokio::test]
12503 async fn test_describe_table_version_latest() {
12504 use arrow::array::{Int32Array, RecordBatchIterator};
12505 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
12506 use arrow::record_batch::RecordBatch;
12507 use lance::dataset::{Dataset, WriteMode, WriteParams};
12508 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableVersionRequest};
12509
12510 let temp_dir = TempStrDir::default();
12511 let temp_path: &str = &temp_dir;
12512
12513 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12514 DirectoryNamespaceBuilder::new(temp_path)
12515 .table_version_tracking_enabled(true)
12516 .build()
12517 .await
12518 .unwrap(),
12519 );
12520
12521 let mut create_ns_req = CreateNamespaceRequest::new();
12523 create_ns_req.id = Some(vec!["workspace".to_string()]);
12524 namespace.create_namespace(create_ns_req).await.unwrap();
12525
12526 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
12528 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
12529 "id",
12530 DataType::Int32,
12531 false,
12532 )]));
12533 let batch = RecordBatch::try_new(
12534 arrow_schema.clone(),
12535 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
12536 )
12537 .unwrap();
12538 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
12539 let write_params = WriteParams {
12540 mode: WriteMode::Create,
12541 ..Default::default()
12542 };
12543 let mut dataset = Dataset::write_into_namespace(
12544 batches,
12545 namespace.clone(),
12546 table_id.clone(),
12547 Some(write_params),
12548 )
12549 .await
12550 .unwrap();
12551
12552 let batch2 = RecordBatch::try_new(
12554 arrow_schema.clone(),
12555 vec![Arc::new(Int32Array::from(vec![100, 200]))],
12556 )
12557 .unwrap();
12558 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema.clone());
12559 dataset.append(batches, None).await.unwrap();
12560
12561 let batch3 = RecordBatch::try_new(
12563 arrow_schema.clone(),
12564 vec![Arc::new(Int32Array::from(vec![300, 400]))],
12565 )
12566 .unwrap();
12567 let batches = RecordBatchIterator::new(vec![Ok(batch3)], arrow_schema);
12568 dataset.append(batches, None).await.unwrap();
12569
12570 let mut describe_req = DescribeTableVersionRequest::new();
12572 describe_req.id = Some(table_id.clone());
12573 describe_req.version = None;
12574 let describe_resp = namespace
12575 .describe_table_version(describe_req)
12576 .await
12577 .unwrap();
12578
12579 assert_eq!(describe_resp.version.version, 3);
12581 }
12582
12583 #[tokio::test]
12584 async fn test_create_table_version() {
12585 use futures::TryStreamExt;
12586 use lance::dataset::builder::DatasetBuilder;
12587 use lance_namespace::models::CreateTableVersionRequest;
12588
12589 let temp_dir = TempStrDir::default();
12590 let temp_path: &str = &temp_dir;
12591
12592 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12593 DirectoryNamespaceBuilder::new(temp_path)
12594 .table_version_tracking_enabled(true)
12595 .build()
12596 .await
12597 .unwrap(),
12598 );
12599
12600 let schema = create_test_schema();
12602 let ipc_data = create_test_ipc_data(&schema);
12603 let mut create_req = CreateTableRequest::new();
12604 create_req.id = Some(vec!["test_table".to_string()]);
12605 namespace
12606 .create_table(create_req, bytes::Bytes::from(ipc_data))
12607 .await
12608 .unwrap();
12609
12610 let table_id = vec!["test_table".to_string()];
12612 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12613 .await
12614 .unwrap()
12615 .load()
12616 .await
12617 .unwrap();
12618
12619 let versions_path = dataset.versions_dir();
12621 let manifest_metas: Vec<_> = dataset
12622 .object_store(None)
12623 .await
12624 .unwrap()
12625 .inner
12626 .list(Some(&versions_path))
12627 .try_collect()
12628 .await
12629 .unwrap();
12630
12631 let manifest_meta = manifest_metas
12632 .iter()
12633 .find(|m| {
12634 m.location
12635 .filename()
12636 .map(|f| f.ends_with(".manifest"))
12637 .unwrap_or(false)
12638 })
12639 .expect("No manifest file found");
12640
12641 let manifest_data = dataset
12643 .object_store(None)
12644 .await
12645 .unwrap()
12646 .inner
12647 .get(&manifest_meta.location)
12648 .await
12649 .unwrap()
12650 .bytes()
12651 .await
12652 .unwrap();
12653
12654 let staging_path = dataset.versions_dir().join("staging_manifest");
12656 dataset
12657 .object_store(None)
12658 .await
12659 .unwrap()
12660 .inner
12661 .put(&staging_path, manifest_data.into())
12662 .await
12663 .unwrap();
12664
12665 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12668 create_version_req.id = Some(table_id.clone());
12669 create_version_req.naming_scheme = Some("V2".to_string());
12670
12671 let result = namespace.create_table_version(create_version_req).await;
12672 assert!(
12673 result.is_ok(),
12674 "create_table_version should succeed: {:?}",
12675 result
12676 );
12677
12678 let response = result.unwrap();
12680 let version_info = response
12681 .version
12682 .expect("response should contain version info");
12683 let version_2_path = Path::parse(&version_info.manifest_path).unwrap();
12684 let head_result = dataset
12685 .object_store(None)
12686 .await
12687 .unwrap()
12688 .inner
12689 .head(&version_2_path)
12690 .await;
12691 assert!(
12692 head_result.is_ok(),
12693 "Version 2 manifest should exist at {}",
12694 version_2_path
12695 );
12696
12697 let staging_head_result = dataset
12699 .object_store(None)
12700 .await
12701 .unwrap()
12702 .inner
12703 .head(&staging_path)
12704 .await;
12705 assert!(
12706 staging_head_result.is_err(),
12707 "Staging manifest should have been deleted after create_table_version"
12708 );
12709 }
12710
12711 #[tokio::test]
12712 async fn test_create_table_version_idempotent() {
12713 use futures::TryStreamExt;
12716 use lance::dataset::builder::DatasetBuilder;
12717 use lance_namespace::models::CreateTableVersionRequest;
12718
12719 let temp_dir = TempStrDir::default();
12720 let temp_path: &str = &temp_dir;
12721
12722 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12723 DirectoryNamespaceBuilder::new(temp_path)
12724 .table_version_tracking_enabled(true)
12725 .build()
12726 .await
12727 .unwrap(),
12728 );
12729
12730 let schema = create_test_schema();
12731 let ipc_data = create_test_ipc_data(&schema);
12732 let mut create_req = CreateTableRequest::new();
12733 create_req.id = Some(vec!["test_table".to_string()]);
12734 namespace
12735 .create_table(create_req, bytes::Bytes::from(ipc_data))
12736 .await
12737 .unwrap();
12738
12739 let table_id = vec!["test_table".to_string()];
12740 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12741 .await
12742 .unwrap()
12743 .load()
12744 .await
12745 .unwrap();
12746
12747 let versions_path = dataset.versions_dir();
12748 let manifest_metas: Vec<_> = dataset
12749 .object_store(None)
12750 .await
12751 .unwrap()
12752 .inner
12753 .list(Some(&versions_path))
12754 .try_collect()
12755 .await
12756 .unwrap();
12757
12758 let manifest_meta = manifest_metas
12759 .iter()
12760 .find(|m| {
12761 m.location
12762 .filename()
12763 .map(|f| f.ends_with(".manifest"))
12764 .unwrap_or(false)
12765 })
12766 .expect("No manifest file found");
12767
12768 let manifest_data = dataset
12769 .object_store(None)
12770 .await
12771 .unwrap()
12772 .inner
12773 .get(&manifest_meta.location)
12774 .await
12775 .unwrap()
12776 .bytes()
12777 .await
12778 .unwrap();
12779
12780 let staging_path = dataset.versions_dir().join("staging_manifest");
12781 dataset
12782 .object_store(None)
12783 .await
12784 .unwrap()
12785 .inner
12786 .put(&staging_path, manifest_data.clone().into())
12787 .await
12788 .unwrap();
12789
12790 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12791 create_version_req.id = Some(table_id.clone());
12792 create_version_req.naming_scheme = Some("V2".to_string());
12793 let first = namespace
12794 .create_table_version(create_version_req)
12795 .await
12796 .expect("first create_table_version should succeed");
12797
12798 let retry_staging = dataset.versions_dir().join("staging_manifest_retry");
12800 dataset
12801 .object_store(None)
12802 .await
12803 .unwrap()
12804 .inner
12805 .put(&retry_staging, manifest_data.into())
12806 .await
12807 .unwrap();
12808
12809 let mut retry_req = CreateTableVersionRequest::new(2, retry_staging.to_string());
12810 retry_req.id = Some(table_id.clone());
12811 retry_req.naming_scheme = Some("V2".to_string());
12812 let second = namespace
12813 .create_table_version(retry_req)
12814 .await
12815 .expect("idempotent retry must succeed");
12816
12817 assert_eq!(
12818 first.version.as_ref().map(|v| v.version),
12819 second.version.as_ref().map(|v| v.version)
12820 );
12821 assert_eq!(
12822 first.version.as_ref().map(|v| &v.manifest_path),
12823 second.version.as_ref().map(|v| &v.manifest_path)
12824 );
12825 }
12826
12827 #[tokio::test]
12828 async fn test_create_table_version_conflict() {
12829 use futures::TryStreamExt;
12831 use lance::dataset::builder::DatasetBuilder;
12832 use lance_namespace::models::CreateTableVersionRequest;
12833
12834 let temp_dir = TempStrDir::default();
12835 let temp_path: &str = &temp_dir;
12836
12837 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12838 DirectoryNamespaceBuilder::new(temp_path)
12839 .table_version_tracking_enabled(true)
12840 .build()
12841 .await
12842 .unwrap(),
12843 );
12844
12845 let schema = create_test_schema();
12847 let ipc_data = create_test_ipc_data(&schema);
12848 let mut create_req = CreateTableRequest::new();
12849 create_req.id = Some(vec!["test_table".to_string()]);
12850 namespace
12851 .create_table(create_req, bytes::Bytes::from(ipc_data))
12852 .await
12853 .unwrap();
12854
12855 let table_id = vec!["test_table".to_string()];
12857 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
12858 .await
12859 .unwrap()
12860 .load()
12861 .await
12862 .unwrap();
12863
12864 let versions_path = dataset.versions_dir();
12866 let manifest_metas: Vec<_> = dataset
12867 .object_store(None)
12868 .await
12869 .unwrap()
12870 .inner
12871 .list(Some(&versions_path))
12872 .try_collect()
12873 .await
12874 .unwrap();
12875
12876 let manifest_meta = manifest_metas
12877 .iter()
12878 .find(|m| {
12879 m.location
12880 .filename()
12881 .map(|f| f.ends_with(".manifest"))
12882 .unwrap_or(false)
12883 })
12884 .expect("No manifest file found");
12885
12886 let manifest_data = dataset
12888 .object_store(None)
12889 .await
12890 .unwrap()
12891 .inner
12892 .get(&manifest_meta.location)
12893 .await
12894 .unwrap()
12895 .bytes()
12896 .await
12897 .unwrap();
12898
12899 let staging_path = dataset.versions_dir().join("staging_manifest");
12901 dataset
12902 .object_store(None)
12903 .await
12904 .unwrap()
12905 .inner
12906 .put(&staging_path, manifest_data.into())
12907 .await
12908 .unwrap();
12909
12910 let mut create_version_req = CreateTableVersionRequest::new(2, staging_path.to_string());
12912 create_version_req.id = Some(table_id.clone());
12913 create_version_req.naming_scheme = Some("V2".to_string());
12914 let first_result = namespace.create_table_version(create_version_req).await;
12915 assert!(
12916 first_result.is_ok(),
12917 "First create_table_version for version 2 should succeed: {:?}",
12918 first_result
12919 );
12920
12921 let version_2_path = Path::parse(
12923 &first_result
12924 .unwrap()
12925 .version
12926 .expect("response should contain version info")
12927 .manifest_path,
12928 )
12929 .unwrap();
12930
12931 let conflict_staging = dataset.versions_dir().join("staging_manifest_conflict");
12933 dataset
12934 .object_store(None)
12935 .await
12936 .unwrap()
12937 .inner
12938 .put(
12939 &conflict_staging,
12940 bytes::Bytes::from_static(b"not-a-real-manifest").into(),
12941 )
12942 .await
12943 .unwrap();
12944
12945 let mut create_version_req =
12946 CreateTableVersionRequest::new(2, conflict_staging.to_string());
12947 create_version_req.id = Some(table_id.clone());
12948 create_version_req.naming_scheme = Some("V2".to_string());
12949
12950 let result = namespace.create_table_version(create_version_req).await;
12951 assert!(
12952 result.is_err(),
12953 "create_table_version should fail for existing version with different content"
12954 );
12955 let err = result.unwrap_err().to_string();
12956 assert!(
12957 err.contains("already exists") || err.contains("ConcurrentModification"),
12958 "expected ConcurrentModification, got: {err}"
12959 );
12960
12961 let head_result = dataset
12963 .object_store(None)
12964 .await
12965 .unwrap()
12966 .inner
12967 .head(&version_2_path)
12968 .await;
12969 assert!(
12970 head_result.is_ok(),
12971 "Version 2 manifest should still exist at {}",
12972 version_2_path
12973 );
12974 }
12975
12976 #[tokio::test]
12977 async fn test_create_table_version_cas_rejects_gap() {
12978 use futures::TryStreamExt;
12980 use lance::dataset::builder::DatasetBuilder;
12981 use lance_namespace::models::CreateTableVersionRequest;
12982
12983 let temp_dir = TempStrDir::default();
12984 let temp_path: &str = &temp_dir;
12985
12986 let namespace: Arc<dyn LanceNamespace> = Arc::new(
12987 DirectoryNamespaceBuilder::new(temp_path)
12988 .table_version_tracking_enabled(true)
12989 .build()
12990 .await
12991 .unwrap(),
12992 );
12993
12994 let schema = create_test_schema();
12995 let ipc_data = create_test_ipc_data(&schema);
12996 let mut create_req = CreateTableRequest::new();
12997 create_req.id = Some(vec!["test_table".to_string()]);
12998 namespace
12999 .create_table(create_req, bytes::Bytes::from(ipc_data))
13000 .await
13001 .unwrap();
13002
13003 let table_id = vec!["test_table".to_string()];
13004 let dataset = DatasetBuilder::from_namespace(namespace.clone(), table_id.clone())
13005 .await
13006 .unwrap()
13007 .load()
13008 .await
13009 .unwrap();
13010
13011 let versions_path = dataset.versions_dir();
13012 let manifest_metas: Vec<_> = dataset
13013 .object_store(None)
13014 .await
13015 .unwrap()
13016 .inner
13017 .list(Some(&versions_path))
13018 .try_collect()
13019 .await
13020 .unwrap();
13021 let manifest_meta = manifest_metas
13022 .iter()
13023 .find(|m| {
13024 m.location
13025 .filename()
13026 .map(|f| f.ends_with(".manifest"))
13027 .unwrap_or(false)
13028 })
13029 .expect("No manifest file found");
13030 let manifest_data = dataset
13031 .object_store(None)
13032 .await
13033 .unwrap()
13034 .inner
13035 .get(&manifest_meta.location)
13036 .await
13037 .unwrap()
13038 .bytes()
13039 .await
13040 .unwrap();
13041
13042 let staging_path = dataset.versions_dir().join("staging_gap");
13043 dataset
13044 .object_store(None)
13045 .await
13046 .unwrap()
13047 .inner
13048 .put(&staging_path, manifest_data.into())
13049 .await
13050 .unwrap();
13051
13052 let mut req = CreateTableVersionRequest::new(5, staging_path.to_string());
13054 req.id = Some(table_id);
13055 req.naming_scheme = Some("V2".to_string());
13056 let err = namespace
13057 .create_table_version(req)
13058 .await
13059 .expect_err("gap create must fail CAS");
13060 let msg = err.to_string();
13061 assert!(
13062 msg.contains("CAS") || msg.contains("ConcurrentModification"),
13063 "expected CAS ConcurrentModification, got: {msg}"
13064 );
13065 }
13066
13067 #[tokio::test]
13068 async fn test_create_table_version_branch_cas_requires_parent_version() {
13069 use futures::TryStreamExt;
13072 use lance_namespace::models::CreateTableVersionRequest;
13073
13074 let (namespace, _temp_dir) = create_test_namespace().await;
13075 create_scalar_table(&namespace, "users").await;
13076 let main_uri = open_dataset(&namespace, "users").await.uri().to_string();
13077 append_scalar_version(&main_uri, 10).await; let mut main = open_dataset(&namespace, "users").await;
13080 let fork_version = main.version().version;
13081 assert_eq!(fork_version, 2);
13082 let branch_uri = main
13083 .create_branch("exp", fork_version, None)
13084 .await
13085 .unwrap()
13086 .uri()
13087 .to_string();
13088
13089 let branch_ds = Dataset::open(&branch_uri).await.unwrap();
13090 let versions_dir = branch_ds.versions_dir();
13091 let store = branch_ds.object_store(None).await.unwrap();
13092 let manifests: Vec<_> = store
13093 .inner
13094 .list(Some(&versions_dir))
13095 .try_collect()
13096 .await
13097 .unwrap();
13098 for meta in &manifests {
13099 if meta
13100 .location
13101 .filename()
13102 .is_some_and(|f| f.ends_with(".manifest"))
13103 {
13104 store.inner.delete(&meta.location).await.unwrap();
13105 }
13106 }
13107 let remaining_manifests = store
13110 .inner
13111 .list(Some(&versions_dir))
13112 .try_collect::<Vec<_>>()
13113 .await
13114 .unwrap()
13115 .into_iter()
13116 .filter(|m| {
13117 m.location
13118 .filename()
13119 .is_some_and(|f| f.ends_with(".manifest"))
13120 })
13121 .count();
13122 assert_eq!(
13123 remaining_manifests, 0,
13124 "branch version chain should be empty after deleting manifests"
13125 );
13126
13127 let main_ds = open_dataset(&namespace, "users").await;
13129 let main_versions = main_ds.versions_dir();
13130 let main_store = main_ds.object_store(None).await.unwrap();
13131 let source_meta = main_store
13132 .inner
13133 .list(Some(&main_versions))
13134 .try_collect::<Vec<_>>()
13135 .await
13136 .unwrap()
13137 .into_iter()
13138 .find(|m| {
13139 m.location
13140 .filename()
13141 .is_some_and(|f| f.ends_with(".manifest"))
13142 })
13143 .expect("main should have a manifest");
13144 let source_bytes = main_store
13145 .inner
13146 .get(&source_meta.location)
13147 .await
13148 .unwrap()
13149 .bytes()
13150 .await
13151 .unwrap();
13152
13153 let staging_wrong = versions_dir.clone().join("staging_wrong");
13154 store
13155 .inner
13156 .put(&staging_wrong, source_bytes.clone().into())
13157 .await
13158 .unwrap();
13159 let err = namespace
13160 .create_table_version(CreateTableVersionRequest {
13161 id: Some(vec!["users".to_string()]),
13162 version: 1,
13163 manifest_path: staging_wrong.to_string(),
13164 naming_scheme: Some("V2".to_string()),
13165 branch: Some("exp".to_string()),
13166 ..Default::default()
13167 })
13168 .await
13169 .expect_err("bootstrap at v1 must fail when parent_version is 2");
13170 let msg = err.to_string();
13171 assert!(
13172 msg.contains("CAS") || msg.contains("ConcurrentModification"),
13173 "expected CAS ConcurrentModification, got: {msg}"
13174 );
13175
13176 let staging_ok = versions_dir.join("staging_ok");
13177 store
13178 .inner
13179 .put(&staging_ok, source_bytes.into())
13180 .await
13181 .unwrap();
13182 let resp = namespace
13183 .create_table_version(CreateTableVersionRequest {
13184 id: Some(vec!["users".to_string()]),
13185 version: 2,
13186 manifest_path: staging_ok.to_string(),
13187 naming_scheme: Some("V2".to_string()),
13188 branch: Some("exp".to_string()),
13189 ..Default::default()
13190 })
13191 .await
13192 .expect("bootstrap at parent_version must succeed");
13193 assert_eq!(resp.version.as_ref().map(|v| v.version), Some(2));
13194 }
13195
13196 #[tokio::test]
13197 async fn test_create_table_version_table_not_found() {
13198 use lance_namespace::models::CreateTableVersionRequest;
13199
13200 let temp_dir = TempStdDir::default();
13201 let temp_path = temp_dir.to_str().unwrap();
13202
13203 let namespace = DirectoryNamespaceBuilder::new(temp_path)
13204 .table_version_tracking_enabled(true)
13205 .build()
13206 .await
13207 .unwrap();
13208
13209 let mut create_version_req =
13211 CreateTableVersionRequest::new(1, "/some/staging/path".to_string());
13212 create_version_req.id = Some(vec!["non_existent_table".to_string()]);
13213
13214 let result = namespace.create_table_version(create_version_req).await;
13215 assert!(
13216 result.is_err(),
13217 "create_table_version should fail for non-existent table"
13218 );
13219 let err_msg = result.unwrap_err().to_string();
13220 assert!(
13221 err_msg.contains("Table not found"),
13222 "Error should mention table not found, got: {}",
13223 err_msg
13224 );
13225 }
13226
13227 mod e2e_table_version_tracking {
13229 use super::*;
13230 use std::sync::atomic::{AtomicUsize, Ordering};
13231
13232 struct TrackingNamespace {
13234 inner: DirectoryNamespace,
13235 create_table_version_count: AtomicUsize,
13236 describe_table_version_count: AtomicUsize,
13237 list_table_versions_count: AtomicUsize,
13238 }
13239
13240 impl TrackingNamespace {
13241 fn new(inner: DirectoryNamespace) -> Self {
13242 Self {
13243 inner,
13244 create_table_version_count: AtomicUsize::new(0),
13245 describe_table_version_count: AtomicUsize::new(0),
13246 list_table_versions_count: AtomicUsize::new(0),
13247 }
13248 }
13249
13250 fn create_table_version_calls(&self) -> usize {
13251 self.create_table_version_count.load(Ordering::SeqCst)
13252 }
13253
13254 fn describe_table_version_calls(&self) -> usize {
13255 self.describe_table_version_count.load(Ordering::SeqCst)
13256 }
13257
13258 fn list_table_versions_calls(&self) -> usize {
13259 self.list_table_versions_count.load(Ordering::SeqCst)
13260 }
13261 }
13262
13263 impl std::fmt::Debug for TrackingNamespace {
13264 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
13265 f.debug_struct("TrackingNamespace")
13266 .field(
13267 "create_table_version_calls",
13268 &self.create_table_version_calls(),
13269 )
13270 .finish()
13271 }
13272 }
13273
13274 #[async_trait]
13275 impl LanceNamespace for TrackingNamespace {
13276 async fn create_namespace(
13277 &self,
13278 request: CreateNamespaceRequest,
13279 ) -> Result<CreateNamespaceResponse> {
13280 self.inner.create_namespace(request).await
13281 }
13282
13283 async fn describe_namespace(
13284 &self,
13285 request: DescribeNamespaceRequest,
13286 ) -> Result<DescribeNamespaceResponse> {
13287 self.inner.describe_namespace(request).await
13288 }
13289
13290 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
13291 self.inner.namespace_exists(request).await
13292 }
13293
13294 async fn list_namespaces(
13295 &self,
13296 request: ListNamespacesRequest,
13297 ) -> Result<ListNamespacesResponse> {
13298 self.inner.list_namespaces(request).await
13299 }
13300
13301 async fn drop_namespace(
13302 &self,
13303 request: DropNamespaceRequest,
13304 ) -> Result<DropNamespaceResponse> {
13305 self.inner.drop_namespace(request).await
13306 }
13307
13308 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
13309 self.inner.list_tables(request).await
13310 }
13311
13312 async fn describe_table(
13313 &self,
13314 request: DescribeTableRequest,
13315 ) -> Result<DescribeTableResponse> {
13316 self.inner.describe_table(request).await
13317 }
13318
13319 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
13320 self.inner.table_exists(request).await
13321 }
13322
13323 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
13324 self.inner.drop_table(request).await
13325 }
13326
13327 async fn create_table(
13328 &self,
13329 request: CreateTableRequest,
13330 request_data: Bytes,
13331 ) -> Result<CreateTableResponse> {
13332 self.inner.create_table(request, request_data).await
13333 }
13334
13335 async fn declare_table(
13336 &self,
13337 request: DeclareTableRequest,
13338 ) -> Result<DeclareTableResponse> {
13339 self.inner.declare_table(request).await
13340 }
13341
13342 async fn list_table_versions(
13343 &self,
13344 request: ListTableVersionsRequest,
13345 ) -> Result<ListTableVersionsResponse> {
13346 self.list_table_versions_count
13347 .fetch_add(1, Ordering::SeqCst);
13348 self.inner.list_table_versions(request).await
13349 }
13350
13351 async fn create_table_version(
13352 &self,
13353 request: CreateTableVersionRequest,
13354 ) -> Result<CreateTableVersionResponse> {
13355 self.create_table_version_count
13356 .fetch_add(1, Ordering::SeqCst);
13357 self.inner.create_table_version(request).await
13358 }
13359
13360 async fn describe_table_version(
13361 &self,
13362 request: DescribeTableVersionRequest,
13363 ) -> Result<DescribeTableVersionResponse> {
13364 self.describe_table_version_count
13365 .fetch_add(1, Ordering::SeqCst);
13366 self.inner.describe_table_version(request).await
13367 }
13368
13369 async fn batch_delete_table_versions(
13370 &self,
13371 request: BatchDeleteTableVersionsRequest,
13372 ) -> Result<BatchDeleteTableVersionsResponse> {
13373 self.inner.batch_delete_table_versions(request).await
13374 }
13375
13376 fn namespace_id(&self) -> String {
13377 self.inner.namespace_id()
13378 }
13379 }
13380
13381 #[tokio::test]
13382 async fn test_describe_table_returns_managed_versioning() {
13383 use lance_namespace::models::{CreateNamespaceRequest, DescribeTableRequest};
13384
13385 let temp_dir = TempStdDir::default();
13386 let temp_path = temp_dir.to_str().unwrap();
13387
13388 let ns = DirectoryNamespaceBuilder::new(temp_path)
13390 .table_version_tracking_enabled(true)
13391 .manifest_enabled(true)
13392 .build()
13393 .await
13394 .unwrap();
13395
13396 let mut create_ns_req = CreateNamespaceRequest::new();
13398 create_ns_req.id = Some(vec!["workspace".to_string()]);
13399 ns.create_namespace(create_ns_req).await.unwrap();
13400
13401 let schema = create_test_schema();
13403 let ipc_data = create_test_ipc_data(&schema);
13404 let mut create_req = CreateTableRequest::new();
13405 create_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
13406 ns.create_table(create_req, bytes::Bytes::from(ipc_data))
13407 .await
13408 .unwrap();
13409
13410 let mut describe_req = DescribeTableRequest::new();
13412 describe_req.id = Some(vec!["workspace".to_string(), "test_table".to_string()]);
13413 let describe_resp = ns.describe_table(describe_req).await.unwrap();
13414
13415 assert_eq!(
13417 describe_resp.managed_versioning,
13418 Some(true),
13419 "managed_versioning should be true when table_version_tracking_enabled=true"
13420 );
13421 }
13422
13423 #[tokio::test]
13424 async fn test_external_manifest_store_invokes_namespace_apis() {
13425 use arrow::array::{Int32Array, StringArray};
13426 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
13427 use arrow::record_batch::RecordBatch;
13428 use lance::Dataset;
13429 use lance::dataset::builder::DatasetBuilder;
13430 use lance::dataset::{WriteMode, WriteParams};
13431 use lance_namespace::models::CreateNamespaceRequest;
13432
13433 let temp_dir = TempStdDir::default();
13434 let temp_path = temp_dir.to_str().unwrap();
13435
13436 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
13438 .table_version_tracking_enabled(true)
13439 .manifest_enabled(true)
13440 .build()
13441 .await
13442 .unwrap();
13443
13444 let tracking_ns = Arc::new(TrackingNamespace::new(inner_ns));
13445 let ns: Arc<dyn LanceNamespace> = tracking_ns.clone();
13446
13447 let mut create_ns_req = CreateNamespaceRequest::new();
13449 create_ns_req.id = Some(vec!["workspace".to_string()]);
13450 ns.create_namespace(create_ns_req).await.unwrap();
13451
13452 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
13454
13455 let arrow_schema = Arc::new(ArrowSchema::new(vec![
13457 Field::new("id", DataType::Int32, false),
13458 Field::new("name", DataType::Utf8, true),
13459 ]));
13460 let batch = RecordBatch::try_new(
13461 arrow_schema.clone(),
13462 vec![
13463 Arc::new(Int32Array::from(vec![1, 2, 3])),
13464 Arc::new(StringArray::from(vec!["a", "b", "c"])),
13465 ],
13466 )
13467 .unwrap();
13468
13469 let batches = RecordBatchIterator::new(vec![Ok(batch.clone())], arrow_schema.clone());
13471 let write_params = WriteParams {
13472 mode: WriteMode::Create,
13473 ..Default::default()
13474 };
13475 let mut dataset = Dataset::write_into_namespace(
13476 batches,
13477 ns.clone(),
13478 table_id.clone(),
13479 Some(write_params),
13480 )
13481 .await
13482 .unwrap();
13483 assert_eq!(dataset.version().version, 1);
13484
13485 assert_eq!(
13487 tracking_ns.create_table_version_calls(),
13488 1,
13489 "create_table_version should have been called once during initial write_into_namespace"
13490 );
13491
13492 let append_batch = RecordBatch::try_new(
13494 arrow_schema.clone(),
13495 vec![
13496 Arc::new(Int32Array::from(vec![4, 5, 6])),
13497 Arc::new(StringArray::from(vec!["d", "e", "f"])),
13498 ],
13499 )
13500 .unwrap();
13501 let append_batches = RecordBatchIterator::new(vec![Ok(append_batch)], arrow_schema);
13502 dataset.append(append_batches, None).await.unwrap();
13503
13504 assert_eq!(
13505 tracking_ns.create_table_version_calls(),
13506 2,
13507 "create_table_version should have been called twice (once for create, once for append)"
13508 );
13509
13510 let initial_list_calls = tracking_ns.list_table_versions_calls();
13512 let latest_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
13513 .await
13514 .unwrap()
13515 .load()
13516 .await
13517 .unwrap();
13518 assert_eq!(latest_dataset.version().version, 2);
13519 assert_eq!(
13520 tracking_ns.list_table_versions_calls(),
13521 initial_list_calls + 1,
13522 "list_table_versions should have been called exactly once during checkout_latest"
13523 );
13524
13525 let initial_describe_calls = tracking_ns.describe_table_version_calls();
13527 let v1_dataset = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
13528 .await
13529 .unwrap()
13530 .with_version(1)
13531 .load()
13532 .await
13533 .unwrap();
13534 assert_eq!(v1_dataset.version().version, 1);
13535 assert_eq!(
13536 tracking_ns.describe_table_version_calls(),
13537 initial_describe_calls + 1,
13538 "describe_table_version should have been called exactly once during checkout to version 1"
13539 );
13540 }
13541
13542 #[tokio::test]
13543 async fn test_dataset_commit_with_external_manifest_store() {
13544 use arrow::array::{Int32Array, StringArray};
13545 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
13546 use arrow::record_batch::RecordBatch;
13547 use futures::TryStreamExt;
13548 use lance::dataset::{Dataset, WriteMode, WriteParams};
13549 use lance_namespace::models::CreateNamespaceRequest;
13550 use lance_table::io::commit::ManifestNamingScheme;
13551
13552 let temp_dir = TempStdDir::default();
13553 let temp_path = temp_dir.to_str().unwrap();
13554
13555 let inner_ns = DirectoryNamespaceBuilder::new(temp_path)
13557 .table_version_tracking_enabled(true)
13558 .manifest_enabled(true)
13559 .build()
13560 .await
13561 .unwrap();
13562
13563 let tracking_ns: Arc<dyn LanceNamespace> = Arc::new(TrackingNamespace::new(inner_ns));
13564
13565 let mut create_ns_req = CreateNamespaceRequest::new();
13567 create_ns_req.id = Some(vec!["workspace".to_string()]);
13568 tracking_ns.create_namespace(create_ns_req).await.unwrap();
13569
13570 let table_id = vec!["workspace".to_string(), "test_table".to_string()];
13572 let arrow_schema = Arc::new(ArrowSchema::new(vec![
13573 Field::new("id", DataType::Int32, false),
13574 Field::new("name", DataType::Utf8, true),
13575 ]));
13576 let batch = RecordBatch::try_new(
13577 arrow_schema.clone(),
13578 vec![
13579 Arc::new(Int32Array::from(vec![1, 2, 3])),
13580 Arc::new(StringArray::from(vec!["a", "b", "c"])),
13581 ],
13582 )
13583 .unwrap();
13584 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
13585 let write_params = WriteParams {
13586 mode: WriteMode::Create,
13587 ..Default::default()
13588 };
13589 let dataset = Dataset::write_into_namespace(
13590 batches,
13591 tracking_ns.clone(),
13592 table_id.clone(),
13593 Some(write_params),
13594 )
13595 .await
13596 .unwrap();
13597 assert_eq!(dataset.version().version, 1);
13598
13599 let batch2 = RecordBatch::try_new(
13601 arrow_schema.clone(),
13602 vec![
13603 Arc::new(Int32Array::from(vec![4, 5, 6])),
13604 Arc::new(StringArray::from(vec!["d", "e", "f"])),
13605 ],
13606 )
13607 .unwrap();
13608 let batches = RecordBatchIterator::new(vec![Ok(batch2)], arrow_schema);
13609 let write_params = WriteParams {
13610 mode: WriteMode::Append,
13611 ..Default::default()
13612 };
13613 Dataset::write_into_namespace(
13614 batches,
13615 tracking_ns.clone(),
13616 table_id.clone(),
13617 Some(write_params),
13618 )
13619 .await
13620 .unwrap();
13621
13622 let manifest_metas: Vec<_> = dataset
13625 .object_store(None)
13626 .await
13627 .unwrap()
13628 .inner
13629 .list(Some(&dataset.versions_dir()))
13630 .try_collect()
13631 .await
13632 .unwrap();
13633 let version_2_found = manifest_metas.iter().any(|m| {
13634 m.location
13635 .filename()
13636 .map(|f| {
13637 f.ends_with(".manifest")
13638 && ManifestNamingScheme::V2.parse_version(f) == Some(2)
13639 })
13640 .unwrap_or(false)
13641 });
13642 assert!(
13643 version_2_found,
13644 "Version 2 manifest should exist in versions directory"
13645 );
13646 }
13647
13648 async fn create_ns_with_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
13650 use arrow::array::{Int32Array, StringArray};
13651 use arrow::ipc::writer::StreamWriter;
13652
13653 let (namespace, temp_dir) = create_test_namespace().await;
13654
13655 let schema = create_test_schema();
13656 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13657 let arrow_schema = Arc::new(arrow_schema);
13658
13659 let id_array = Int32Array::from(vec![1, 2, 3]);
13660 let name_array = StringArray::from(vec!["Alice", "Bob", "Charlie"]);
13661 let batch = arrow::record_batch::RecordBatch::try_new(
13662 arrow_schema.clone(),
13663 vec![Arc::new(id_array), Arc::new(name_array)],
13664 )
13665 .unwrap();
13666
13667 let mut buffer = Vec::new();
13668 {
13669 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13670 writer.write(&batch).unwrap();
13671 writer.finish().unwrap();
13672 }
13673
13674 let mut request = CreateTableRequest::new();
13675 let table_id = vec!["test_ops_table".to_string()];
13676 request.id = Some(table_id.clone());
13677
13678 namespace
13679 .create_table(request, Bytes::from(buffer))
13680 .await
13681 .unwrap();
13682
13683 (namespace, temp_dir, table_id)
13684 }
13685
13686 #[tokio::test]
13687 async fn test_count_table_rows_basic() {
13688 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13689
13690 let request = CountTableRowsRequest {
13691 id: Some(table_id),
13692 version: None,
13693 predicate: None,
13694 ..Default::default()
13695 };
13696
13697 let count = namespace.count_table_rows(request).await.unwrap();
13698 assert_eq!(count, 3);
13699 }
13700
13701 #[tokio::test]
13702 async fn test_count_table_rows_with_predicate() {
13703 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13704
13705 let request = CountTableRowsRequest {
13706 id: Some(table_id),
13707 version: None,
13708 predicate: Some("id > 1".to_string()),
13709 ..Default::default()
13710 };
13711
13712 let count = namespace.count_table_rows(request).await.unwrap();
13713 assert_eq!(count, 2);
13714 }
13715
13716 #[tokio::test]
13717 async fn test_query_table_invalid_distance_type() {
13718 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
13719
13720 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
13721 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
13722 multi_vector: None,
13723 });
13724
13725 let request = QueryTableRequest {
13726 id: Some(table_id),
13727 k: 2,
13728 vector,
13729 vector_column: Some("vector".to_string()),
13730 distance_type: Some("invalid_metric".to_string()),
13731 filter: None,
13732 offset: None,
13733 version: None,
13734 ..Default::default()
13735 };
13736
13737 let result = namespace.query_table(request).await;
13738 assert!(result.is_err());
13739 let err_msg = result.unwrap_err().to_string();
13740 assert!(
13741 err_msg.contains("Unknown distance type"),
13742 "Expected error about unknown distance type, got: {}",
13743 err_msg
13744 );
13745 }
13746
13747 #[tokio::test]
13748 async fn test_insert_into_table_append() {
13749 use arrow::array::{Int32Array, StringArray};
13750 use arrow::ipc::writer::StreamWriter;
13751
13752 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13753
13754 let schema = create_test_schema();
13756 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13757 let arrow_schema = Arc::new(arrow_schema);
13758
13759 let id_array = Int32Array::from(vec![4, 5]);
13760 let name_array = StringArray::from(vec!["Dave", "Eve"]);
13761 let batch = arrow::record_batch::RecordBatch::try_new(
13762 arrow_schema.clone(),
13763 vec![Arc::new(id_array), Arc::new(name_array)],
13764 )
13765 .unwrap();
13766
13767 let mut buffer = Vec::new();
13768 {
13769 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13770 writer.write(&batch).unwrap();
13771 writer.finish().unwrap();
13772 }
13773
13774 let request = InsertIntoTableRequest {
13775 id: Some(table_id.clone()),
13776 mode: Some("append".to_string()),
13777 ..Default::default()
13778 };
13779
13780 let response = namespace
13781 .insert_into_table(request, Bytes::from(buffer))
13782 .await
13783 .unwrap();
13784 assert!(response.transaction_id.is_none());
13785
13786 let count_req = CountTableRowsRequest {
13788 id: Some(table_id),
13789 version: None,
13790 predicate: None,
13791 ..Default::default()
13792 };
13793 let count = namespace.count_table_rows(count_req).await.unwrap();
13794 assert_eq!(count, 5);
13795 }
13796
13797 #[tokio::test]
13798 async fn test_insert_into_table_overwrite() {
13799 use arrow::array::{Int32Array, StringArray};
13800 use arrow::ipc::writer::StreamWriter;
13801
13802 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13803
13804 let schema = create_test_schema();
13805 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13806 let arrow_schema = Arc::new(arrow_schema);
13807
13808 let id_array = Int32Array::from(vec![10, 20]);
13809 let name_array = StringArray::from(vec!["X", "Y"]);
13810 let batch = arrow::record_batch::RecordBatch::try_new(
13811 arrow_schema.clone(),
13812 vec![Arc::new(id_array), Arc::new(name_array)],
13813 )
13814 .unwrap();
13815
13816 let mut buffer = Vec::new();
13817 {
13818 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13819 writer.write(&batch).unwrap();
13820 writer.finish().unwrap();
13821 }
13822
13823 let request = InsertIntoTableRequest {
13824 id: Some(table_id.clone()),
13825 mode: Some("overwrite".to_string()),
13826 ..Default::default()
13827 };
13828
13829 namespace
13830 .insert_into_table(request, Bytes::from(buffer))
13831 .await
13832 .unwrap();
13833
13834 let count_req = CountTableRowsRequest {
13836 id: Some(table_id),
13837 version: None,
13838 predicate: None,
13839 ..Default::default()
13840 };
13841 let count = namespace.count_table_rows(count_req).await.unwrap();
13842 assert_eq!(count, 2);
13843 }
13844
13845 #[tokio::test]
13846 async fn test_insert_into_table_empty_data() {
13847 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13848
13849 let request = InsertIntoTableRequest {
13850 id: Some(table_id),
13851 mode: None,
13852 ..Default::default()
13853 };
13854
13855 let result = namespace.insert_into_table(request, Bytes::new()).await;
13856 assert!(result.is_err());
13857 assert!(
13858 result
13859 .unwrap_err()
13860 .to_string()
13861 .contains("Arrow IPC stream) is required")
13862 );
13863 }
13864
13865 #[tokio::test]
13866 async fn test_insert_into_table_with_storage_options() {
13867 use arrow::array::{Int32Array, StringArray};
13868 use arrow::ipc::writer::StreamWriter;
13869
13870 let temp_dir = TempStdDir::default();
13871
13872 let namespace = DirectoryNamespaceBuilder::new(temp_dir.to_str().unwrap())
13874 .storage_option("allow_http", "true")
13875 .build()
13876 .await
13877 .unwrap();
13878
13879 let schema = create_test_schema();
13881 let ipc_data = create_test_ipc_data(&schema);
13882 let mut create_req = CreateTableRequest::new();
13883 let table_id = vec!["so_table".to_string()];
13884 create_req.id = Some(table_id.clone());
13885 namespace
13886 .create_table(create_req, Bytes::from(ipc_data))
13887 .await
13888 .unwrap();
13889
13890 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
13892 let arrow_schema = Arc::new(arrow_schema);
13893
13894 let id_array = Int32Array::from(vec![10, 20]);
13895 let name_array = StringArray::from(vec!["X", "Y"]);
13896 let batch = arrow::record_batch::RecordBatch::try_new(
13897 arrow_schema.clone(),
13898 vec![Arc::new(id_array), Arc::new(name_array)],
13899 )
13900 .unwrap();
13901
13902 let mut buffer = Vec::new();
13903 {
13904 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
13905 writer.write(&batch).unwrap();
13906 writer.finish().unwrap();
13907 }
13908
13909 let request = InsertIntoTableRequest {
13910 id: Some(table_id.clone()),
13911 mode: Some("append".to_string()),
13912 ..Default::default()
13913 };
13914
13915 let response = namespace
13916 .insert_into_table(request, Bytes::from(buffer))
13917 .await
13918 .unwrap();
13919 assert!(response.transaction_id.is_none());
13920
13921 let count_req = CountTableRowsRequest {
13923 id: Some(table_id),
13924 version: None,
13925 predicate: None,
13926 ..Default::default()
13927 };
13928 let count = namespace.count_table_rows(count_req).await.unwrap();
13929 assert_eq!(count, 2);
13930 }
13931
13932 #[tokio::test]
13933 async fn test_query_table_basic() {
13934 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13935
13936 let request = QueryTableRequest {
13937 id: Some(table_id),
13938 k: 10,
13939 filter: None,
13940 offset: None,
13941 version: None,
13942 ..Default::default()
13943 };
13944
13945 let bytes = namespace.query_table(request).await.unwrap();
13946
13947 let cursor = Cursor::new(bytes.to_vec());
13949 let reader = FileReader::try_new(cursor, None).unwrap();
13950 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13951 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13952 assert_eq!(total_rows, 3);
13953 }
13954
13955 #[tokio::test]
13956 async fn test_query_table_with_filter() {
13957 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13958
13959 let request = QueryTableRequest {
13960 id: Some(table_id),
13961 k: 10,
13962 filter: Some("id <= 2".to_string()),
13963 offset: None,
13964 version: None,
13965 ..Default::default()
13966 };
13967
13968 let bytes = namespace.query_table(request).await.unwrap();
13969
13970 let cursor = Cursor::new(bytes.to_vec());
13971 let reader = FileReader::try_new(cursor, None).unwrap();
13972 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13973 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13974 assert_eq!(total_rows, 2);
13975 }
13976
13977 #[tokio::test]
13978 async fn test_query_table_with_limit_and_offset() {
13979 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
13980
13981 let request = QueryTableRequest {
13982 id: Some(table_id),
13983 k: 2,
13984 filter: None,
13985 offset: Some(1),
13986 version: None,
13987 ..Default::default()
13988 };
13989
13990 let bytes = namespace.query_table(request).await.unwrap();
13991
13992 let cursor = Cursor::new(bytes.to_vec());
13993 let reader = FileReader::try_new(cursor, None).unwrap();
13994 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
13995 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
13996 assert_eq!(total_rows, 2);
13997 }
13998
13999 #[tokio::test]
14000 async fn test_query_table_no_limit() {
14001 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14002
14003 let request = QueryTableRequest {
14005 id: Some(table_id),
14006 k: 0,
14007 filter: None,
14008 offset: None,
14009 version: None,
14010 ..Default::default()
14011 };
14012
14013 let bytes = namespace.query_table(request).await.unwrap();
14014
14015 let cursor = Cursor::new(bytes.to_vec());
14016 let reader = FileReader::try_new(cursor, None).unwrap();
14017 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14018 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14019 assert_eq!(total_rows, 3);
14020 }
14021
14022 #[tokio::test]
14023 async fn test_query_table_with_columns() {
14024 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14025
14026 let columns = Box::new(lance_namespace::models::QueryTableRequestColumns {
14027 column_names: Some(vec!["id".to_string()]),
14028 column_aliases: None,
14029 });
14030
14031 let request = QueryTableRequest {
14032 id: Some(table_id),
14033 k: 10,
14034 filter: None,
14035 offset: None,
14036 version: None,
14037 columns: Some(columns),
14038 ..Default::default()
14039 };
14040
14041 let bytes = namespace.query_table(request).await.unwrap();
14042
14043 let cursor = Cursor::new(bytes.to_vec());
14044 let reader = FileReader::try_new(cursor, None).unwrap();
14045 let schema = reader.schema();
14046 assert_eq!(schema.fields().len(), 1);
14047 assert_eq!(schema.field(0).name(), "id");
14048 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14049 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14050 assert_eq!(total_rows, 3);
14051 }
14052
14053 #[tokio::test]
14054 async fn test_count_table_rows_with_version() {
14055 use arrow::array::{Int32Array, StringArray};
14056 use arrow::ipc::writer::StreamWriter;
14057
14058 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14059
14060 let schema = create_test_schema();
14062 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
14063 let arrow_schema = Arc::new(arrow_schema);
14064
14065 let id_array = Int32Array::from(vec![4, 5]);
14066 let name_array = StringArray::from(vec!["Dave", "Eve"]);
14067 let batch = arrow::record_batch::RecordBatch::try_new(
14068 arrow_schema.clone(),
14069 vec![Arc::new(id_array), Arc::new(name_array)],
14070 )
14071 .unwrap();
14072
14073 let mut buffer = Vec::new();
14074 {
14075 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
14076 writer.write(&batch).unwrap();
14077 writer.finish().unwrap();
14078 }
14079
14080 let request = InsertIntoTableRequest {
14081 id: Some(table_id.clone()),
14082 mode: None,
14083 ..Default::default()
14084 };
14085 namespace
14086 .insert_into_table(request, Bytes::from(buffer))
14087 .await
14088 .unwrap();
14089
14090 let count_req = CountTableRowsRequest {
14092 id: Some(table_id.clone()),
14093 version: Some(1),
14094 predicate: None,
14095 ..Default::default()
14096 };
14097 let count = namespace.count_table_rows(count_req).await.unwrap();
14098 assert_eq!(count, 3);
14099
14100 let count_req = CountTableRowsRequest {
14102 id: Some(table_id),
14103 version: None,
14104 predicate: None,
14105 ..Default::default()
14106 };
14107 let count = namespace.count_table_rows(count_req).await.unwrap();
14108 assert_eq!(count, 5);
14109 }
14110
14111 #[tokio::test]
14112 async fn test_query_table_with_version() {
14113 use arrow::array::{Int32Array, StringArray};
14114 use arrow::ipc::writer::StreamWriter;
14115
14116 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14117
14118 let schema = create_test_schema();
14120 let arrow_schema = convert_json_arrow_schema(&schema).unwrap();
14121 let arrow_schema = Arc::new(arrow_schema);
14122
14123 let id_array = Int32Array::from(vec![4, 5]);
14124 let name_array = StringArray::from(vec!["Dave", "Eve"]);
14125 let batch = arrow::record_batch::RecordBatch::try_new(
14126 arrow_schema.clone(),
14127 vec![Arc::new(id_array), Arc::new(name_array)],
14128 )
14129 .unwrap();
14130
14131 let mut buffer = Vec::new();
14132 {
14133 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
14134 writer.write(&batch).unwrap();
14135 writer.finish().unwrap();
14136 }
14137
14138 let request = InsertIntoTableRequest {
14139 id: Some(table_id.clone()),
14140 mode: None,
14141 ..Default::default()
14142 };
14143 namespace
14144 .insert_into_table(request, Bytes::from(buffer))
14145 .await
14146 .unwrap();
14147
14148 let request = QueryTableRequest {
14150 id: Some(table_id.clone()),
14151 k: 100,
14152 filter: None,
14153 offset: None,
14154 version: Some(1),
14155 ..Default::default()
14156 };
14157
14158 let bytes = namespace.query_table(request).await.unwrap();
14159 let cursor = Cursor::new(bytes.to_vec());
14160 let reader = FileReader::try_new(cursor, None).unwrap();
14161 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14162 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14163 assert_eq!(total_rows, 3);
14164
14165 let request = QueryTableRequest {
14167 id: Some(table_id),
14168 k: 100,
14169 filter: None,
14170 offset: None,
14171 version: None,
14172 ..Default::default()
14173 };
14174
14175 let bytes = namespace.query_table(request).await.unwrap();
14176 let cursor = Cursor::new(bytes.to_vec());
14177 let reader = FileReader::try_new(cursor, None).unwrap();
14178 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14179 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14180 assert_eq!(total_rows, 5);
14181 }
14182
14183 async fn create_ns_with_vector_table() -> (DirectoryNamespace, TempStdDir, Vec<String>) {
14186 use arrow::array::{FixedSizeListArray, Float32Array, Int32Array};
14187 use arrow::ipc::writer::StreamWriter;
14188
14189 let (namespace, temp_dir) = create_test_namespace().await;
14190
14191 let arrow_schema = Arc::new(arrow::datatypes::Schema::new(vec![
14193 arrow::datatypes::Field::new("id", arrow::datatypes::DataType::Int32, false),
14194 arrow::datatypes::Field::new(
14195 "vector",
14196 arrow::datatypes::DataType::FixedSizeList(
14197 Arc::new(arrow::datatypes::Field::new(
14198 "item",
14199 arrow::datatypes::DataType::Float32,
14200 true,
14201 )),
14202 4,
14203 ),
14204 true,
14205 ),
14206 ]));
14207
14208 let id_array = Int32Array::from(vec![1, 2, 3]);
14209 let values = Float32Array::from(vec![
14210 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, ]);
14214 let vector_array = FixedSizeListArray::try_new(
14215 Arc::new(arrow::datatypes::Field::new(
14216 "item",
14217 arrow::datatypes::DataType::Float32,
14218 true,
14219 )),
14220 4,
14221 Arc::new(values),
14222 None,
14223 )
14224 .unwrap();
14225
14226 let batch = arrow::record_batch::RecordBatch::try_new(
14227 arrow_schema.clone(),
14228 vec![Arc::new(id_array), Arc::new(vector_array)],
14229 )
14230 .unwrap();
14231
14232 let mut buffer = Vec::new();
14233 {
14234 let mut writer = StreamWriter::try_new(&mut buffer, &arrow_schema).unwrap();
14235 writer.write(&batch).unwrap();
14236 writer.finish().unwrap();
14237 }
14238
14239 let table_name = "vector_table";
14241 let table_uri = format!("{}/{}.lance", temp_dir.to_str().unwrap(), table_name);
14242 let reader = arrow::record_batch::RecordBatchIterator::new(
14243 vec![Ok(batch)],
14244 arrow_schema.clone(),
14245 );
14246 Dataset::write(reader, &table_uri, None).await.unwrap();
14247
14248 let table_id = vec![table_name.to_string()];
14249 (namespace, temp_dir, table_id)
14250 }
14251
14252 #[tokio::test]
14253 async fn test_query_table_vector_search() {
14254 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14255
14256 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14257 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
14258 multi_vector: None,
14259 });
14260
14261 let request = QueryTableRequest {
14262 id: Some(table_id),
14263 k: 2,
14264 vector,
14265 filter: None,
14266 offset: None,
14267 version: None,
14268 ..Default::default()
14269 };
14270
14271 let bytes = namespace.query_table(request).await.unwrap();
14272
14273 let cursor = Cursor::new(bytes.to_vec());
14274 let reader = FileReader::try_new(cursor, None).unwrap();
14275 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14276 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14277 assert_eq!(total_rows, 2);
14278 }
14279
14280 #[tokio::test]
14281 async fn test_query_table_vector_search_with_distance_type() {
14282 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14283
14284 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14285 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
14286 multi_vector: None,
14287 });
14288
14289 let request = QueryTableRequest {
14290 id: Some(table_id),
14291 k: 3,
14292 vector,
14293 filter: None,
14294 offset: None,
14295 version: None,
14296 distance_type: Some("cosine".to_string()),
14297 ..Default::default()
14298 };
14299
14300 let bytes = namespace.query_table(request).await.unwrap();
14301
14302 let cursor = Cursor::new(bytes.to_vec());
14303 let reader = FileReader::try_new(cursor, None).unwrap();
14304 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14305 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14306 assert_eq!(total_rows, 3);
14307 }
14308
14309 #[tokio::test]
14310 async fn test_query_table_vector_search_with_filter() {
14311 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14312
14313 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14314 single_vector: Some(vec![1.0, 0.0, 0.0, 0.0]),
14315 multi_vector: None,
14316 });
14317
14318 let request = QueryTableRequest {
14319 id: Some(table_id),
14320 k: 10,
14321 vector,
14322 filter: Some("id <= 2".to_string()),
14323 offset: None,
14324 version: None,
14325 ..Default::default()
14326 };
14327
14328 let bytes = namespace.query_table(request).await.unwrap();
14329
14330 let cursor = Cursor::new(bytes.to_vec());
14331 let reader = FileReader::try_new(cursor, None).unwrap();
14332 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14333 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14334 assert!(total_rows <= 2);
14335 }
14336
14337 #[tokio::test]
14338 async fn test_query_table_vector_search_with_nprobes_and_refine() {
14339 let (namespace, _temp_dir, table_id) = create_ns_with_vector_table().await;
14340
14341 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14342 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
14343 multi_vector: None,
14344 });
14345
14346 let request = QueryTableRequest {
14347 id: Some(table_id),
14348 k: 2,
14349 vector,
14350 filter: None,
14351 offset: None,
14352 version: None,
14353 nprobes: Some(1),
14354 refine_factor: Some(1),
14355 prefilter: Some(true),
14356 ..Default::default()
14357 };
14358
14359 let bytes = namespace.query_table(request).await.unwrap();
14360
14361 let cursor = Cursor::new(bytes.to_vec());
14362 let reader = FileReader::try_new(cursor, None).unwrap();
14363 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14364 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14365 assert_eq!(total_rows, 2);
14366 }
14367
14368 #[tokio::test]
14369 async fn test_explain_vector_probe_fields_are_applied_independently() {
14370 use lance_namespace::models::ExplainTableQueryPlanRequest;
14371
14372 let (namespace, temp_dir, table_id) = create_ns_with_vector_table().await;
14373 let table_uri = format!("{}/vector_table.lance", temp_dir.to_str().unwrap());
14374 let mut dataset = Dataset::open(&table_uri).await.unwrap();
14375 dataset
14376 .create_index(
14377 &["vector"],
14378 IndexType::Vector,
14379 Some("vector_idx".to_string()),
14380 &VectorIndexParams::ivf_flat(1, MetricType::L2),
14381 false,
14382 )
14383 .await
14384 .unwrap();
14385 let vector = || {
14386 Box::new(lance_namespace::models::QueryTableRequestVector {
14387 single_vector: Some(vec![0.0, 1.0, 0.0, 0.0]),
14388 multi_vector: None,
14389 })
14390 };
14391
14392 let query = QueryTableRequest {
14393 id: None,
14394 k: 2,
14395 vector: vector(),
14396 nprobes: Some(20),
14397 minimum_nprobes: Some(3),
14398 maximum_nprobes: Some(10),
14399 ..Default::default()
14400 };
14401 let mut request = ExplainTableQueryPlanRequest::new(query);
14402 request.id = Some(table_id.clone());
14403
14404 let plan = namespace.explain_table_query_plan(request).await.unwrap();
14405 assert!(plan.contains("minimum_nprobes=3"), "{plan}");
14406 assert!(plan.contains("maximum_nprobes=Some(10)"), "{plan}");
14407
14408 let query = QueryTableRequest {
14409 id: None,
14410 k: 2,
14411 vector: vector(),
14412 nprobes: Some(20),
14413 ..Default::default()
14414 };
14415 let mut request = ExplainTableQueryPlanRequest::new(query);
14416 request.id = Some(table_id.clone());
14417
14418 let plan = namespace.explain_table_query_plan(request).await.unwrap();
14419 assert!(plan.contains("minimum_nprobes=20"), "{plan}");
14420 assert!(plan.contains("maximum_nprobes=Some(20)"), "{plan}");
14421
14422 let query = QueryTableRequest {
14423 id: None,
14424 k: 2,
14425 vector: vector(),
14426 nprobes: Some(-1),
14427 ..Default::default()
14428 };
14429 let mut request = ExplainTableQueryPlanRequest::new(query);
14430 request.id = Some(table_id);
14431
14432 let err = namespace
14433 .explain_table_query_plan(request)
14434 .await
14435 .unwrap_err();
14436 assert!(err.to_string().contains("nprobes must be non-negative"));
14437 }
14438
14439 #[tokio::test]
14440 async fn test_namespace_id() {
14441 let (namespace, _temp_dir) = create_test_namespace().await;
14442 let id = namespace.namespace_id();
14443 assert!(id.contains("DirectoryNamespace"));
14444 assert!(id.contains("root"));
14445 }
14446
14447 #[tokio::test]
14448 async fn test_query_table_empty_table() {
14449 let (namespace, _temp_dir) = create_test_namespace().await;
14450
14451 let schema = create_test_schema();
14453 let ipc_data = create_test_ipc_data(&schema);
14454 let mut create_request = CreateTableRequest::new();
14455 create_request.id = Some(vec!["empty_table".to_string()]);
14456 namespace
14457 .create_table(create_request, bytes::Bytes::from(ipc_data))
14458 .await
14459 .unwrap();
14460
14461 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14463 single_vector: None,
14464 multi_vector: None,
14465 });
14466 let request = QueryTableRequest {
14467 id: Some(vec!["empty_table".to_string()]),
14468 k: 10,
14469 vector,
14470 ..Default::default()
14471 };
14472 let bytes = namespace.query_table(request).await.unwrap();
14473
14474 let cursor = Cursor::new(bytes.to_vec());
14475 let reader = FileReader::try_new(cursor, None).unwrap();
14476 let batches: Vec<_> = reader.collect::<std::result::Result<Vec<_>, _>>().unwrap();
14477 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14478 assert_eq!(total_rows, 0, "empty table should yield no rows");
14479 }
14480
14481 #[tokio::test]
14482 async fn test_query_table_with_plain_filter_no_vector() {
14483 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14484
14485 let vector = Box::new(lance_namespace::models::QueryTableRequestVector {
14487 single_vector: None,
14488 multi_vector: None,
14489 });
14490 let request = QueryTableRequest {
14491 id: Some(table_id),
14492 k: 0,
14493 vector,
14494 filter: Some("id > 1".to_string()),
14495 ..Default::default()
14496 };
14497 let bytes = namespace.query_table(request).await.unwrap();
14498
14499 let cursor = Cursor::new(bytes.to_vec());
14500 let reader = FileReader::try_new(cursor, None).unwrap();
14501 let batches: Vec<_> = reader.into_iter().map(|b| b.unwrap()).collect();
14502 let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
14503 assert!(total_rows > 0);
14504 assert!(total_rows < 3);
14505 }
14506
14507 #[tokio::test]
14510 async fn test_update_full_table() {
14511 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14512
14513 let base_version = open_dataset(&namespace, &table_id[0])
14515 .await
14516 .version()
14517 .version;
14518
14519 let request = UpdateTableRequest {
14520 id: Some(table_id.clone()),
14521 updates: vec![vec!["name".to_string(), "'updated'".to_string()]],
14522 predicate: None,
14523 ..Default::default()
14524 };
14525
14526 let response = namespace.update_table(request).await.unwrap();
14527 assert_eq!(response.updated_rows, 3);
14528 assert!(response.version as u64 > base_version);
14529
14530 let count_req = CountTableRowsRequest {
14532 id: Some(table_id),
14533 version: None,
14534 predicate: Some("name = 'updated'".to_string()),
14535 ..Default::default()
14536 };
14537 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 3);
14538 }
14539
14540 #[tokio::test]
14541 async fn test_update_with_predicate() {
14542 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14543
14544 let request = UpdateTableRequest {
14545 id: Some(table_id.clone()),
14546 updates: vec![vec!["name".to_string(), "'matched'".to_string()]],
14547 predicate: Some("id > 1".to_string()),
14548 ..Default::default()
14549 };
14550
14551 let response = namespace.update_table(request).await.unwrap();
14552 assert_eq!(response.updated_rows, 2);
14553
14554 let untouched = CountTableRowsRequest {
14556 id: Some(table_id.clone()),
14557 version: None,
14558 predicate: Some("name = 'Alice'".to_string()),
14559 ..Default::default()
14560 };
14561 assert_eq!(namespace.count_table_rows(untouched).await.unwrap(), 1);
14562
14563 let touched = CountTableRowsRequest {
14564 id: Some(table_id),
14565 version: None,
14566 predicate: Some("name = 'matched'".to_string()),
14567 ..Default::default()
14568 };
14569 assert_eq!(namespace.count_table_rows(touched).await.unwrap(), 2);
14570 }
14571
14572 #[tokio::test]
14573 async fn test_update_invalid_expression_returns_invalid_input() {
14574 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14575
14576 let request = UpdateTableRequest {
14577 id: Some(table_id),
14578 updates: vec![vec!["name".to_string(), "no_such_column + 1".to_string()]],
14580 predicate: None,
14581 ..Default::default()
14582 };
14583
14584 let err = namespace.update_table(request).await.unwrap_err();
14585 let msg = err.to_string();
14586 assert!(
14587 msg.contains("Invalid input"),
14588 "expected Invalid input error, got: {}",
14589 msg
14590 );
14591 }
14592
14593 #[tokio::test]
14594 async fn test_update_rejects_duplicate_columns() {
14595 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14596
14597 let request = UpdateTableRequest {
14598 id: Some(table_id),
14599 updates: vec![
14600 vec!["name".to_string(), "'a'".to_string()],
14601 vec!["name".to_string(), "'b'".to_string()],
14602 ],
14603 predicate: None,
14604 ..Default::default()
14605 };
14606
14607 let err = namespace.update_table(request).await.unwrap_err();
14608 let msg = err.to_string();
14609 assert!(
14610 msg.contains("Invalid input") && msg.contains("more than once"),
14611 "expected duplicate column InvalidInput error, got: {}",
14612 msg
14613 );
14614 }
14615
14616 #[tokio::test]
14617 async fn test_delete_with_predicate() {
14618 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14619
14620 let request = DeleteFromTableRequest {
14621 id: Some(table_id.clone()),
14622 predicate: "id > 1".to_string(),
14623 ..Default::default()
14624 };
14625
14626 let response = namespace.delete_from_table(request).await.unwrap();
14627 assert!(response.version.is_some());
14628
14629 let count_req = CountTableRowsRequest {
14630 id: Some(table_id),
14631 version: None,
14632 predicate: None,
14633 ..Default::default()
14634 };
14635 assert_eq!(namespace.count_table_rows(count_req).await.unwrap(), 1);
14637 }
14638
14639 #[tokio::test]
14640 async fn test_delete_empty_predicate_returns_invalid_input() {
14641 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14642
14643 let request = DeleteFromTableRequest {
14644 id: Some(table_id),
14645 predicate: " ".to_string(),
14646 ..Default::default()
14647 };
14648
14649 let err = namespace.delete_from_table(request).await.unwrap_err();
14650 let msg = err.to_string();
14651 assert!(
14652 msg.contains("Invalid input") && msg.contains("non-empty predicate"),
14653 "expected non-empty predicate InvalidInput error, got: {}",
14654 msg
14655 );
14656 }
14657
14658 #[tokio::test]
14659 async fn test_delete_table_not_found() {
14660 let (namespace, _temp_dir) = create_test_namespace().await;
14661
14662 let request = DeleteFromTableRequest {
14663 id: Some(vec!["does_not_exist".to_string()]),
14664 predicate: "id = 1".to_string(),
14665 ..Default::default()
14666 };
14667
14668 let err = namespace.delete_from_table(request).await.unwrap_err();
14669 let msg = err.to_string();
14670 assert!(
14671 msg.contains("Table not found"),
14672 "expected TableNotFound for missing table, got: {}",
14673 msg
14674 );
14675 }
14676
14677 #[tokio::test]
14678 async fn test_delete_invalid_predicate_returns_invalid_input() {
14679 let (namespace, _temp_dir, table_id) = create_ns_with_table().await;
14680
14681 let request = DeleteFromTableRequest {
14685 id: Some(table_id),
14686 predicate: "no_such_column = 1".to_string(),
14687 ..Default::default()
14688 };
14689
14690 let err = namespace.delete_from_table(request).await.unwrap_err();
14691 let lance_core::Error::Namespace { source, .. } = &err else {
14692 panic!("expected a Namespace error, got: {}", err);
14693 };
14694 let ns_err = source
14695 .downcast_ref::<NamespaceError>()
14696 .expect("expected a NamespaceError source");
14697 assert_eq!(
14698 ns_err.code(),
14699 lance_namespace::ErrorCode::InvalidInput,
14700 "expected InvalidInput for an invalid delete predicate, got: {}",
14701 err
14702 );
14703 }
14704 }
14705
14706 #[tokio::test]
14707 async fn test_list_all_tables() {
14708 use lance_namespace::models::ListTablesRequest;
14709
14710 let (namespace, _temp_dir) = create_test_namespace().await;
14711 create_scalar_table(&namespace, "alpha").await;
14712 create_scalar_table(&namespace, "beta").await;
14713
14714 let request = ListTablesRequest {
14715 id: Some(vec![]),
14716 page_token: None,
14717 limit: None,
14718 ..Default::default()
14719 };
14720 let response = namespace.list_all_tables(request).await.unwrap();
14721 let mut tables = response.tables;
14722 tables.sort();
14723 assert_eq!(tables, vec!["alpha", "beta"]);
14724 }
14725
14726 #[tokio::test]
14727 async fn test_restore_table() {
14728 use lance_namespace::models::RestoreTableRequest;
14729
14730 let (namespace, _temp_dir) = create_test_namespace().await;
14731 create_scalar_table(&namespace, "users").await;
14732
14733 create_scalar_index(&namespace, "users", "users_id_idx").await;
14735
14736 let dataset = open_dataset(&namespace, "users").await;
14737 let current_version = dataset.version().version;
14738 assert!(current_version >= 2, "Should have at least 2 versions");
14739
14740 let mut restore_req = RestoreTableRequest::new(1);
14742 restore_req.id = Some(vec!["users".to_string()]);
14743 let response = namespace.restore_table(restore_req).await.unwrap();
14744
14745 assert!(
14747 response.transaction_id.is_some(),
14748 "restore_table should return a transaction_id"
14749 );
14750
14751 let dataset_after = open_dataset(&namespace, "users").await;
14753 assert!(
14754 dataset_after.version().version > current_version,
14755 "Restore should create a new version"
14756 );
14757 }
14758
14759 #[tokio::test]
14760 async fn test_alter_table_add_columns() {
14761 use lance_namespace::models::{
14762 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
14763 };
14764
14765 let (namespace, _temp_dir) = create_test_namespace().await;
14766
14767 let schema = create_test_schema();
14769 let ipc_data = create_test_ipc_data(&schema);
14770 let mut create_request = CreateTableRequest::new();
14771 create_request.id = Some(vec!["test_table".to_string()]);
14772 namespace
14773 .create_table(create_request, bytes::Bytes::from(ipc_data))
14774 .await
14775 .unwrap();
14776
14777 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
14779 new_col.expression = Some(Some("id * 2".to_string()));
14780 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
14781 add_request.id = Some(vec!["test_table".to_string()]);
14782
14783 let response = namespace
14784 .alter_table_add_columns(add_request)
14785 .await
14786 .unwrap();
14787 assert!(
14788 response.version > 1,
14789 "Version should increment after adding columns"
14790 );
14791
14792 let mut describe_request = DescribeTableRequest::new();
14794 describe_request.id = Some(vec!["test_table".to_string()]);
14795 describe_request.load_detailed_metadata = Some(true);
14796 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14797 assert!(describe_response.schema.is_some());
14798
14799 let resp_schema = describe_response.schema.unwrap();
14800 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14801 assert!(
14802 field_names.contains(&"doubled_id"),
14803 "Column 'doubled_id' should exist, got: {:?}",
14804 field_names
14805 );
14806 }
14807
14808 #[tokio::test]
14809 async fn test_update_table_schema_metadata() {
14810 use lance_namespace::models::UpdateTableSchemaMetadataRequest;
14811
14812 let (namespace, _temp_dir) = create_test_namespace().await;
14813 create_scalar_table(&namespace, "products").await;
14814
14815 let mut metadata = HashMap::new();
14816 metadata.insert("owner".to_string(), "team_a".to_string());
14817 metadata.insert("version".to_string(), "1.0".to_string());
14818
14819 let mut req = UpdateTableSchemaMetadataRequest::new();
14820 req.id = Some(vec!["products".to_string()]);
14821 req.metadata = Some(metadata.clone());
14822
14823 let response = namespace.update_table_schema_metadata(req).await.unwrap();
14824
14825 assert!(response.metadata.is_some());
14826 let returned = response.metadata.unwrap();
14827 assert_eq!(returned.get("owner"), Some(&"team_a".to_string()));
14828 assert_eq!(returned.get("version"), Some(&"1.0".to_string()));
14829 assert!(
14830 response.transaction_id.is_some(),
14831 "update_table_schema_metadata should return a transaction_id"
14832 );
14833 }
14834
14835 #[tokio::test]
14836 async fn test_alter_table_add_columns_missing_id() {
14837 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
14838
14839 let (namespace, _temp_dir) = create_test_namespace().await;
14840
14841 let new_col = AddColumnsEntry::new("col".to_string());
14842 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
14843 let result = namespace.alter_table_add_columns(request).await;
14844 assert!(result.is_err(), "Should fail when table ID is missing");
14845 }
14846
14847 #[tokio::test]
14848 async fn test_alter_table_alter_columns_rename() {
14849 use lance_namespace::models::{
14850 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
14851 };
14852
14853 let (namespace, _temp_dir) = create_test_namespace().await;
14854
14855 let schema = create_test_schema();
14857 let ipc_data = create_test_ipc_data(&schema);
14858 let mut create_request = CreateTableRequest::new();
14859 create_request.id = Some(vec!["test_table".to_string()]);
14860 namespace
14861 .create_table(create_request, bytes::Bytes::from(ipc_data))
14862 .await
14863 .unwrap();
14864
14865 let mut entry = AlterColumnsEntry::new("name".to_string());
14867 entry.rename = Some(Some("full_name".to_string()));
14868 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
14869 alter_request.id = Some(vec!["test_table".to_string()]);
14870
14871 let response = namespace
14872 .alter_table_alter_columns(alter_request)
14873 .await
14874 .unwrap();
14875 assert!(
14876 response.version > 1,
14877 "Version should increment after altering columns"
14878 );
14879
14880 let mut describe_request = DescribeTableRequest::new();
14882 describe_request.id = Some(vec!["test_table".to_string()]);
14883 describe_request.load_detailed_metadata = Some(true);
14884 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14885 assert!(describe_response.schema.is_some());
14886
14887 let resp_schema = describe_response.schema.unwrap();
14888 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
14889 assert!(
14890 field_names.contains(&"full_name"),
14891 "Column should be renamed to 'full_name', got: {:?}",
14892 field_names
14893 );
14894 assert!(
14895 !field_names.contains(&"name"),
14896 "Old column 'name' should not exist, got: {:?}",
14897 field_names
14898 );
14899 }
14900
14901 #[tokio::test]
14902 async fn test_get_table_stats() {
14903 use lance_namespace::models::GetTableStatsRequest;
14904
14905 let (namespace, _temp_dir) = create_test_namespace().await;
14906 create_scalar_table(&namespace, "items").await;
14907 create_scalar_index(&namespace, "items", "items_id_idx").await;
14908
14909 let mut req = GetTableStatsRequest::new();
14910 req.id = Some(vec!["items".to_string()]);
14911
14912 let response = namespace.get_table_stats(req).await.unwrap();
14913 assert_eq!(response.num_rows, 3);
14914 assert_eq!(response.num_indices, 1);
14915 }
14916
14917 #[tokio::test]
14918 async fn test_explain_table_query_plan() {
14919 use lance_namespace::models::QueryTableRequestVector;
14920 use lance_namespace::models::{ExplainTableQueryPlanRequest, QueryTableRequest};
14921
14922 let (namespace, _temp_dir) = create_test_namespace().await;
14923 create_scalar_table(&namespace, "catalog").await;
14924
14925 let mut query = QueryTableRequest::new(1, QueryTableRequestVector::new());
14926 query.filter = Some("id > 1".to_string());
14927 query.columns = Some(Box::new(QueryTableRequestColumns {
14928 column_names: Some(vec!["id".to_string(), "name".to_string()]),
14929 column_aliases: None,
14930 }));
14931 query.with_row_id = Some(true);
14932
14933 let mut req = ExplainTableQueryPlanRequest::new(query);
14934 req.id = Some(vec!["catalog".to_string()]);
14935
14936 let plan_str = namespace.explain_table_query_plan(req).await.unwrap();
14937 assert_plan_contains_all(
14938 &plan_str,
14939 &[
14940 "ProjectionExec: expr=[id@0 as id, name@2 as name",
14941 "projection=[name], source=stream(_rowid)",
14942 "LanceRead: uri=",
14943 "projection=[id]",
14944 "row_id=true, row_addr=false",
14945 "full_filter=id > Int32(1)",
14946 "refine_filter=id > Int32(1)",
14947 ],
14948 "Filtered explain plan should preserve late materialization and filter pushdown",
14949 );
14950 }
14951
14952 #[tokio::test]
14953 async fn test_alter_table_alter_columns_missing_id() {
14954 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
14955
14956 let (namespace, _temp_dir) = create_test_namespace().await;
14957
14958 let entry = AlterColumnsEntry::new("name".to_string());
14959 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
14960 let result = namespace.alter_table_alter_columns(request).await;
14961 assert!(result.is_err(), "Should fail when table ID is missing");
14962 }
14963
14964 #[tokio::test]
14965 async fn test_alter_table_drop_columns() {
14966 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
14967
14968 let (namespace, _temp_dir) = create_test_namespace().await;
14969
14970 let schema = create_test_schema();
14972 let ipc_data = create_test_ipc_data(&schema);
14973 let mut create_request = CreateTableRequest::new();
14974 create_request.id = Some(vec!["test_table".to_string()]);
14975 namespace
14976 .create_table(create_request, bytes::Bytes::from(ipc_data))
14977 .await
14978 .unwrap();
14979
14980 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
14982 drop_request.id = Some(vec!["test_table".to_string()]);
14983
14984 let response = namespace
14985 .alter_table_drop_columns(drop_request)
14986 .await
14987 .unwrap();
14988 assert!(
14989 response.version > 1,
14990 "Version should increment after dropping columns"
14991 );
14992
14993 let mut describe_request = DescribeTableRequest::new();
14995 describe_request.id = Some(vec!["test_table".to_string()]);
14996 describe_request.load_detailed_metadata = Some(true);
14997 let describe_response = namespace.describe_table(describe_request).await.unwrap();
14998 assert!(describe_response.schema.is_some());
14999
15000 let resp_schema = describe_response.schema.unwrap();
15001 let field_names: Vec<&str> = resp_schema.fields.iter().map(|f| f.name.as_str()).collect();
15002 assert!(
15003 !field_names.contains(&"name"),
15004 "Column 'name' should be dropped, got: {:?}",
15005 field_names
15006 );
15007 assert!(
15008 field_names.contains(&"id"),
15009 "Column 'id' should still exist, got: {:?}",
15010 field_names
15011 );
15012 }
15013
15014 #[tokio::test]
15015 async fn test_analyze_table_query_plan() {
15016 use lance_namespace::models::AnalyzeTableQueryPlanRequest;
15017 use lance_namespace::models::QueryTableRequestVector;
15018
15019 let (namespace, _temp_dir) = create_test_namespace().await;
15020 create_scalar_table(&namespace, "catalog").await;
15021
15022 let mut req = AnalyzeTableQueryPlanRequest::new(1, QueryTableRequestVector::new());
15023 req.id = Some(vec!["catalog".to_string()]);
15024 req.filter = Some("id > 0".to_string());
15025 req.columns = Some(Box::new(QueryTableRequestColumns {
15026 column_names: Some(vec!["id".to_string(), "name".to_string()]),
15027 column_aliases: None,
15028 }));
15029 req.with_row_id = Some(true);
15030
15031 let analysis_str = namespace.analyze_table_query_plan(req).await.unwrap();
15032 assert_plan_contains_all(
15033 &analysis_str,
15034 &[
15035 "AnalyzeExec verbose=true",
15036 "ProjectionExec: elapsed=",
15037 "expr=[id@0 as id, name@2 as name",
15038 "projection=[name], source=stream(_rowid)",
15039 "LanceRead: elapsed=",
15040 "projection=[id]",
15041 "row_id=true, row_addr=false",
15042 "full_filter=id > Int32(0)",
15043 "refine_filter=id > Int32(0)",
15044 "metrics=[output_rows=",
15045 ],
15046 "Filtered analyze plan should preserve late materialization and filter pushdown",
15047 );
15048 }
15049
15050 #[tokio::test]
15051 async fn test_dir_listing_no_extra_calls_without_migration() {
15052 let temp_dir = TempStdDir::default();
15053 let temp_path = temp_dir.to_str().unwrap();
15054 let root_uri = file_object_store_uri(temp_path);
15055 let listing_count = Arc::new(AtomicUsize::new(0));
15056 let session = build_listing_counting_session(listing_count.clone());
15057
15058 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
15060 .session(session.clone())
15061 .manifest_enabled(false)
15062 .dir_listing_enabled(true)
15063 .build()
15064 .await
15065 .unwrap();
15066
15067 let schema = create_test_schema();
15068 let ipc_data = create_test_ipc_data(&schema);
15069 let mut create_req = CreateTableRequest::new();
15070 create_req.id = Some(vec!["test_table".to_string()]);
15071 dir_only_ns
15072 .create_table(create_req, Bytes::from(ipc_data))
15073 .await
15074 .unwrap();
15075
15076 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
15078 .session(session)
15079 .manifest_enabled(true)
15080 .dir_listing_enabled(true)
15081 .dir_listing_to_manifest_migration_enabled(false)
15082 .build()
15083 .await
15084 .unwrap();
15085
15086 listing_count.store(0, Ordering::SeqCst);
15088
15089 let mut exists_req = TableExistsRequest::new();
15091 exists_req.id = Some(vec!["test_table".to_string()]);
15092 hybrid_ns.table_exists(exists_req).await.unwrap();
15093
15094 let count = listing_count.load(Ordering::SeqCst);
15095 assert_eq!(
15096 count, 1,
15097 "Expected exactly 1 listing call for table_exists \
15098 without migration mode, but got {}",
15099 count
15100 );
15101
15102 listing_count.store(0, Ordering::SeqCst);
15104
15105 let mut describe_req = DescribeTableRequest::new();
15106 describe_req.id = Some(vec!["test_table".to_string()]);
15107 hybrid_ns.describe_table(describe_req).await.unwrap();
15108
15109 let count = listing_count.load(Ordering::SeqCst);
15110 assert_eq!(
15111 count, 1,
15112 "Expected exactly 1 listing call for describe_table \
15113 without migration mode, but got {}",
15114 count
15115 );
15116 }
15117
15118 #[tokio::test]
15119 async fn test_build_and_root_reads_do_not_create_manifest() {
15120 let temp_dir = TempStdDir::default();
15121 let temp_path = temp_dir.to_str().unwrap();
15122 let manifest_path = std::path::Path::new(temp_path).join("__manifest");
15123
15124 let dir_only_ns = DirectoryNamespaceBuilder::new(temp_path)
15125 .manifest_enabled(false)
15126 .dir_listing_enabled(true)
15127 .build()
15128 .await
15129 .unwrap();
15130 create_scalar_table(&dir_only_ns, "catalog").await;
15131 assert!(!manifest_path.exists());
15132
15133 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15134 .manifest_enabled(true)
15135 .dir_listing_enabled(true)
15136 .build()
15137 .await
15138 .unwrap();
15139 assert!(!manifest_path.exists());
15140
15141 let mut exists_req = TableExistsRequest::new();
15142 exists_req.id = Some(vec!["catalog".to_string()]);
15143 namespace.table_exists(exists_req).await.unwrap();
15144 assert!(!manifest_path.exists());
15145
15146 let mut describe_req = DescribeTableRequest::new();
15147 describe_req.id = Some(vec!["catalog".to_string()]);
15148 namespace.describe_table(describe_req).await.unwrap();
15149 assert!(!manifest_path.exists());
15150
15151 let list_response = namespace
15152 .list_tables(ListTablesRequest {
15153 id: Some(vec![]),
15154 ..Default::default()
15155 })
15156 .await
15157 .unwrap();
15158 assert_eq!(list_response.tables, vec!["catalog".to_string()]);
15159 assert!(!manifest_path.exists());
15160
15161 let mut list_namespaces_req = ListNamespacesRequest::new();
15162 list_namespaces_req.id = Some(vec!["workspace".to_string()]);
15163 let err = namespace
15164 .list_namespaces(list_namespaces_req)
15165 .await
15166 .unwrap_err();
15167 assert!(err.to_string().contains("__manifest"));
15168 assert!(!manifest_path.exists());
15169
15170 let err = namespace
15171 .list_tables(ListTablesRequest {
15172 id: Some(vec!["workspace".to_string()]),
15173 ..Default::default()
15174 })
15175 .await
15176 .unwrap_err();
15177 assert!(err.to_string().contains("__manifest"));
15178 assert!(!manifest_path.exists());
15179
15180 let mut child_describe_req = DescribeTableRequest::new();
15181 child_describe_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
15182 let err = namespace
15183 .describe_table(child_describe_req)
15184 .await
15185 .unwrap_err();
15186 assert!(err.to_string().contains("__manifest"));
15187 assert!(!manifest_path.exists());
15188
15189 let mut child_exists_req = TableExistsRequest::new();
15190 child_exists_req.id = Some(vec!["workspace".to_string(), "catalog".to_string()]);
15191 let err = namespace.table_exists(child_exists_req).await.unwrap_err();
15192 assert!(err.to_string().contains("__manifest"));
15193 assert!(!manifest_path.exists());
15194
15195 let mut create_ns_req = CreateNamespaceRequest::new();
15196 create_ns_req.id = Some(vec!["workspace".to_string()]);
15197 namespace.create_namespace(create_ns_req).await.unwrap();
15198 assert!(manifest_path.exists());
15199 }
15200
15201 #[tokio::test]
15202 async fn test_migrate_updates_read_opened_legacy_manifest() {
15203 let temp_dir = TempStdDir::default();
15204 let temp_path = temp_dir.to_str().unwrap();
15205 create_legacy_manifest_without_primary_key_metadata(temp_path).await;
15206 assert!(!manifest_has_primary_key_metadata(temp_path).await);
15207
15208 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15209 .manifest_enabled(true)
15210 .dir_listing_enabled(true)
15211 .build()
15212 .await
15213 .unwrap();
15214 assert!(!manifest_has_primary_key_metadata(temp_path).await);
15215
15216 let migrated = namespace.migrate().await.unwrap();
15217 assert_eq!(migrated, 0);
15218 assert!(manifest_has_primary_key_metadata(temp_path).await);
15219 }
15220
15221 #[tokio::test]
15222 async fn test_describe_declared_table_checks_versions_only_when_requested() {
15223 let temp_dir = TempStdDir::default();
15224 let temp_path = temp_dir.to_str().unwrap();
15225 let root_uri = file_object_store_uri(temp_path);
15226 let listing_count = Arc::new(AtomicUsize::new(0));
15227 let session = build_listing_counting_session(listing_count.clone());
15228
15229 let namespace = DirectoryNamespaceBuilder::new(root_uri)
15230 .session(session)
15231 .manifest_enabled(false)
15232 .dir_listing_enabled(true)
15233 .build()
15234 .await
15235 .unwrap();
15236
15237 let mut declare_req = DeclareTableRequest::new();
15238 declare_req.id = Some(vec!["test_table".to_string()]);
15239 namespace.declare_table(declare_req).await.unwrap();
15240
15241 listing_count.store(0, Ordering::SeqCst);
15242
15243 let mut describe_req = DescribeTableRequest::new();
15244 describe_req.id = Some(vec!["test_table".to_string()]);
15245 let describe_response = namespace.describe_table(describe_req).await.unwrap();
15246
15247 assert_eq!(describe_response.is_only_declared, None);
15248 assert_eq!(
15249 listing_count.load(Ordering::SeqCst),
15250 1,
15251 "Default describe_table should only list the table directory"
15252 );
15253
15254 listing_count.store(0, Ordering::SeqCst);
15255
15256 let mut describe_req = DescribeTableRequest::new();
15257 describe_req.id = Some(vec!["test_table".to_string()]);
15258 describe_req.check_declared = Some(true);
15259 let describe_response = namespace.describe_table(describe_req).await.unwrap();
15260
15261 assert_eq!(describe_response.is_only_declared, Some(true));
15262 assert_eq!(
15263 listing_count.load(Ordering::SeqCst),
15264 2,
15265 "check_declared describe_table should list the table directory and _versions"
15266 );
15267 }
15268
15269 #[tokio::test]
15270 async fn test_dir_listing_extra_calls_with_migration() {
15271 let temp_dir = TempStdDir::default();
15272 let temp_path = temp_dir.to_str().unwrap();
15273 let root_uri = file_object_store_uri(temp_path);
15274 let listing_count = Arc::new(AtomicUsize::new(0));
15275 let session = build_listing_counting_session(listing_count.clone());
15276
15277 let dir_only_ns = DirectoryNamespaceBuilder::new(root_uri.clone())
15279 .session(session.clone())
15280 .manifest_enabled(false)
15281 .dir_listing_enabled(true)
15282 .build()
15283 .await
15284 .unwrap();
15285
15286 let schema = create_test_schema();
15287 let ipc_data = create_test_ipc_data(&schema);
15288 let mut create_req = CreateTableRequest::new();
15289 create_req.id = Some(vec!["test_table".to_string()]);
15290 dir_only_ns
15291 .create_table(create_req, Bytes::from(ipc_data))
15292 .await
15293 .unwrap();
15294
15295 let hybrid_ns = DirectoryNamespaceBuilder::new(root_uri)
15296 .session(session)
15297 .manifest_enabled(true)
15298 .dir_listing_enabled(true)
15299 .dir_listing_to_manifest_migration_enabled(true)
15300 .build()
15301 .await
15302 .unwrap();
15303
15304 listing_count.store(0, Ordering::SeqCst);
15312
15313 let mut exists_req = TableExistsRequest::new();
15314 exists_req.id = Some(vec!["test_table".to_string()]);
15315 hybrid_ns.table_exists(exists_req).await.unwrap();
15316
15317 let count = listing_count.load(Ordering::SeqCst);
15318 assert_eq!(
15319 count, 2,
15320 "Expected 2 listing calls for table_exists with migration mode \
15321 (absent-__manifest probe + table directory fallback), but got {}",
15322 count
15323 );
15324
15325 listing_count.store(0, Ordering::SeqCst);
15328
15329 let mut describe_req = DescribeTableRequest::new();
15330 describe_req.id = Some(vec!["test_table".to_string()]);
15331 hybrid_ns.describe_table(describe_req).await.unwrap();
15332
15333 let count = listing_count.load(Ordering::SeqCst);
15334 assert_eq!(
15335 count, 2,
15336 "Expected 2 listing calls for describe_table with migration mode \
15337 (absent-__manifest probe + table directory fallback), but got {}",
15338 count
15339 );
15340 }
15341
15342 #[tokio::test]
15343 async fn test_manifest_reload_observes_new_version_from_other_namespace() {
15344 let temp_dir = TempStdDir::default();
15345 let temp_path = temp_dir.to_str().unwrap();
15346
15347 let namespace_a = DirectoryNamespaceBuilder::new(temp_path)
15348 .manifest_enabled(true)
15349 .dir_listing_enabled(false)
15350 .build()
15351 .await
15352 .unwrap();
15353 create_scalar_table(&namespace_a, "alpha").await;
15354
15355 let namespace_b = DirectoryNamespaceBuilder::new(temp_path)
15356 .manifest_enabled(true)
15357 .dir_listing_enabled(false)
15358 .build()
15359 .await
15360 .unwrap();
15361 create_scalar_table(&namespace_b, "beta").await;
15362
15363 let response = namespace_a
15364 .list_tables(ListTablesRequest {
15365 id: Some(vec![]),
15366 ..Default::default()
15367 })
15368 .await
15369 .unwrap();
15370
15371 let mut tables = response.tables;
15372 tables.sort();
15373 assert_eq!(tables, vec!["alpha", "beta"]);
15374 }
15375
15376 #[tokio::test]
15377 async fn test_migration_not_found_errors_include_table_id() {
15378 let temp_dir = TempStdDir::default();
15379 let temp_path = temp_dir.to_str().unwrap();
15380
15381 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15382 .manifest_enabled(true)
15383 .dir_listing_enabled(true)
15384 .dir_listing_to_manifest_migration_enabled(true)
15385 .build()
15386 .await
15387 .unwrap();
15388
15389 let mut exists_req = TableExistsRequest::new();
15390 exists_req.id = Some(vec!["missing_table".to_string()]);
15391 let err = namespace.table_exists(exists_req).await.unwrap_err();
15392 assert!(matches!(err, Error::Namespace { .. }));
15393 let err_msg = err.to_string();
15394 assert!(err_msg.contains("Table not found"));
15395 assert!(err_msg.contains("table id 'missing_table'"));
15396
15397 let mut describe_req = DescribeTableRequest::new();
15398 describe_req.id = Some(vec!["missing_table".to_string()]);
15399 let err = namespace.describe_table(describe_req).await.unwrap_err();
15400 assert!(matches!(err, Error::Namespace { .. }));
15401 let err_msg = err.to_string();
15402 assert!(err_msg.contains("Table not found"));
15403 assert!(err_msg.contains("table id 'missing_table'"));
15404 }
15405
15406 #[tokio::test]
15407 async fn test_manifest_not_found_errors_include_full_table_id() {
15408 use lance_namespace::models::CreateNamespaceRequest;
15409
15410 let temp_dir = TempStdDir::default();
15411 let temp_path = temp_dir.to_str().unwrap();
15412
15413 let namespace = DirectoryNamespaceBuilder::new(temp_path)
15414 .manifest_enabled(true)
15415 .dir_listing_enabled(true)
15416 .build()
15417 .await
15418 .unwrap();
15419
15420 let mut create_ns_req = CreateNamespaceRequest::new();
15421 create_ns_req.id = Some(vec!["workspace".to_string()]);
15422 namespace.create_namespace(create_ns_req).await.unwrap();
15423
15424 let missing_table_id = vec!["workspace".to_string(), "missing_table".to_string()];
15425
15426 let mut exists_req = TableExistsRequest::new();
15427 exists_req.id = Some(missing_table_id.clone());
15428 let err = namespace.table_exists(exists_req).await.unwrap_err();
15429 assert!(matches!(err, Error::Namespace { .. }));
15430 let err_msg = err.to_string();
15431 assert!(err_msg.contains("Table not found"));
15432 assert!(err_msg.contains("table id 'workspace$missing_table'"));
15433
15434 let mut describe_req = DescribeTableRequest::new();
15435 describe_req.id = Some(missing_table_id);
15436 let err = namespace.describe_table(describe_req).await.unwrap_err();
15437 assert!(matches!(err, Error::Namespace { .. }));
15438 let err_msg = err.to_string();
15439 assert!(err_msg.contains("Table not found"));
15440 assert!(err_msg.contains("table id 'workspace$missing_table'"));
15441 }
15442
15443 async fn create_tagged_test_table(
15446 versions: u32,
15447 ) -> (Arc<DirectoryNamespace>, TempStdDir, Vec<String>) {
15448 use arrow::array::{Int32Array, RecordBatchIterator};
15449 use arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
15450 use arrow::record_batch::RecordBatch;
15451 use lance::dataset::{Dataset, WriteMode, WriteParams};
15452
15453 assert!(versions >= 1, "versions must be at least 1");
15454
15455 let temp_dir = TempStdDir::default();
15456 let temp_path = temp_dir.to_str().unwrap();
15457
15458 let namespace = Arc::new(
15459 DirectoryNamespaceBuilder::new(temp_path)
15460 .build()
15461 .await
15462 .unwrap(),
15463 );
15464 let table_id = vec!["tag_table".to_string()];
15465 let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new(
15466 "id",
15467 DataType::Int32,
15468 false,
15469 )]));
15470 let initial_batch = RecordBatch::try_new(
15471 arrow_schema.clone(),
15472 vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
15473 )
15474 .unwrap();
15475 let batches = RecordBatchIterator::new(vec![Ok(initial_batch)], arrow_schema.clone());
15476 let write_params = WriteParams {
15477 mode: WriteMode::Create,
15478 ..Default::default()
15479 };
15480
15481 let mut dataset = Dataset::write_into_namespace(
15482 batches,
15483 namespace.clone() as Arc<dyn LanceNamespace>,
15484 table_id.clone(),
15485 Some(write_params),
15486 )
15487 .await
15488 .unwrap();
15489
15490 for i in 1..versions {
15491 let value_start = (i as i32) * 10;
15492 let batch = RecordBatch::try_new(
15493 arrow_schema.clone(),
15494 vec![Arc::new(Int32Array::from(vec![
15495 value_start,
15496 value_start + 1,
15497 ]))],
15498 )
15499 .unwrap();
15500 let batches = RecordBatchIterator::new(vec![Ok(batch)], arrow_schema.clone());
15501 dataset.append(batches, None).await.unwrap();
15502 }
15503
15504 (namespace, temp_dir, table_id)
15505 }
15506
15507 fn namespace_code(err: &Error) -> Option<ErrorCode> {
15509 match err {
15510 Error::Namespace { source, .. } => {
15511 source.downcast_ref::<NamespaceError>().map(|e| e.code())
15512 }
15513 _ => None,
15514 }
15515 }
15516
15517 #[tokio::test]
15518 async fn test_create_and_list_branches() {
15519 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15520
15521 namespace
15522 .create_table_branch(CreateTableBranchRequest {
15523 id: Some(table_id.clone()),
15524 name: "dev".to_string(),
15525 ..Default::default()
15526 })
15527 .await
15528 .unwrap();
15529 namespace
15530 .create_table_branch(CreateTableBranchRequest {
15531 id: Some(table_id.clone()),
15532 name: "staging".to_string(),
15533 ..Default::default()
15534 })
15535 .await
15536 .unwrap();
15537
15538 let resp = namespace
15539 .list_table_branches(ListTableBranchesRequest {
15540 id: Some(table_id.clone()),
15541 ..Default::default()
15542 })
15543 .await
15544 .unwrap();
15545 assert_eq!(
15546 resp.branches.len(),
15547 2,
15548 "expected 2 branches, got: {:?}",
15549 resp.branches
15550 );
15551 assert!(resp.branches.contains_key("dev"));
15552 assert!(resp.branches.contains_key("staging"));
15553 assert!(resp.page_token.is_none());
15554
15555 namespace
15557 .delete_table_branch(DeleteTableBranchRequest {
15558 id: Some(table_id.clone()),
15559 name: "dev".to_string(),
15560 ..Default::default()
15561 })
15562 .await
15563 .unwrap();
15564
15565 let resp = namespace
15566 .list_table_branches(ListTableBranchesRequest {
15567 id: Some(table_id),
15568 ..Default::default()
15569 })
15570 .await
15571 .unwrap();
15572 assert_eq!(resp.branches.len(), 1, "expected 1 branch after delete");
15573 assert!(!resp.branches.contains_key("dev"));
15574 assert!(resp.branches.contains_key("staging"));
15575 }
15576
15577 #[tokio::test]
15578 async fn test_create_branch_from_version() {
15579 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15580
15581 namespace
15583 .create_table_branch(CreateTableBranchRequest {
15584 id: Some(table_id.clone()),
15585 name: "from-v1".to_string(),
15586 from_version: Some(1),
15587 ..Default::default()
15588 })
15589 .await
15590 .unwrap();
15591
15592 let resp = namespace
15593 .list_table_branches(ListTableBranchesRequest {
15594 id: Some(table_id),
15595 ..Default::default()
15596 })
15597 .await
15598 .unwrap();
15599 let branch = resp
15600 .branches
15601 .get("from-v1")
15602 .expect("forked branch should be listed");
15603 assert_eq!(
15604 branch.parent_version, 1,
15605 "branch should fork from version 1"
15606 );
15607 assert!(
15608 branch.parent_branch.is_none(),
15609 "a branch forked from main has no parent branch"
15610 );
15611 }
15612
15613 #[tokio::test]
15618 async fn test_create_branch_from_other_branch() {
15619 use lance::dataset::builder::DatasetBuilder;
15620
15621 let (namespace, _temp_dir) = create_test_namespace().await;
15622 create_scalar_table(&namespace, "users").await; create_branch_with_commits(&namespace, "users", "dev", 1).await;
15625 let main_ds = open_dataset(&namespace, "users").await;
15627 append_scalar_version(main_ds.uri(), 500).await; namespace
15630 .create_table_branch(CreateTableBranchRequest {
15631 id: Some(vec!["users".to_string()]),
15632 name: "child".to_string(),
15633 from_branch: Some("dev".to_string()),
15634 from_version: Some(2),
15635 ..Default::default()
15636 })
15637 .await
15638 .unwrap();
15639
15640 let child_ds = DatasetBuilder::from_uri(main_ds.uri())
15641 .with_branch("child", None)
15642 .load()
15643 .await
15644 .unwrap();
15645 let ids = scan_id_column(&child_ds).await;
15646 assert!(
15647 ids.contains(&100) && ids.contains(&101),
15648 "child must contain dev's appended rows, got: {:?}",
15649 ids
15650 );
15651 assert!(
15652 !ids.contains(&500),
15653 "child must not contain main's diverged rows, got: {:?}",
15654 ids
15655 );
15656
15657 let listed = namespace
15659 .list_table_branches(ListTableBranchesRequest {
15660 id: Some(vec!["users".to_string()]),
15661 ..Default::default()
15662 })
15663 .await
15664 .unwrap();
15665 assert_eq!(
15666 listed
15667 .branches
15668 .get("child")
15669 .unwrap()
15670 .parent_branch
15671 .as_deref(),
15672 Some("dev")
15673 );
15674 }
15675
15676 #[tokio::test]
15677 async fn test_create_existing_branch_conflict() {
15678 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15679
15680 namespace
15681 .create_table_branch(CreateTableBranchRequest {
15682 id: Some(table_id.clone()),
15683 name: "dev".to_string(),
15684 ..Default::default()
15685 })
15686 .await
15687 .unwrap();
15688
15689 let err = namespace
15690 .create_table_branch(CreateTableBranchRequest {
15691 id: Some(table_id),
15692 name: "dev".to_string(),
15693 ..Default::default()
15694 })
15695 .await
15696 .unwrap_err();
15697 assert_eq!(
15698 namespace_code(&err),
15699 Some(ErrorCode::TableBranchAlreadyExists),
15700 "expected TableBranchAlreadyExists, got: {}",
15701 err
15702 );
15703 assert!(
15704 err.to_string().to_lowercase().contains("already exists"),
15705 "expected already-exists message, got: {}",
15706 err
15707 );
15708 }
15709
15710 #[tokio::test]
15711 async fn test_delete_unknown_branch() {
15712 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15713
15714 let err = namespace
15715 .delete_table_branch(DeleteTableBranchRequest {
15716 id: Some(table_id),
15717 name: "does-not-exist".to_string(),
15718 ..Default::default()
15719 })
15720 .await
15721 .unwrap_err();
15722 assert_eq!(
15723 namespace_code(&err),
15724 Some(ErrorCode::TableBranchNotFound),
15725 "expected TableBranchNotFound, got: {}",
15726 err
15727 );
15728 assert!(
15729 err.to_string().to_lowercase().contains("not found"),
15730 "expected not-found message, got: {}",
15731 err
15732 );
15733 }
15734
15735 #[tokio::test]
15736 async fn test_delete_referenced_branch_conflict() {
15737 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15738
15739 namespace
15741 .create_table_branch(CreateTableBranchRequest {
15742 id: Some(table_id.clone()),
15743 name: "parent".to_string(),
15744 ..Default::default()
15745 })
15746 .await
15747 .unwrap();
15748 namespace
15749 .create_table_branch(CreateTableBranchRequest {
15750 id: Some(table_id.clone()),
15751 name: "child".to_string(),
15752 from_branch: Some("parent".to_string()),
15753 ..Default::default()
15754 })
15755 .await
15756 .unwrap();
15757
15758 let listed = namespace
15760 .list_table_branches(ListTableBranchesRequest {
15761 id: Some(table_id.clone()),
15762 ..Default::default()
15763 })
15764 .await
15765 .unwrap();
15766 let child = listed
15767 .branches
15768 .get("child")
15769 .expect("child branch should be listed");
15770 assert_eq!(
15771 child.parent_branch.as_deref(),
15772 Some("parent"),
15773 "child should record parent branch as its fork point"
15774 );
15775 assert!(
15776 child.parent_version >= 1,
15777 "child should record the parent version it forked from, got {}",
15778 child.parent_version
15779 );
15780
15781 let err = namespace
15784 .delete_table_branch(DeleteTableBranchRequest {
15785 id: Some(table_id),
15786 name: "parent".to_string(),
15787 ..Default::default()
15788 })
15789 .await
15790 .unwrap_err();
15791 assert_eq!(
15792 namespace_code(&err),
15793 Some(ErrorCode::InvalidInput),
15794 "expected InvalidInput for deleting a referenced branch, got: {}",
15795 err
15796 );
15797 assert!(
15798 err.to_string().to_lowercase().contains("referenced"),
15799 "error should explain the branch is still referenced, got: {}",
15800 err
15801 );
15802 }
15803
15804 #[tokio::test]
15805 async fn test_branch_name_required() {
15806 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15807
15808 let create_err = namespace
15809 .create_table_branch(CreateTableBranchRequest {
15810 id: Some(table_id.clone()),
15811 name: String::new(),
15812 ..Default::default()
15813 })
15814 .await
15815 .unwrap_err();
15816 assert_eq!(
15817 namespace_code(&create_err),
15818 Some(ErrorCode::InvalidInput),
15819 "empty name on create should be InvalidInput, got: {}",
15820 create_err
15821 );
15822 assert!(
15823 create_err
15824 .to_string()
15825 .to_lowercase()
15826 .contains("must not be empty")
15827 );
15828
15829 let delete_err = namespace
15830 .delete_table_branch(DeleteTableBranchRequest {
15831 id: Some(table_id),
15832 name: String::new(),
15833 ..Default::default()
15834 })
15835 .await
15836 .unwrap_err();
15837 assert_eq!(
15838 namespace_code(&delete_err),
15839 Some(ErrorCode::InvalidInput),
15840 "empty name on delete should be InvalidInput, got: {}",
15841 delete_err
15842 );
15843 assert!(
15844 delete_err
15845 .to_string()
15846 .to_lowercase()
15847 .contains("must not be empty")
15848 );
15849 }
15850
15851 #[tokio::test]
15852 async fn test_create_branch_rejects_negative_from_version() {
15853 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15854
15855 let err = namespace
15856 .create_table_branch(CreateTableBranchRequest {
15857 id: Some(table_id),
15858 name: "dev".to_string(),
15859 from_version: Some(-1),
15860 ..Default::default()
15861 })
15862 .await
15863 .unwrap_err();
15864 assert_eq!(
15865 namespace_code(&err),
15866 Some(ErrorCode::InvalidInput),
15867 "negative from_version should be InvalidInput, got: {}",
15868 err
15869 );
15870 assert!(err.to_string().to_lowercase().contains("from_version"));
15871 }
15872
15873 #[tokio::test]
15874 async fn test_create_branch_nonexistent_from_version() {
15875 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15876
15877 let err = namespace
15880 .create_table_branch(CreateTableBranchRequest {
15881 id: Some(table_id),
15882 name: "dev".to_string(),
15883 from_version: Some(999),
15884 ..Default::default()
15885 })
15886 .await
15887 .unwrap_err();
15888 assert_eq!(
15889 namespace_code(&err),
15890 Some(ErrorCode::InvalidInput),
15891 "non-existent from_version should map to InvalidInput, got: {}",
15892 err
15893 );
15894 assert!(
15895 err.to_string().to_lowercase().contains("does not exist"),
15896 "error should name the missing source, got: {}",
15897 err
15898 );
15899 }
15900
15901 #[tokio::test]
15902 async fn test_create_and_list_tags() {
15903 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15904
15905 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15906 req.id = Some(table_id.clone());
15907 namespace.create_table_tag(req).await.unwrap();
15908
15909 let mut req = CreateTableTagRequest::new("v2".to_string(), 2);
15910 req.id = Some(table_id.clone());
15911 namespace.create_table_tag(req).await.unwrap();
15912
15913 let mut list_req = ListTableTagsRequest::new();
15914 list_req.id = Some(table_id);
15915 let resp = namespace.list_table_tags(list_req).await.unwrap();
15916
15917 assert_eq!(resp.tags.len(), 2, "expected 2 tags, got: {:?}", resp.tags);
15918 assert_eq!(resp.tags.get("v1").unwrap().version, 1);
15919 assert_eq!(resp.tags.get("v2").unwrap().version, 2);
15920 assert!(resp.page_token.is_none());
15921 }
15922
15923 #[tokio::test]
15924 async fn test_create_existing_tag_conflict() {
15925 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15926
15927 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
15928 req.id = Some(table_id.clone());
15929 namespace.create_table_tag(req).await.unwrap();
15930
15931 let mut req = CreateTableTagRequest::new("v1".to_string(), 2);
15932 req.id = Some(table_id);
15933 let err = namespace.create_table_tag(req).await.unwrap_err();
15934 assert!(
15935 err.to_string().to_lowercase().contains("already exists"),
15936 "expected already-exists error, got: {}",
15937 err
15938 );
15939 }
15940
15941 #[tokio::test]
15942 async fn test_get_tag_version() {
15943 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15944
15945 let mut req = CreateTableTagRequest::new("release".to_string(), 2);
15946 req.id = Some(table_id.clone());
15947 namespace.create_table_tag(req).await.unwrap();
15948
15949 let mut get_req = GetTableTagVersionRequest::new("release".to_string());
15950 get_req.id = Some(table_id);
15951 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15952 assert_eq!(resp.version, 2);
15953 assert_eq!(resp.branch, None);
15954 }
15955
15956 #[tokio::test]
15957 async fn test_get_unknown_tag() {
15958 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15959
15960 let mut get_req = GetTableTagVersionRequest::new("does-not-exist".to_string());
15961 get_req.id = Some(table_id);
15962 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
15963 assert!(
15964 err.to_string().to_lowercase().contains("not found"),
15965 "expected not-found error, got: {}",
15966 err
15967 );
15968 }
15969
15970 #[tokio::test]
15971 async fn test_update_tag_to_new_version() {
15972 let (namespace, _temp_dir, table_id) = create_tagged_test_table(3).await;
15973
15974 let mut req = CreateTableTagRequest::new("rolling".to_string(), 1);
15975 req.id = Some(table_id.clone());
15976 namespace.create_table_tag(req).await.unwrap();
15977
15978 let mut update_req = UpdateTableTagRequest::new("rolling".to_string(), 3);
15979 update_req.id = Some(table_id.clone());
15980 namespace.update_table_tag(update_req).await.unwrap();
15981
15982 let mut get_req = GetTableTagVersionRequest::new("rolling".to_string());
15983 get_req.id = Some(table_id);
15984 let resp = namespace.get_table_tag_version(get_req).await.unwrap();
15985 assert_eq!(resp.version, 3);
15986 }
15987
15988 #[tokio::test]
15989 async fn test_update_unknown_tag() {
15990 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
15991
15992 let mut update_req = UpdateTableTagRequest::new("ghost".to_string(), 1);
15993 update_req.id = Some(table_id);
15994 let err = namespace.update_table_tag(update_req).await.unwrap_err();
15995 assert!(
15996 err.to_string().to_lowercase().contains("not found"),
15997 "expected not-found error, got: {}",
15998 err
15999 );
16000 }
16001
16002 #[tokio::test]
16003 async fn test_delete_tag() {
16004 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
16005
16006 let mut req = CreateTableTagRequest::new("doomed".to_string(), 1);
16007 req.id = Some(table_id.clone());
16008 namespace.create_table_tag(req).await.unwrap();
16009
16010 let mut delete_req = DeleteTableTagRequest::new("doomed".to_string());
16011 delete_req.id = Some(table_id.clone());
16012 namespace.delete_table_tag(delete_req).await.unwrap();
16013
16014 let mut list_req = ListTableTagsRequest::new();
16015 list_req.id = Some(table_id.clone());
16016 let resp = namespace.list_table_tags(list_req).await.unwrap();
16017 assert!(resp.tags.is_empty(), "tag should be removed after delete");
16018
16019 let mut get_req = GetTableTagVersionRequest::new("doomed".to_string());
16021 get_req.id = Some(table_id);
16022 let err = namespace.get_table_tag_version(get_req).await.unwrap_err();
16023 assert!(err.to_string().to_lowercase().contains("not found"));
16024 }
16025
16026 #[tokio::test]
16027 async fn test_delete_unknown_tag() {
16028 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
16029
16030 let mut delete_req = DeleteTableTagRequest::new("nope".to_string());
16031 delete_req.id = Some(table_id);
16032 let err = namespace.delete_table_tag(delete_req).await.unwrap_err();
16033 assert!(
16034 err.to_string().to_lowercase().contains("not found"),
16035 "expected not-found error, got: {}",
16036 err
16037 );
16038 }
16039
16040 #[tokio::test]
16041 async fn test_create_tag_invalid_version() {
16042 let (namespace, _temp_dir, table_id) = create_tagged_test_table(2).await;
16043
16044 let mut req = CreateTableTagRequest::new("v0".to_string(), 0);
16046 req.id = Some(table_id.clone());
16047 let err = namespace.create_table_tag(req).await.unwrap_err();
16048 assert!(
16049 err.to_string().to_lowercase().contains("positive"),
16050 "expected positive-version error, got: {}",
16051 err
16052 );
16053
16054 let mut req = CreateTableTagRequest::new(String::new(), 1);
16056 req.id = Some(table_id);
16057 let err = namespace.create_table_tag(req).await.unwrap_err();
16058 assert!(
16059 err.to_string().to_lowercase().contains("must not be empty"),
16060 "expected empty-tag-name error, got: {}",
16061 err
16062 );
16063 }
16064
16065 #[tokio::test]
16066 async fn test_create_tag_table_not_found() {
16067 let (namespace, _temp_dir) = create_test_namespace().await;
16068
16069 let mut req = CreateTableTagRequest::new("v1".to_string(), 1);
16070 req.id = Some(vec!["does_not_exist".to_string()]);
16071 let err = namespace.create_table_tag(req).await.unwrap_err();
16072 let msg = err.to_string();
16073 assert!(
16074 msg.contains("Table") && msg.to_lowercase().contains("not found"),
16075 "expected TableNotFound error, got: {}",
16076 err
16077 );
16078 }
16079 #[tokio::test]
16080 async fn test_alter_table_drop_columns_missing_id() {
16081 use lance_namespace::models::AlterTableDropColumnsRequest;
16082
16083 let (namespace, _temp_dir) = create_test_namespace().await;
16084
16085 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
16086 let result = namespace.alter_table_drop_columns(request).await;
16087 assert!(result.is_err(), "Should fail when table ID is missing");
16088 }
16089
16090 #[tokio::test]
16091 async fn test_alter_table_drop_columns_nonexistent_table() {
16092 use lance_namespace::models::AlterTableDropColumnsRequest;
16093
16094 let (namespace, _temp_dir) = create_test_namespace().await;
16095
16096 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
16097 request.id = Some(vec!["nonexistent".to_string()]);
16098 let result = namespace.alter_table_drop_columns(request).await;
16099 assert!(result.is_err(), "Should fail when table does not exist");
16100 }
16101
16102 #[tokio::test]
16103 async fn test_create_branch_on_managed_dataset_succeeds() {
16104 use lance::dataset::builder::DatasetBuilder;
16105
16106 let temp = TempStdDir::default();
16107 let ns = create_managed_namespace(temp.to_str().unwrap()).await;
16108 let table_id = vec!["t".to_string()];
16109 let mut main = create_managed_table(&ns, &table_id).await;
16110
16111 let fork_version = main.version().version;
16112 let branch = main
16113 .create_branch("exp", fork_version, None)
16114 .await
16115 .expect("create_branch failed");
16116 assert_eq!(branch.manifest.branch.as_deref(), Some("exp"));
16117 assert_eq!(scan_id_column(&branch).await, vec![1, 2]);
16118
16119 let reopened = DatasetBuilder::from_namespace(ns.clone(), table_id.clone())
16120 .await
16121 .unwrap()
16122 .with_branch("exp", None)
16123 .load()
16124 .await
16125 .expect("reopen branch failed");
16126 assert_eq!(scan_id_column(&reopened).await, vec![1, 2]);
16127 }
16128
16129 #[tokio::test]
16130 async fn test_alter_transaction_set_status() {
16131 use lance_namespace::models::{
16132 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16133 DescribeTransactionRequest,
16134 };
16135
16136 let (namespace, _temp_dir) = create_test_namespace().await;
16137 create_scalar_table(&namespace, "users").await;
16138 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16139 .await
16140 .expect("create_scalar_index should return a transaction id");
16141
16142 let describe_resp = namespace
16144 .describe_transaction(DescribeTransactionRequest {
16145 id: Some(vec!["users".to_string(), txn_id.clone()]),
16146 ..Default::default()
16147 })
16148 .await
16149 .unwrap();
16150 assert_eq!(describe_resp.status, "SUCCEEDED");
16151
16152 let response = namespace
16154 .alter_transaction(AlterTransactionRequest {
16155 id: Some(vec!["users".to_string(), txn_id.clone()]),
16156 actions: vec![AlterTransactionAction {
16157 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16158 status: Some("Canceled".to_string()),
16159 })),
16160 set_property_action: None,
16161 unset_property_action: None,
16162 }],
16163 ..Default::default()
16164 })
16165 .await
16166 .unwrap();
16167 assert_eq!(response.status, "Canceled");
16168 assert!(response.properties.is_some());
16169 let props = response.properties.unwrap();
16170 assert_eq!(props.get("uuid"), Some(&txn_id));
16171 assert_eq!(props.get("operation"), Some(&"CreateIndex".to_string()));
16172 }
16173
16174 #[tokio::test]
16175 async fn test_alter_transaction_set_property() {
16176 use lance_namespace::models::{
16177 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16178 };
16179
16180 let (namespace, _temp_dir) = create_test_namespace().await;
16181 create_scalar_table(&namespace, "users").await;
16182 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16183 .await
16184 .expect("create_scalar_index should return a transaction id");
16185
16186 let response = namespace
16187 .alter_transaction(AlterTransactionRequest {
16188 id: Some(vec!["users".to_string(), txn_id.clone()]),
16189 actions: vec![AlterTransactionAction {
16190 set_status_action: None,
16191 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16192 key: Some("custom_key".to_string()),
16193 value: Some("custom_value".to_string()),
16194 mode: None,
16195 })),
16196 unset_property_action: None,
16197 }],
16198 ..Default::default()
16199 })
16200 .await
16201 .unwrap();
16202 assert_eq!(response.status, "SUCCEEDED");
16203 let props = response.properties.unwrap();
16204 assert_eq!(props.get("custom_key"), Some(&"custom_value".to_string()));
16205 }
16206
16207 #[tokio::test]
16208 async fn test_alter_transaction_set_property_fail_mode() {
16209 use lance_namespace::models::{
16210 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16211 };
16212
16213 let (namespace, _temp_dir) = create_test_namespace().await;
16214 create_scalar_table(&namespace, "users").await;
16215 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16216 .await
16217 .expect("create_scalar_index should return a transaction id");
16218
16219 namespace
16221 .alter_transaction(AlterTransactionRequest {
16222 id: Some(vec!["users".to_string(), txn_id.clone()]),
16223 actions: vec![AlterTransactionAction {
16224 set_status_action: None,
16225 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16226 key: Some("custom_key".to_string()),
16227 value: Some("initial_value".to_string()),
16228 mode: None,
16229 })),
16230 unset_property_action: None,
16231 }],
16232 ..Default::default()
16233 })
16234 .await
16235 .unwrap();
16236
16237 let result = namespace
16240 .alter_transaction(AlterTransactionRequest {
16241 id: Some(vec!["users".to_string(), txn_id.clone()]),
16242 actions: vec![AlterTransactionAction {
16243 set_status_action: None,
16244 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16245 key: Some("custom_key".to_string()),
16246 value: Some("new_value".to_string()),
16247 mode: Some("Fail".to_string()),
16248 })),
16249 unset_property_action: None,
16250 }],
16251 ..Default::default()
16252 })
16253 .await;
16254 assert!(result.is_err());
16255 }
16256
16257 #[tokio::test]
16258 async fn test_alter_transaction_unset_property() {
16259 use lance_namespace::models::{
16260 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16261 AlterTransactionUnsetProperty,
16262 };
16263
16264 let (namespace, _temp_dir) = create_test_namespace().await;
16265 create_scalar_table(&namespace, "users").await;
16266 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16267 .await
16268 .expect("create_scalar_index should return a transaction id");
16269
16270 let response = namespace
16272 .alter_transaction(AlterTransactionRequest {
16273 id: Some(vec!["users".to_string(), txn_id.clone()]),
16274 actions: vec![
16275 AlterTransactionAction {
16276 set_status_action: None,
16277 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16278 key: Some("temp_key".to_string()),
16279 value: Some("temp_value".to_string()),
16280 mode: None,
16281 })),
16282 unset_property_action: None,
16283 },
16284 AlterTransactionAction {
16285 set_status_action: None,
16286 set_property_action: None,
16287 unset_property_action: Some(Box::new(AlterTransactionUnsetProperty {
16288 key: Some("temp_key".to_string()),
16289 mode: None,
16290 })),
16291 },
16292 ],
16293 ..Default::default()
16294 })
16295 .await
16296 .unwrap();
16297 assert_eq!(response.status, "SUCCEEDED");
16298 let props = response.properties.unwrap();
16299 assert!(!props.contains_key("temp_key"));
16300 }
16301
16302 #[tokio::test]
16303 async fn test_alter_transaction_invalid_status() {
16304 use lance_namespace::models::{
16305 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16306 };
16307
16308 let (namespace, _temp_dir) = create_test_namespace().await;
16309 create_scalar_table(&namespace, "users").await;
16310 let txn_id = create_scalar_index(&namespace, "users", "users_id_idx")
16311 .await
16312 .expect("create_scalar_index should return a transaction id");
16313
16314 let result = namespace
16315 .alter_transaction(AlterTransactionRequest {
16316 id: Some(vec!["users".to_string(), txn_id.clone()]),
16317 actions: vec![AlterTransactionAction {
16318 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16319 status: Some("InvalidStatus".to_string()),
16320 })),
16321 set_property_action: None,
16322 unset_property_action: None,
16323 }],
16324 ..Default::default()
16325 })
16326 .await;
16327 assert!(result.is_err());
16328 }
16329
16330 #[tokio::test]
16331 async fn test_alter_transaction_not_found() {
16332 use lance_namespace::models::{
16333 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16334 };
16335
16336 let (namespace, _temp_dir) = create_test_namespace().await;
16337 create_scalar_table(&namespace, "users").await;
16338
16339 let result = namespace
16341 .alter_transaction(AlterTransactionRequest {
16342 id: Some(vec!["users".to_string(), "non_existent_txn".to_string()]),
16343 actions: vec![AlterTransactionAction {
16344 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16345 status: Some("Canceled".to_string()),
16346 })),
16347 set_property_action: None,
16348 unset_property_action: None,
16349 }],
16350 ..Default::default()
16351 })
16352 .await;
16353 assert!(result.is_err());
16354 }
16355
16356 #[tokio::test]
16357 async fn test_alter_transaction_missing_id() {
16358 use lance_namespace::models::{
16359 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetStatus,
16360 };
16361
16362 let (namespace, _temp_dir) = create_test_namespace().await;
16363
16364 let result = namespace
16366 .alter_transaction(AlterTransactionRequest {
16367 id: None,
16368 actions: vec![AlterTransactionAction {
16369 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16370 status: Some("Canceled".to_string()),
16371 })),
16372 set_property_action: None,
16373 unset_property_action: None,
16374 }],
16375 ..Default::default()
16376 })
16377 .await;
16378 assert!(result.is_err());
16379
16380 let result = namespace
16382 .alter_transaction(AlterTransactionRequest {
16383 id: Some(vec!["users".to_string()]),
16384 actions: vec![AlterTransactionAction {
16385 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16386 status: Some("Canceled".to_string()),
16387 })),
16388 set_property_action: None,
16389 unset_property_action: None,
16390 }],
16391 ..Default::default()
16392 })
16393 .await;
16394 assert!(result.is_err());
16395 }
16396
16397 #[tokio::test]
16398 async fn test_alter_transaction_persists_changes() {
16399 use lance_namespace::models::{
16400 AlterTransactionAction, AlterTransactionRequest, AlterTransactionSetProperty,
16401 AlterTransactionSetStatus, DescribeTransactionRequest,
16402 };
16403
16404 let (namespace, _temp_dir) = create_test_namespace().await;
16405 create_scalar_table(&namespace, "users").await;
16406 let transaction_id = create_scalar_index(&namespace, "users", "users_id_idx").await;
16407
16408 let txn_id = transaction_id.expect("scalar index should produce a transaction id");
16409
16410 namespace
16412 .alter_transaction(AlterTransactionRequest {
16413 id: Some(vec!["users".to_string(), txn_id.clone()]),
16414 actions: vec![
16415 AlterTransactionAction {
16416 set_status_action: Some(Box::new(AlterTransactionSetStatus {
16417 status: Some("Canceled".to_string()),
16418 })),
16419 set_property_action: None,
16420 unset_property_action: None,
16421 },
16422 AlterTransactionAction {
16423 set_status_action: None,
16424 set_property_action: Some(Box::new(AlterTransactionSetProperty {
16425 key: Some("owner".to_string()),
16426 value: Some("alice".to_string()),
16427 mode: None,
16428 })),
16429 unset_property_action: None,
16430 },
16431 ],
16432 ..Default::default()
16433 })
16434 .await
16435 .unwrap();
16436
16437 let describe_resp = namespace
16440 .describe_transaction(DescribeTransactionRequest {
16441 id: Some(vec!["users".to_string(), txn_id.clone()]),
16442 ..Default::default()
16443 })
16444 .await
16445 .unwrap();
16446 let props = describe_resp.properties.expect("properties should be set");
16447 assert_eq!(props.get("owner"), Some(&"alice".to_string()));
16448 assert!(!props.contains_key("_status"));
16452
16453 let follow_up = namespace
16454 .alter_transaction(AlterTransactionRequest {
16455 id: Some(vec!["users".to_string(), txn_id.clone()]),
16456 actions: vec![],
16457 ..Default::default()
16458 })
16459 .await
16460 .unwrap();
16461 assert_eq!(follow_up.status, "Canceled");
16462 let follow_up_props = follow_up.properties.unwrap();
16463 assert_eq!(follow_up_props.get("owner"), Some(&"alice".to_string()));
16464 }
16465}