1use super::manifest_feature_flags::{ensure_readable, ensure_writable};
10use arrow::array::builder::{ListBuilder, StringBuilder};
11use arrow::array::{Array, ListArray, RecordBatch, RecordBatchIterator, StringArray, UInt64Array};
12use arrow::datatypes::{DataType, Field, Schema as ArrowSchema, SchemaRef};
13use arrow_ipc::reader::StreamReader;
14use async_trait::async_trait;
15use bytes::Bytes;
16use datafusion_common::DataFusionError;
17use datafusion_physical_plan::{
18 SendableRecordBatchStream,
19 stream::RecordBatchStreamAdapter as DatafusionRecordBatchStreamAdapter,
20};
21use futures::{
22 FutureExt, TryStreamExt,
23 stream::{self, StreamExt},
24};
25use lance::dataset::index::LanceIndexStoreExt;
26use lance::dataset::transaction::{Operation, Transaction};
27use lance::dataset::{
28 InsertBuilder, ReadParams, WhenMatched, WriteMode, WriteParams, builder::DatasetBuilder,
29};
30use lance::session::Session;
31use lance::{Dataset, dataset::scanner::Scanner};
32use lance_core::Error as LanceError;
33use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION;
34use lance_core::{Error, ROW_ID, Result, box_error};
35use lance_index::progress::noop_progress;
36use lance_index::registry::IndexPluginRegistry;
37use lance_index::scalar::lance_format::LanceIndexStore;
38use lance_index::scalar::registry::VALUE_COLUMN_NAME;
39use lance_index::scalar::{
40 BuiltinIndexType, CreatedIndex, ScalarIndexParams, index_files_to_table,
41};
42use lance_io::object_store::{ObjectStore, ObjectStoreParams};
43use lance_io::stream::RecordBatchStream as LanceRecordBatchStream;
44use lance_namespace::LanceNamespace;
45use lance_namespace::error::NamespaceError;
46use lance_namespace::models::{
47 AlterTableAddColumnsRequest, AlterTableAddColumnsResponse, AlterTableAlterColumnsRequest,
48 AlterTableAlterColumnsResponse, AlterTableDropColumnsRequest, AlterTableDropColumnsResponse,
49 CreateNamespaceRequest, CreateNamespaceResponse, CreateTableRequest, CreateTableResponse,
50 DeclareTableRequest, DeclareTableResponse, DeregisterTableRequest, DeregisterTableResponse,
51 DescribeNamespaceRequest, DescribeNamespaceResponse, DescribeTableRequest,
52 DescribeTableResponse, DropNamespaceRequest, DropNamespaceResponse, DropTableRequest,
53 DropTableResponse, ListNamespacesRequest, ListNamespacesResponse, ListTablesRequest,
54 ListTablesResponse, NamespaceExistsRequest, RegisterTableRequest, RegisterTableResponse,
55 TableExistsRequest,
56};
57use lance_namespace::schema::arrow_schema_to_json;
58use lance_table::feature_flags::apply_feature_flags;
59use lance_table::format::{Fragment, IndexMetadata, Manifest};
60use lance_table::io::commit::{
61 CommitError, CommitHandler, commit_handler_from_url, write_manifest_file_to_path,
62};
63use object_store::{Error as ObjectStoreError, path::Path};
64use roaring::RoaringBitmap;
65use std::io::Cursor;
66use std::time::{SystemTime, UNIX_EPOCH};
67use std::{
68 collections::{BTreeMap, HashMap, HashSet},
69 hash::{DefaultHasher, Hash, Hasher},
70 ops::{Deref, DerefMut},
71 sync::{Arc, LazyLock, Mutex as StdMutex, MutexGuard as StdMutexGuard},
72};
73use tokio::sync::{Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard};
74use uuid::Uuid;
75
76const MANIFEST_TABLE_NAME: &str = "__manifest";
77const LANCE_DATA_DIR: &str = "data";
78const LANCE_INDICES_DIR: &str = "_indices";
79const DELIMITER: &str = "$";
80pub(crate) const DECLARED_FILTER_CONCURRENCY: usize = 16;
83
84const OBJECT_ID_INDEX_NAME: &str = "object_id_btree";
87const OBJECT_TYPE_INDEX_NAME: &str = "object_type_bitmap";
89const BASE_OBJECTS_INDEX_NAME: &str = "base_objects_label_list";
91static BASE_OBJECTS_VALUE_FIELD: LazyLock<Field> = LazyLock::new(|| {
94 Field::new(
95 VALUE_COLUMN_NAME,
96 DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))),
97 true,
98 )
99});
100const DEFAULT_MANIFEST_REWRITE_COMMIT_RETRIES: u32 = 20;
103const MANIFEST_INDEX_BATCH_SIZE: usize = 8192;
104
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
107pub enum ObjectType {
108 Namespace,
109 Table,
110}
111
112impl ObjectType {
113 pub fn as_str(&self) -> &'static str {
114 match self {
115 Self::Namespace => "namespace",
116 Self::Table => "table",
117 }
118 }
119
120 pub fn parse(s: &str) -> Result<Self> {
121 match s {
122 "namespace" => Ok(Self::Namespace),
123 "table" => Ok(Self::Table),
124 _ => Err(NamespaceError::Internal {
125 message: format!("Invalid object type: {}", s),
126 }
127 .into()),
128 }
129 }
130}
131
132#[derive(Debug, Clone, Copy, PartialEq, Eq)]
133enum CreateTableMode {
134 Create,
135 ExistOk,
136 Overwrite,
137}
138
139impl CreateTableMode {
140 fn parse(mode: Option<&str>) -> Result<Self> {
141 match mode {
142 None => Ok(Self::Create),
143 Some(mode) if mode.eq_ignore_ascii_case("create") => Ok(Self::Create),
144 Some(mode)
145 if mode.eq_ignore_ascii_case("existok")
146 || mode.eq_ignore_ascii_case("exist_ok") =>
147 {
148 Ok(Self::ExistOk)
149 }
150 Some(mode) if mode.eq_ignore_ascii_case("overwrite") => Ok(Self::Overwrite),
151 Some(mode) => Err(NamespaceError::InvalidInput {
152 message: format!(
153 "Unsupported create_table mode '{}'. Supported modes are: 'Create', 'ExistOk', 'Overwrite'",
154 mode
155 ),
156 }
157 .into()),
158 }
159 }
160
161 fn write_mode(self) -> WriteMode {
162 match self {
163 Self::Overwrite => WriteMode::Overwrite,
164 Self::Create | Self::ExistOk => WriteMode::Create,
165 }
166 }
167}
168
169#[derive(Debug, Clone)]
171pub struct TableInfo {
172 pub namespace: Vec<String>,
173 pub name: String,
174 pub location: String,
175 pub metadata: Option<HashMap<String, String>>,
176}
177
178#[derive(Debug, Clone)]
183pub struct ManifestEntry {
184 pub object_id: String,
186 pub object_type: ObjectType,
188 pub location: Option<String>,
190 pub metadata: Option<String>,
192}
193
194struct CopyOnWriteMutation<T> {
195 result: T,
196 has_changes: bool,
197}
198
199impl<T> CopyOnWriteMutation<T> {
200 fn updated(result: T) -> Self {
201 Self {
202 result,
203 has_changes: true,
204 }
205 }
206
207 fn unchanged(result: T) -> Self {
208 Self {
209 result,
210 has_changes: false,
211 }
212 }
213}
214
215struct ManifestIndexBuildInput {
216 index_name: &'static str,
217 column_name: &'static str,
218 params: ScalarIndexParams,
219 field: Field,
220 stream: SendableRecordBatchStream,
221}
222
223struct ManifestTrainedIndex {
224 index_name: &'static str,
225 column_name: &'static str,
226 uuid: Uuid,
227 created_index: CreatedIndex,
228}
229
230struct ManifestRowValue {
231 object_id: String,
232 object_type: ObjectType,
233 location: Option<String>,
234 metadata: Option<String>,
235 base_objects: Option<Vec<String>>,
236}
237
238struct ManifestOutputRow<'a> {
239 object_id: &'a str,
240 object_type: ObjectType,
241 location: Option<&'a str>,
242 metadata: Option<&'a str>,
243 base_objects: Option<&'a [String]>,
244}
245
246#[derive(Default)]
247struct ManifestIndexAccumulator {
248 object_ids: BTreeMap<Arc<str>, u64>,
249 object_types: BTreeMap<&'static str, RoaringBitmap>,
250 base_objects_values: Vec<Option<Vec<String>>>,
251 base_objects_row_ids: Vec<u64>,
252 row_count: u64,
253}
254
255impl ManifestIndexAccumulator {
256 fn next_row_id(&self) -> Result<u64> {
257 if self.row_count >= u64::from(u32::MAX) {
258 return Err(NamespaceError::Internal {
259 message: format!(
260 "Manifest rewrite exceeded maximum single-fragment row count: {}",
261 self.row_count
262 ),
263 }
264 .into());
265 }
266 Ok(self.row_count)
267 }
268
269 fn push(&mut self, row: &ManifestOutputRow<'_>) -> Result<u64> {
270 let row_id = self.next_row_id()?;
271 if self
272 .object_ids
273 .insert(Arc::<str>::from(row.object_id), row_id)
274 .is_some()
275 {
276 return Err(NamespaceError::Internal {
277 message: format!("Manifest contains duplicate object_id '{}'", row.object_id),
278 }
279 .into());
280 }
281 self.object_types
282 .entry(row.object_type.as_str())
283 .or_default()
284 .insert(row_id as u32);
285 self.base_objects_values
286 .push(row.base_objects.map(|objects| objects.to_vec()));
287 self.base_objects_row_ids.push(row_id);
288 self.row_count += 1;
289 Ok(row_id)
290 }
291}
292
293struct ManifestBatchBuilder {
294 object_ids: Vec<String>,
295 object_types: Vec<&'static str>,
296 locations: Vec<Option<String>>,
297 metadatas: Vec<Option<String>>,
298 base_objects: Vec<Option<Vec<String>>>,
299}
300
301impl ManifestBatchBuilder {
302 fn new() -> Self {
303 Self {
304 object_ids: Vec::new(),
305 object_types: Vec::new(),
306 locations: Vec::new(),
307 metadatas: Vec::new(),
308 base_objects: Vec::new(),
309 }
310 }
311
312 fn is_empty(&self) -> bool {
313 self.object_ids.is_empty()
314 }
315
316 fn append(
317 &mut self,
318 index_data: &mut ManifestIndexAccumulator,
319 row: ManifestOutputRow<'_>,
320 ) -> Result<()> {
321 index_data.push(&row)?;
322 self.object_ids.push(row.object_id.to_string());
323 self.object_types.push(row.object_type.as_str());
324 self.locations.push(row.location.map(ToString::to_string));
325 self.metadatas.push(row.metadata.map(ToString::to_string));
326 self.base_objects
327 .push(row.base_objects.map(|objects| objects.to_vec()));
328 Ok(())
329 }
330
331 fn finish(self) -> Result<RecordBatch> {
332 let base_objects_array = ManifestNamespace::base_objects_array(&self.base_objects);
333 RecordBatch::try_new(
334 ManifestNamespace::manifest_schema(),
335 vec![
336 Arc::new(StringArray::from(self.object_ids)),
337 Arc::new(StringArray::from(self.object_types)),
338 Arc::new(StringArray::from(self.locations)),
339 Arc::new(StringArray::from(self.metadatas)),
340 Arc::new(base_objects_array),
341 ],
342 )
343 .map_err(|e| {
344 lance_core::Error::from(NamespaceError::Internal {
345 message: format!("Failed to create manifest snapshot batch: {:?}", e),
346 })
347 })
348 }
349}
350
351enum ConflictResolution<O> {
354 Retry,
357 FailIfExists(Vec<String>),
361 SucceedIfAbsent { object_id: String, output: O },
364}
365
366trait ManifestStreamMutation: Send {
367 type Output: Clone + Send + 'static;
368
369 fn process_existing_row(
370 &mut self,
371 row: ManifestRowValue,
372 output: &mut ManifestBatchBuilder,
373 index_data: &mut ManifestIndexAccumulator,
374 ) -> Result<()>;
375
376 fn append_rows(
377 &mut self,
378 output: &mut ManifestBatchBuilder,
379 index_data: &mut ManifestIndexAccumulator,
380 ) -> Result<()>;
381
382 fn finish(&self) -> CopyOnWriteMutation<Self::Output>;
383
384 fn conflict_resolution(&self) -> ConflictResolution<Self::Output> {
387 ConflictResolution::Retry
388 }
389}
390
391struct ManifestRewriteShared<M: ManifestStreamMutation> {
392 mutation: M,
393 index_data: Option<ManifestIndexAccumulator>,
394 result: Option<CopyOnWriteMutation<M::Output>>,
395 error: Option<LanceError>,
396}
397
398impl<M: ManifestStreamMutation> ManifestRewriteShared<M> {
399 fn new(mutation: M) -> Self {
400 Self {
401 mutation,
402 index_data: Some(ManifestIndexAccumulator::default()),
403 result: None,
404 error: None,
405 }
406 }
407}
408
409struct UpsertManifestMutation {
410 entries: Vec<ManifestEntry>,
411 base_objects: Vec<Option<Vec<String>>>,
412 entry_positions: HashMap<String, usize>,
413 matched: Vec<bool>,
414 when_matched: WhenMatched,
415}
416
417impl UpsertManifestMutation {
418 fn new(
419 entries: Vec<ManifestEntry>,
420 base_objects: Option<Vec<String>>,
421 when_matched: WhenMatched,
422 ) -> Self {
423 let entry_positions = entries
424 .iter()
425 .enumerate()
426 .map(|(index, entry)| (entry.object_id.clone(), index))
427 .collect();
428 let matched = vec![false; entries.len()];
429 let mut entry_base_objects = vec![None; entries.len()];
430 if !entry_base_objects.is_empty() {
431 entry_base_objects[0] = base_objects;
432 }
433 Self {
434 entries,
435 base_objects: entry_base_objects,
436 entry_positions,
437 matched,
438 when_matched,
439 }
440 }
441
442 fn entry_row(&self, index: usize) -> ManifestOutputRow<'_> {
443 let entry = &self.entries[index];
444 ManifestOutputRow {
445 object_id: &entry.object_id,
446 object_type: entry.object_type,
447 location: entry.location.as_deref(),
448 metadata: entry.metadata.as_deref(),
449 base_objects: self.base_objects[index].as_deref(),
450 }
451 }
452}
453
454impl ManifestStreamMutation for UpsertManifestMutation {
455 type Output = ();
456
457 fn process_existing_row(
458 &mut self,
459 row: ManifestRowValue,
460 output: &mut ManifestBatchBuilder,
461 index_data: &mut ManifestIndexAccumulator,
462 ) -> Result<()> {
463 if let Some(index) = self.entry_positions.get(&row.object_id).copied() {
464 match self.when_matched {
465 WhenMatched::Fail => {
466 return Err(NamespaceError::ConcurrentModification {
467 message: format!(
468 "Object '{}' was concurrently created by another operation",
469 row.object_id
470 ),
471 }
472 .into());
473 }
474 WhenMatched::UpdateAll => {
475 self.matched[index] = true;
476 output.append(index_data, self.entry_row(index))?;
477 return Ok(());
478 }
479 _ => {
480 return Err(NamespaceError::Internal {
481 message: format!(
482 "Unsupported manifest rewrite matched action: {:?}",
483 self.when_matched
484 ),
485 }
486 .into());
487 }
488 }
489 }
490
491 output.append(
492 index_data,
493 ManifestOutputRow {
494 object_id: &row.object_id,
495 object_type: row.object_type,
496 location: row.location.as_deref(),
497 metadata: row.metadata.as_deref(),
498 base_objects: row.base_objects.as_deref(),
499 },
500 )
501 }
502
503 fn append_rows(
504 &mut self,
505 output: &mut ManifestBatchBuilder,
506 index_data: &mut ManifestIndexAccumulator,
507 ) -> Result<()> {
508 for index in 0..self.entries.len() {
509 if !self.matched[index] {
510 output.append(index_data, self.entry_row(index))?;
511 }
512 }
513 Ok(())
514 }
515
516 fn finish(&self) -> CopyOnWriteMutation<Self::Output> {
517 CopyOnWriteMutation::updated(())
518 }
519
520 fn conflict_resolution(&self) -> ConflictResolution<Self::Output> {
521 match self.when_matched {
522 WhenMatched::Fail => ConflictResolution::FailIfExists(
525 self.entries.iter().map(|e| e.object_id.clone()).collect(),
526 ),
527 _ => ConflictResolution::Retry,
529 }
530 }
531}
532
533struct DeleteObjectMutation {
534 object_id: String,
535 deleted: bool,
536}
537
538impl ManifestStreamMutation for DeleteObjectMutation {
539 type Output = ();
540
541 fn process_existing_row(
542 &mut self,
543 row: ManifestRowValue,
544 output: &mut ManifestBatchBuilder,
545 index_data: &mut ManifestIndexAccumulator,
546 ) -> Result<()> {
547 if row.object_id == self.object_id {
548 self.deleted = true;
549 return Ok(());
550 }
551
552 output.append(
553 index_data,
554 ManifestOutputRow {
555 object_id: &row.object_id,
556 object_type: row.object_type,
557 location: row.location.as_deref(),
558 metadata: row.metadata.as_deref(),
559 base_objects: row.base_objects.as_deref(),
560 },
561 )
562 }
563
564 fn append_rows(
565 &mut self,
566 _output: &mut ManifestBatchBuilder,
567 _index_data: &mut ManifestIndexAccumulator,
568 ) -> Result<()> {
569 Ok(())
570 }
571
572 fn finish(&self) -> CopyOnWriteMutation<Self::Output> {
573 if self.deleted {
574 CopyOnWriteMutation::updated(())
575 } else {
576 CopyOnWriteMutation::unchanged(())
577 }
578 }
579
580 fn conflict_resolution(&self) -> ConflictResolution<Self::Output> {
581 ConflictResolution::SucceedIfAbsent {
583 object_id: self.object_id.clone(),
584 output: (),
585 }
586 }
587}
588
589#[derive(Debug, Clone)]
591pub struct NamespaceInfo {
592 pub namespace: Vec<String>,
593 pub name: String,
594 pub metadata: Option<HashMap<String, String>>,
595}
596
597#[derive(Debug, Clone)]
604pub struct DatasetConsistencyWrapper(Arc<RwLock<Dataset>>);
605
606impl DatasetConsistencyWrapper {
607 pub fn new(dataset: Dataset) -> Self {
609 debug_assert!(
610 !dataset
611 .manifest()
612 .config
613 .keys()
614 .any(|key| key.starts_with("lance.auto_cleanup.")),
615 "the directory manifest dataset must not enable old-version cleanup"
616 );
617 Self(Arc::new(RwLock::new(dataset)))
618 }
619
620 pub async fn get(&self) -> Result<DatasetReadGuard<'_>> {
623 self.reload().await?;
624 let guard = DatasetReadGuard {
625 guard: self.0.read().await,
626 };
627 ensure_readable(guard.metadata())?;
630 Ok(guard)
631 }
632
633 pub async fn get_refreshed(&self) -> Result<DatasetReadGuard<'_>> {
635 self.reload().await?;
636 let guard = DatasetReadGuard {
637 guard: self.0.read().await,
638 };
639 ensure_readable(guard.metadata())?;
640 Ok(guard)
641 }
642
643 pub async fn get_mut(&self) -> Result<DatasetWriteGuard<'_>> {
646 self.reload().await?;
647 let guard = DatasetWriteGuard {
648 guard: self.0.write().await,
649 };
650 ensure_readable(guard.metadata())?;
651 ensure_writable(guard.metadata())?;
652 Ok(guard)
653 }
654
655 pub async fn set_latest(&self, dataset: Dataset) {
660 let mut write_guard = self.0.write().await;
661 if dataset.manifest().version > write_guard.manifest().version {
662 *write_guard = dataset;
663 }
664 }
665
666 async fn reload(&self) -> Result<()> {
668 let read_guard = self.0.read().await;
670 let dataset_uri = read_guard.uri().to_string();
671 let current_version = read_guard.version().version;
672 log::debug!(
673 "Reload starting for uri={}, current_version={}",
674 dataset_uri,
675 current_version
676 );
677 let has_successor_version = read_guard.has_successor_version().await.map_err(|e| {
682 lance_core::Error::from(NamespaceError::Internal {
683 message: format!("Failed to check dataset staleness: {:?}", e),
684 })
685 })?;
686 log::debug!(
687 "Reload checked successor_version_exists={} for uri={}, current_version={}",
688 has_successor_version,
689 dataset_uri,
690 current_version
691 );
692 drop(read_guard);
693
694 if !has_successor_version {
696 log::debug!("Already up-to-date for uri={}", dataset_uri);
697 return Ok(());
698 }
699
700 let mut write_guard = self.0.write().await;
702
703 let has_successor_version = write_guard.has_successor_version().await.map_err(|e| {
705 lance_core::Error::from(NamespaceError::Internal {
706 message: format!("Failed to check dataset staleness: {:?}", e),
707 })
708 })?;
709
710 if has_successor_version {
711 write_guard.checkout_latest().await.map_err(|e| {
712 lance_core::Error::from(NamespaceError::Internal {
713 message: format!("Failed to checkout latest: {:?}", e),
714 })
715 })?;
716 }
717
718 Ok(())
719 }
720}
721
722pub struct DatasetReadGuard<'a> {
723 guard: RwLockReadGuard<'a, Dataset>,
724}
725
726impl Deref for DatasetReadGuard<'_> {
727 type Target = Dataset;
728
729 fn deref(&self) -> &Self::Target {
730 &self.guard
731 }
732}
733
734pub struct DatasetWriteGuard<'a> {
735 guard: RwLockWriteGuard<'a, Dataset>,
736}
737
738impl Deref for DatasetWriteGuard<'_> {
739 type Target = Dataset;
740
741 fn deref(&self) -> &Self::Target {
742 &self.guard
743 }
744}
745
746impl DerefMut for DatasetWriteGuard<'_> {
747 fn deref_mut(&mut self) -> &mut Self::Target {
748 &mut self.guard
749 }
750}
751
752pub struct ManifestNamespace {
756 root: String,
757 storage_options: Option<HashMap<String, String>>,
758 session: Option<Arc<Session>>,
759 object_store: Arc<ObjectStore>,
760 base_path: Path,
761 manifest_dataset: DatasetConsistencyWrapper,
762 dir_listing_enabled: bool,
766 inline_optimization_enabled: bool,
769 commit_retries: Option<u32>,
772 manifest_mutation_lock: Arc<Mutex<()>>,
775}
776
777impl std::fmt::Debug for ManifestNamespace {
778 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
779 f.debug_struct("ManifestNamespace")
780 .field("root", &self.root)
781 .field("storage_options", &self.storage_options)
782 .field("dir_listing_enabled", &self.dir_listing_enabled)
783 .field(
784 "inline_optimization_enabled",
785 &self.inline_optimization_enabled,
786 )
787 .finish()
788 }
789}
790
791fn convert_lance_commit_error(e: &LanceError, operation: &str, object_id: Option<&str>) -> Error {
800 match e {
801 LanceError::CommitConflict { .. } => NamespaceError::Throttling {
803 message: format!("Too many concurrent writes, please retry later: {:?}", e),
804 }
805 .into(),
806 LanceError::TooMuchWriteContention { .. } | LanceError::IncompatibleTransaction { .. } => {
809 let message = if let Some(id) = object_id {
810 format!(
811 "Object '{}' was concurrently modified by another operation: {:?}",
812 id, e
813 )
814 } else {
815 format!(
816 "Object was concurrently modified by another operation: {:?}",
817 e
818 )
819 };
820 NamespaceError::ConcurrentModification { message }.into()
821 }
822 _ => {
824 let error_msg = e.to_string();
825 if error_msg.contains("matched")
826 || error_msg.contains("duplicate")
827 || error_msg.contains("already exists")
828 {
829 let message = if let Some(id) = object_id {
830 format!(
831 "Object '{}' was concurrently created by another operation: {:?}",
832 id, e
833 )
834 } else {
835 format!(
836 "Object was concurrently created by another operation: {:?}",
837 e
838 )
839 };
840 return NamespaceError::ConcurrentModification { message }.into();
841 }
842 lance_core::Error::from(NamespaceError::Internal {
843 message: format!("{}: {:?}", operation, e),
844 })
845 }
846 }
847}
848
849impl ManifestNamespace {
850 #[allow(clippy::too_many_arguments)]
852 pub async fn from_directory(
853 root: String,
854 storage_options: Option<HashMap<String, String>>,
855 session: Option<Arc<Session>>,
856 object_store: Arc<ObjectStore>,
857 base_path: Path,
858 dir_listing_enabled: bool,
859 inline_optimization_enabled: bool,
860 commit_retries: Option<u32>,
861 ) -> Result<Self> {
862 let manifest_dataset =
863 Self::ensure_manifest_table_up_to_date(&root, &storage_options, session.clone())
864 .await?;
865
866 Ok(Self::new(
867 root,
868 storage_options,
869 session,
870 object_store,
871 base_path,
872 manifest_dataset,
873 dir_listing_enabled,
874 inline_optimization_enabled,
875 commit_retries,
876 ))
877 }
878
879 #[allow(clippy::too_many_arguments)]
881 pub async fn open_from_directory(
882 root: String,
883 storage_options: Option<HashMap<String, String>>,
884 session: Option<Arc<Session>>,
885 object_store: Arc<ObjectStore>,
886 base_path: Path,
887 dir_listing_enabled: bool,
888 inline_optimization_enabled: bool,
889 commit_retries: Option<u32>,
890 ) -> Result<Self> {
891 let manifest_dataset =
892 Self::open_manifest_table(&root, &storage_options, session.clone()).await?;
893
894 Ok(Self::new(
895 root,
896 storage_options,
897 session,
898 object_store,
899 base_path,
900 manifest_dataset,
901 dir_listing_enabled,
902 inline_optimization_enabled,
903 commit_retries,
904 ))
905 }
906
907 #[allow(clippy::too_many_arguments)]
908 fn new(
909 root: String,
910 storage_options: Option<HashMap<String, String>>,
911 session: Option<Arc<Session>>,
912 object_store: Arc<ObjectStore>,
913 base_path: Path,
914 manifest_dataset: DatasetConsistencyWrapper,
915 dir_listing_enabled: bool,
916 inline_optimization_enabled: bool,
917 commit_retries: Option<u32>,
918 ) -> Self {
919 Self {
920 root,
921 storage_options,
922 session,
923 object_store,
924 base_path,
925 manifest_dataset,
926 dir_listing_enabled,
927 inline_optimization_enabled,
928 commit_retries,
929 manifest_mutation_lock: Arc::new(Mutex::new(())),
930 }
931 }
932
933 pub fn build_object_id(namespace: &[String], name: &str) -> String {
935 if namespace.is_empty() {
936 name.to_string()
937 } else {
938 let mut id = namespace.join(DELIMITER);
939 id.push_str(DELIMITER);
940 id.push_str(name);
941 id
942 }
943 }
944
945 pub fn parse_object_id(object_id: &str) -> (Vec<String>, String) {
947 let parts: Vec<&str> = object_id.split(DELIMITER).collect();
948 if parts.len() == 1 {
949 (Vec::new(), parts[0].to_string())
950 } else {
951 let namespace = parts[..parts.len() - 1]
952 .iter()
953 .map(|s| s.to_string())
954 .collect();
955 let name = parts[parts.len() - 1].to_string();
956 (namespace, name)
957 }
958 }
959
960 pub fn split_object_id(object_id: &[String]) -> (Vec<String>, String) {
962 if object_id.len() == 1 {
963 (vec![], object_id[0].clone())
964 } else {
965 (
966 object_id[..object_id.len() - 1].to_vec(),
967 object_id[object_id.len() - 1].clone(),
968 )
969 }
970 }
971
972 pub fn str_object_id(object_id: &[String]) -> String {
974 object_id.join(DELIMITER)
975 }
976
977 fn format_table_id(table_id: &[String]) -> String {
978 format!("table id '{}'", Self::str_object_id(table_id))
979 }
980
981 pub fn generate_dir_name(object_id: &str) -> String {
988 let random_num: u64 = rand::random();
990
991 let mut hasher = DefaultHasher::new();
993 random_num.hash(&mut hasher);
994 object_id.hash(&mut hasher);
995 let hash = hasher.finish();
996
997 format!("{:08x}_{}", (hash & 0xFFFFFFFF) as u32, object_id)
999 }
1000
1001 pub(crate) fn construct_full_uri(root: &str, relative_location: &str) -> Result<String> {
1003 let mut base_url = lance_io::object_store::uri_to_url(root)?;
1004
1005 if !base_url.path().ends_with('/') {
1010 base_url.set_path(&format!("{}/", base_url.path()));
1011 }
1012
1013 let mut full_url = base_url.clone();
1014 full_url
1015 .path_segments_mut()
1016 .map_err(|_| {
1017 lance_core::Error::from(NamespaceError::InvalidInput {
1018 message: format!("Cannot modify path segments for URI '{}'", root),
1019 })
1020 })?
1021 .pop_if_empty()
1022 .extend(
1023 relative_location
1024 .split('/')
1025 .filter(|segment| !segment.is_empty()),
1026 );
1027
1028 full_url.set_query(None);
1032
1033 Ok(full_url.to_string())
1034 }
1035
1036 fn string_list_array(values: &[Option<Vec<String>>], child_name: &str) -> ListArray {
1037 let string_builder = StringBuilder::new();
1038 let mut list_builder = ListBuilder::new(string_builder).with_field(Arc::new(Field::new(
1039 child_name,
1040 DataType::Utf8,
1041 true,
1042 )));
1043 for value in values {
1044 match value {
1045 Some(objects) => {
1046 for object in objects {
1047 list_builder.values().append_value(object);
1048 }
1049 list_builder.append(true);
1050 }
1051 None => list_builder.append_null(),
1052 }
1053 }
1054 list_builder.finish()
1055 }
1056
1057 fn base_objects_array(values: &[Option<Vec<String>>]) -> ListArray {
1058 Self::string_list_array(values, "object_id")
1059 }
1060
1061 fn value_row_id_schema(value_field: Field) -> SchemaRef {
1062 Arc::new(ArrowSchema::new(vec![
1063 value_field,
1064 Field::new(ROW_ID, DataType::UInt64, false),
1065 ]))
1066 }
1067
1068 fn string_row_id_batch(
1069 schema: SchemaRef,
1070 values: Vec<String>,
1071 row_ids: Vec<u64>,
1072 ) -> Result<RecordBatch> {
1073 RecordBatch::try_new(
1074 schema,
1075 vec![
1076 Arc::new(StringArray::from(values)),
1077 Arc::new(UInt64Array::from(row_ids)),
1078 ],
1079 )
1080 .map_err(Into::into)
1081 }
1082
1083 fn list_row_id_batch(
1084 schema: SchemaRef,
1085 values: Vec<Option<Vec<String>>>,
1086 row_ids: Vec<u64>,
1087 ) -> Result<RecordBatch> {
1088 RecordBatch::try_new(
1089 schema,
1090 vec![
1091 Arc::new(Self::string_list_array(&values, "item")),
1092 Arc::new(UInt64Array::from(row_ids)),
1093 ],
1094 )
1095 .map_err(Into::into)
1096 }
1097
1098 fn object_id_index_stream(object_ids: BTreeMap<Arc<str>, u64>) -> SendableRecordBatchStream {
1099 let schema =
1100 Self::value_row_id_schema(Field::new(VALUE_COLUMN_NAME, DataType::Utf8, false));
1101 let stream_schema = schema.clone();
1102 let stream = stream::unfold(
1103 (object_ids.into_iter(), false, schema),
1104 |(mut iter, emitted, schema)| async move {
1105 let mut values = Vec::with_capacity(MANIFEST_INDEX_BATCH_SIZE);
1106 let mut row_ids = Vec::with_capacity(MANIFEST_INDEX_BATCH_SIZE);
1107 for _ in 0..MANIFEST_INDEX_BATCH_SIZE {
1108 let Some((value, row_id)) = iter.next() else {
1109 break;
1110 };
1111 values.push(value.to_string());
1112 row_ids.push(row_id);
1113 }
1114 if values.is_empty() {
1115 if emitted {
1116 None
1117 } else {
1118 let batch = Self::string_row_id_batch(schema.clone(), values, row_ids)
1119 .map_err(|err| DataFusionError::External(Box::new(err)));
1120 Some((batch, (iter, true, schema)))
1121 }
1122 } else {
1123 let batch = Self::string_row_id_batch(schema.clone(), values, row_ids)
1124 .map_err(|err| DataFusionError::External(Box::new(err)));
1125 Some((batch, (iter, true, schema)))
1126 }
1127 },
1128 );
1129 Box::pin(DatafusionRecordBatchStreamAdapter::new(
1130 stream_schema,
1131 stream.fuse(),
1132 ))
1133 }
1134
1135 fn object_type_index_stream(
1136 object_types: BTreeMap<&'static str, RoaringBitmap>,
1137 ) -> SendableRecordBatchStream {
1138 let schema =
1139 Self::value_row_id_schema(Field::new(VALUE_COLUMN_NAME, DataType::Utf8, false));
1140 let stream_schema = schema.clone();
1141 let entries = object_types
1142 .into_iter()
1143 .map(|(value, bitmap)| {
1144 (
1145 value,
1146 Box::new(bitmap.into_iter()) as Box<dyn Iterator<Item = u32> + Send>,
1147 )
1148 })
1149 .collect::<Vec<_>>()
1150 .into_iter();
1151 let stream = stream::unfold(
1152 (entries, None, false, schema),
1153 |(mut entries, mut current, emitted, schema)| async move {
1154 let mut values = Vec::with_capacity(MANIFEST_INDEX_BATCH_SIZE);
1155 let mut row_ids = Vec::with_capacity(MANIFEST_INDEX_BATCH_SIZE);
1156 while values.len() < MANIFEST_INDEX_BATCH_SIZE {
1157 if current.is_none() {
1158 current = entries.next();
1159 }
1160 let Some((value, iter)) = current.as_mut() else {
1161 break;
1162 };
1163 if let Some(row_id) = iter.next() {
1164 values.push((*value).to_string());
1165 row_ids.push(u64::from(row_id));
1166 } else {
1167 current = None;
1168 }
1169 }
1170
1171 if values.is_empty() {
1172 if emitted {
1173 None
1174 } else {
1175 let batch = Self::string_row_id_batch(schema.clone(), values, row_ids)
1176 .map_err(|err| DataFusionError::External(Box::new(err)));
1177 Some((batch, (entries, current, true, schema)))
1178 }
1179 } else {
1180 let batch = Self::string_row_id_batch(schema.clone(), values, row_ids)
1181 .map_err(|err| DataFusionError::External(Box::new(err)));
1182 Some((batch, (entries, current, true, schema)))
1183 }
1184 },
1185 );
1186 Box::pin(DatafusionRecordBatchStreamAdapter::new(
1187 stream_schema,
1188 stream.fuse(),
1189 ))
1190 }
1191
1192 fn base_objects_index_stream(
1193 base_objects_values: Vec<Option<Vec<String>>>,
1194 base_objects_row_ids: Vec<u64>,
1195 ) -> SendableRecordBatchStream {
1196 let schema = Self::value_row_id_schema(BASE_OBJECTS_VALUE_FIELD.clone());
1197 let stream_schema = schema.clone();
1198 let stream = stream::unfold(
1199 (
1200 base_objects_values.into_iter().zip(base_objects_row_ids),
1201 false,
1202 schema,
1203 ),
1204 |(mut iter, emitted, schema)| async move {
1205 let mut values = Vec::with_capacity(MANIFEST_INDEX_BATCH_SIZE);
1206 let mut row_ids = Vec::with_capacity(MANIFEST_INDEX_BATCH_SIZE);
1207 for _ in 0..MANIFEST_INDEX_BATCH_SIZE {
1208 let Some((value, row_id)) = iter.next() else {
1209 break;
1210 };
1211 values.push(value);
1212 row_ids.push(row_id);
1213 }
1214 if values.is_empty() {
1215 if emitted {
1216 None
1217 } else {
1218 let batch = Self::list_row_id_batch(schema.clone(), values, row_ids)
1219 .map_err(|err| DataFusionError::External(Box::new(err)));
1220 Some((batch, (iter, true, schema)))
1221 }
1222 } else {
1223 let batch = Self::list_row_id_batch(schema.clone(), values, row_ids)
1224 .map_err(|err| DataFusionError::External(Box::new(err)));
1225 Some((batch, (iter, true, schema)))
1226 }
1227 },
1228 );
1229 Box::pin(DatafusionRecordBatchStreamAdapter::new(
1230 stream_schema,
1231 stream.fuse(),
1232 ))
1233 }
1234
1235 async fn train_manifest_index(
1236 dataset: &Dataset,
1237 registry: Arc<IndexPluginRegistry>,
1238 input: ManifestIndexBuildInput,
1239 index_uuid: Uuid,
1240 ) -> Result<ManifestTrainedIndex> {
1241 let index_store = LanceIndexStore::from_dataset_for_new(dataset, &index_uuid)?;
1242 let trainer = registry
1243 .get_plugin_by_name(&input.params.index_type)?
1244 .basic_trainer()
1245 .ok_or_else(|| {
1246 lance_core::Error::invalid_input_source(
1247 format!(
1248 "The '{}' index type does not support basic training, please refer to the index's documentation for more details on how to create this index.",
1249 input.params.index_type
1250 )
1251 .into(),
1252 )
1253 })?;
1254 let training_request = trainer
1255 .new_training_request(input.params.params.as_deref().unwrap_or("{}"), &input.field)?;
1256 let created_index = trainer
1257 .train_index(
1258 input.stream,
1259 &index_store,
1260 training_request,
1261 None,
1262 noop_progress(),
1263 )
1264 .await?;
1265 Ok(ManifestTrainedIndex {
1266 index_name: input.index_name,
1267 column_name: input.column_name,
1268 uuid: index_uuid,
1269 created_index,
1270 })
1271 }
1272
1273 fn manifest_index_metadata(
1274 lance_schema: &lance_core::datatypes::Schema,
1275 fragment_bitmap: &RoaringBitmap,
1276 dataset_version: u64,
1277 trained_index: ManifestTrainedIndex,
1278 ) -> Result<IndexMetadata> {
1279 Ok(IndexMetadata {
1280 uuid: trained_index.uuid,
1281 fields: vec![lance_schema.field_id(trained_index.column_name)?],
1282 covering_fields: vec![],
1283 name: trained_index.index_name.to_string(),
1284 dataset_version,
1285 fragment_bitmap: Some(fragment_bitmap.clone()),
1286 index_details: Some(Arc::new(trained_index.created_index.index_details)),
1287 index_version: trained_index.created_index.index_version as i32,
1288 created_at: None,
1289 base_id: None,
1290 files: Some(index_files_to_table(trained_index.created_index.files)),
1291 })
1292 }
1293
1294 fn manifest_fragment_bitmap(manifest: &Manifest) -> Result<RoaringBitmap> {
1295 let mut bitmap = RoaringBitmap::new();
1296 for fragment in manifest.fragments.iter() {
1297 let fragment_id = u32::try_from(fragment.id).map_err(|_| {
1298 lance_core::Error::from(NamespaceError::Internal {
1299 message: format!("Manifest fragment id {} exceeds u32", fragment.id),
1300 })
1301 })?;
1302 bitmap.insert(fragment_id);
1303 }
1304 Ok(bitmap)
1305 }
1306
1307 fn manifest_from_overwrite_transaction(
1308 previous: &Manifest,
1309 schema: lance_core::datatypes::Schema,
1310 fragments: &[Fragment],
1311 ) -> Manifest {
1312 let mut next_fragment_id = 0;
1313 let mut fragments = fragments
1314 .iter()
1315 .cloned()
1316 .map(|mut fragment| {
1317 if fragment.id == 0 {
1318 fragment.id = next_fragment_id;
1319 next_fragment_id += 1;
1320 }
1321 fragment
1322 })
1323 .collect::<Vec<_>>();
1324 fragments.sort_by_key(|fragment| fragment.id);
1325 Manifest::new_from_previous(previous, schema, Arc::new(fragments))
1326 }
1327
1328 async fn build_manifest_indices(
1329 dataset: &Dataset,
1330 manifest: &Manifest,
1331 index_data: ManifestIndexAccumulator,
1332 index_uuids: [Uuid; 3],
1333 ) -> Result<Vec<IndexMetadata>> {
1334 let fragment_bitmap = Self::manifest_fragment_bitmap(manifest)?;
1335 let schema = &manifest.schema;
1336 let ManifestIndexAccumulator {
1337 object_ids,
1338 object_types,
1339 base_objects_values,
1340 base_objects_row_ids,
1341 ..
1342 } = index_data;
1343 let [object_id_uuid, object_type_uuid, base_objects_uuid] = index_uuids;
1344 let registry = IndexPluginRegistry::with_default_plugins();
1345
1346 let dataset_version = manifest.version;
1347 let object_id_index_fut = Self::build_manifest_index(
1348 dataset,
1349 registry.clone(),
1350 schema,
1351 ManifestIndexBuildInput {
1352 index_name: OBJECT_ID_INDEX_NAME,
1353 column_name: "object_id",
1354 params: ScalarIndexParams::for_builtin(BuiltinIndexType::BTree),
1355 field: Field::new(VALUE_COLUMN_NAME, DataType::Utf8, false),
1356 stream: Self::object_id_index_stream(object_ids),
1357 },
1358 &fragment_bitmap,
1359 dataset_version,
1360 object_id_uuid,
1361 );
1362 let object_type_index_fut = Self::build_manifest_index(
1363 dataset,
1364 registry.clone(),
1365 schema,
1366 ManifestIndexBuildInput {
1367 index_name: OBJECT_TYPE_INDEX_NAME,
1368 column_name: "object_type",
1369 params: ScalarIndexParams::for_builtin(BuiltinIndexType::Bitmap),
1370 field: Field::new(VALUE_COLUMN_NAME, DataType::Utf8, false),
1371 stream: Self::object_type_index_stream(object_types),
1372 },
1373 &fragment_bitmap,
1374 dataset_version,
1375 object_type_uuid,
1376 );
1377 let base_objects_index_fut = Self::build_manifest_index(
1378 dataset,
1379 registry,
1380 schema,
1381 ManifestIndexBuildInput {
1382 index_name: BASE_OBJECTS_INDEX_NAME,
1383 column_name: "base_objects",
1384 params: ScalarIndexParams::for_builtin(BuiltinIndexType::LabelList),
1385 field: BASE_OBJECTS_VALUE_FIELD.clone(),
1386 stream: Self::base_objects_index_stream(base_objects_values, base_objects_row_ids),
1387 },
1388 &fragment_bitmap,
1389 dataset_version,
1390 base_objects_uuid,
1391 );
1392
1393 let (object_id_index, object_type_index, base_objects_index) = futures::join!(
1394 object_id_index_fut,
1395 object_type_index_fut,
1396 base_objects_index_fut
1397 );
1398
1399 Ok(vec![
1400 object_id_index?,
1401 object_type_index?,
1402 base_objects_index?,
1403 ])
1404 }
1405
1406 async fn build_manifest_index(
1407 dataset: &Dataset,
1408 registry: Arc<IndexPluginRegistry>,
1409 lance_schema: &lance_core::datatypes::Schema,
1410 input: ManifestIndexBuildInput,
1411 fragment_bitmap: &RoaringBitmap,
1412 dataset_version: u64,
1413 index_uuid: Uuid,
1414 ) -> Result<IndexMetadata> {
1415 let trained_index =
1416 Self::train_manifest_index(dataset, registry, input, index_uuid).await?;
1417 Self::manifest_index_metadata(
1418 lance_schema,
1419 fragment_bitmap,
1420 dataset_version,
1421 trained_index,
1422 )
1423 }
1424
1425 fn manifest_schema() -> Arc<ArrowSchema> {
1427 Arc::new(ArrowSchema::new(vec![
1428 Field::new("object_id", DataType::Utf8, false).with_metadata(
1430 [(
1431 LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(),
1432 "0".to_string(),
1433 )]
1434 .into_iter()
1435 .collect(),
1436 ),
1437 Field::new("object_type", DataType::Utf8, false),
1438 Field::new("location", DataType::Utf8, true),
1439 Field::new("metadata", DataType::Utf8, true),
1440 Field::new(
1441 "base_objects",
1442 DataType::List(Arc::new(Field::new("object_id", DataType::Utf8, true))),
1443 true,
1444 ),
1445 ]))
1446 }
1447
1448 async fn manifest_scanner(&self) -> Result<Scanner> {
1450 let dataset_guard = self.manifest_dataset.get().await?;
1451 Ok(dataset_guard.scan())
1452 }
1453
1454 async fn execute_scanner(scanner: Scanner) -> Result<Vec<RecordBatch>> {
1456 let mut stream = scanner.try_into_stream().await.map_err(|e| {
1457 lance_core::Error::from(NamespaceError::Internal {
1458 message: format!("Failed to create stream: {:?}", e),
1459 })
1460 })?;
1461
1462 let mut batches = Vec::new();
1463 while let Some(batch) = stream.next().await {
1464 batches.push(batch.map_err(|e| {
1465 lance_core::Error::from(NamespaceError::Internal {
1466 message: format!("Failed to read batch: {:?}", e),
1467 })
1468 })?);
1469 }
1470
1471 Ok(batches)
1472 }
1473
1474 fn get_string_column<'a>(batch: &'a RecordBatch, column_name: &str) -> Result<&'a StringArray> {
1476 let column = batch.column_by_name(column_name).ok_or_else(|| {
1477 lance_core::Error::from(NamespaceError::Internal {
1478 message: format!("Column '{}' not found", column_name),
1479 })
1480 })?;
1481 column
1482 .as_any()
1483 .downcast_ref::<StringArray>()
1484 .ok_or_else(|| {
1485 lance_core::Error::from(NamespaceError::Internal {
1486 message: format!("Column '{}' is not a string array", column_name),
1487 })
1488 })
1489 }
1490
1491 fn required_string_value<'a>(
1492 array: &'a StringArray,
1493 row: usize,
1494 column_name: &str,
1495 ) -> Result<&'a str> {
1496 if array.is_null(row) {
1497 return Err(NamespaceError::Internal {
1498 message: format!("Manifest column '{}' has null at row {}", column_name, row),
1499 }
1500 .into());
1501 }
1502 Ok(array.value(row))
1503 }
1504
1505 fn optional_string_value(array: &StringArray, row: usize) -> Option<String> {
1506 (!array.is_null(row)).then(|| array.value(row).to_string())
1507 }
1508
1509 fn base_objects_column_values(batch: &RecordBatch) -> Result<Vec<Option<Vec<String>>>> {
1510 let Some(column) = batch.column_by_name("base_objects") else {
1511 return Ok(vec![None; batch.num_rows()]);
1512 };
1513 let array = column.as_any().downcast_ref::<ListArray>().ok_or_else(|| {
1514 lance_core::Error::from(NamespaceError::Internal {
1515 message: format!(
1516 "Column 'base_objects' is not a list array: {:?}",
1517 column.data_type()
1518 ),
1519 })
1520 })?;
1521
1522 let mut values = Vec::with_capacity(batch.num_rows());
1523 for row in 0..batch.num_rows() {
1524 if array.is_null(row) {
1525 values.push(None);
1526 continue;
1527 }
1528 let row_values = array.value(row);
1529 let row_values = row_values
1530 .as_any()
1531 .downcast_ref::<StringArray>()
1532 .ok_or_else(|| {
1533 lance_core::Error::from(NamespaceError::Internal {
1534 message: "Column 'base_objects' values are not strings".to_string(),
1535 })
1536 })?;
1537 let mut objects = Vec::with_capacity(row_values.len());
1538 for value_index in 0..row_values.len() {
1539 if row_values.is_null(value_index) {
1540 return Err(NamespaceError::Internal {
1541 message: format!(
1542 "Manifest column 'base_objects' has null item at row {} item {}",
1543 row, value_index
1544 ),
1545 }
1546 .into());
1547 }
1548 objects.push(row_values.value(value_index).to_string());
1549 }
1550 values.push(Some(objects));
1551 }
1552 Ok(values)
1553 }
1554
1555 async fn manifest_projected_stream(dataset: &Dataset) -> Result<SendableRecordBatchStream> {
1556 let mut scanner = dataset.scan();
1557 scanner
1558 .project(&[
1559 "object_id",
1560 "object_type",
1561 "location",
1562 "metadata",
1563 "base_objects",
1564 ])
1565 .map_err(|e| {
1566 lance_core::Error::from(NamespaceError::Internal {
1567 message: format!("Failed to project manifest columns: {:?}", e),
1568 })
1569 })?;
1570 let stream = scanner.try_into_stream().await.map_err(|e| {
1571 lance_core::Error::from(NamespaceError::Internal {
1572 message: format!("Failed to create manifest stream: {:?}", e),
1573 })
1574 })?;
1575 let schema = stream.schema();
1576 let stream = stream.map_err(|err| DataFusionError::External(Box::new(err)));
1577 Ok(Box::pin(DatafusionRecordBatchStreamAdapter::new(
1578 schema,
1579 stream.fuse(),
1580 )))
1581 }
1582
1583 fn manifest_rewrite_commit_retries(&self) -> u32 {
1584 self.commit_retries
1585 .unwrap_or(DEFAULT_MANIFEST_REWRITE_COMMIT_RETRIES)
1586 }
1587
1588 fn lock_manifest_rewrite_shared<M: ManifestStreamMutation>(
1589 shared: &Arc<StdMutex<ManifestRewriteShared<M>>>,
1590 ) -> Result<StdMutexGuard<'_, ManifestRewriteShared<M>>> {
1591 shared.lock().map_err(|_| {
1592 lance_core::Error::from(NamespaceError::Internal {
1593 message: "Manifest rewrite state mutex was poisoned".to_string(),
1594 })
1595 })
1596 }
1597
1598 fn set_manifest_rewrite_error<M: ManifestStreamMutation>(
1599 shared: &Arc<StdMutex<ManifestRewriteShared<M>>>,
1600 err: LanceError,
1601 ) {
1602 match shared.lock() {
1603 Ok(mut guard) => {
1604 guard.error = Some(err);
1605 }
1606 Err(poisoned) => {
1607 let mut guard = poisoned.into_inner();
1608 guard.error = Some(err);
1609 }
1610 }
1611 }
1612
1613 fn take_manifest_rewrite_error<M: ManifestStreamMutation>(
1614 shared: &Arc<StdMutex<ManifestRewriteShared<M>>>,
1615 ) -> Result<Option<LanceError>> {
1616 let mut guard = Self::lock_manifest_rewrite_shared(shared)?;
1617 Ok(guard.error.take())
1618 }
1619
1620 fn process_manifest_rewrite_batch<M: ManifestStreamMutation>(
1621 batch: RecordBatch,
1622 shared: &Arc<StdMutex<ManifestRewriteShared<M>>>,
1623 ) -> Result<Option<RecordBatch>> {
1624 let object_ids = Self::get_string_column(&batch, "object_id")?;
1625 let object_types = Self::get_string_column(&batch, "object_type")?;
1626 let locations = Self::get_string_column(&batch, "location")?;
1627 let metadatas = Self::get_string_column(&batch, "metadata")?;
1628 let base_objects = Self::base_objects_column_values(&batch)?;
1629 let mut output = ManifestBatchBuilder::new();
1630 let mut guard = Self::lock_manifest_rewrite_shared(shared)?;
1631 let mut index_data = guard.index_data.take().ok_or_else(|| {
1632 lance_core::Error::from(NamespaceError::Internal {
1633 message: "Manifest rewrite index state is unavailable".to_string(),
1634 })
1635 })?;
1636 for (row, base_objects) in base_objects.into_iter().enumerate().take(batch.num_rows()) {
1637 let row_value = ManifestRowValue {
1638 object_id: Self::required_string_value(object_ids, row, "object_id")?.to_string(),
1639 object_type: ObjectType::parse(Self::required_string_value(
1640 object_types,
1641 row,
1642 "object_type",
1643 )?)?,
1644 location: Self::optional_string_value(locations, row),
1645 metadata: Self::optional_string_value(metadatas, row),
1646 base_objects,
1647 };
1648 guard
1649 .mutation
1650 .process_existing_row(row_value, &mut output, &mut index_data)?;
1651 }
1652 guard.index_data = Some(index_data);
1653 if output.is_empty() {
1654 return Ok(None);
1655 }
1656 Ok(Some(output.finish()?))
1657 }
1658
1659 fn finish_manifest_rewrite_stream<M: ManifestStreamMutation>(
1660 shared: &Arc<StdMutex<ManifestRewriteShared<M>>>,
1661 ) -> Result<Option<RecordBatch>> {
1662 let mut output = ManifestBatchBuilder::new();
1663 let mut guard = Self::lock_manifest_rewrite_shared(shared)?;
1664 let mut index_data = guard.index_data.take().ok_or_else(|| {
1665 lance_core::Error::from(NamespaceError::Internal {
1666 message: "Manifest rewrite index state is unavailable".to_string(),
1667 })
1668 })?;
1669 guard.mutation.append_rows(&mut output, &mut index_data)?;
1670 let result = guard.mutation.finish();
1671 let force_empty_batch = index_data.row_count == 0;
1672 guard.result = Some(result);
1673 guard.index_data = Some(index_data);
1674 if output.is_empty() && !force_empty_batch {
1675 Ok(None)
1676 } else {
1677 Ok(Some(output.finish()?))
1678 }
1679 }
1680
1681 fn manifest_rewrite_output_stream<M: ManifestStreamMutation + 'static>(
1682 source: SendableRecordBatchStream,
1683 shared: Arc<StdMutex<ManifestRewriteShared<M>>>,
1684 ) -> SendableRecordBatchStream {
1685 enum Phase {
1686 Source,
1687 Finish,
1688 Done,
1689 }
1690
1691 let schema = Self::manifest_schema();
1692 let stream = stream::unfold(
1693 (source, shared, Phase::Source),
1694 |(mut source, shared, mut phase)| async move {
1695 loop {
1696 match phase {
1697 Phase::Source => match source.next().await {
1698 Some(Ok(batch)) => {
1699 match Self::process_manifest_rewrite_batch(batch, &shared) {
1700 Ok(Some(batch)) => {
1701 return Some((Ok(batch), (source, shared, phase)));
1702 }
1703 Ok(None) => continue,
1704 Err(err) => {
1705 let message = err.to_string();
1706 Self::set_manifest_rewrite_error(&shared, err);
1707 return Some((
1708 Err(DataFusionError::External(Box::new(
1709 std::io::Error::other(message),
1710 ))),
1711 (source, shared, Phase::Done),
1712 ));
1713 }
1714 }
1715 }
1716 Some(Err(err)) => {
1717 return Some((Err(err), (source, shared, Phase::Done)));
1718 }
1719 None => phase = Phase::Finish,
1720 },
1721 Phase::Finish => {
1722 phase = Phase::Done;
1723 match Self::finish_manifest_rewrite_stream(&shared) {
1724 Ok(Some(batch)) => {
1725 return Some((Ok(batch), (source, shared, phase)));
1726 }
1727 Ok(None) => continue,
1728 Err(err) => {
1729 let message = err.to_string();
1730 Self::set_manifest_rewrite_error(&shared, err);
1731 return Some((
1732 Err(DataFusionError::External(Box::new(
1733 std::io::Error::other(message),
1734 ))),
1735 (source, shared, Phase::Done),
1736 ));
1737 }
1738 }
1739 }
1740 Phase::Done => return None,
1741 }
1742 }
1743 },
1744 );
1745 Box::pin(DatafusionRecordBatchStreamAdapter::new(
1746 schema,
1747 stream.fuse(),
1748 ))
1749 }
1750
1751 fn take_manifest_rewrite_result<M: ManifestStreamMutation>(
1752 shared: &Arc<StdMutex<ManifestRewriteShared<M>>>,
1753 ) -> Result<(CopyOnWriteMutation<M::Output>, ManifestIndexAccumulator)> {
1754 let mut guard = Self::lock_manifest_rewrite_shared(shared)?;
1755 let result = guard.result.take().ok_or_else(|| {
1756 lance_core::Error::from(NamespaceError::Internal {
1757 message: "Manifest rewrite stream did not finish".to_string(),
1758 })
1759 })?;
1760 let index_data = guard.index_data.take().ok_or_else(|| {
1761 lance_core::Error::from(NamespaceError::Internal {
1762 message: "Manifest rewrite index state is unavailable".to_string(),
1763 })
1764 })?;
1765 Ok((result, index_data))
1766 }
1767
1768 async fn cleanup_staged_manifest_files(
1773 &self,
1774 object_store: &ObjectStore,
1775 data_files: &HashSet<String>,
1776 index_uuids: &[Uuid],
1777 ) {
1778 let data_dir = self
1779 .base_path
1780 .clone()
1781 .join(MANIFEST_TABLE_NAME)
1782 .join(LANCE_DATA_DIR);
1783 for path in data_files {
1784 let data_path = data_dir.clone().join(path.as_str());
1785 if let Err(err) = object_store.delete(&data_path).await {
1786 log::warn!(
1787 "Failed to clean up uncommitted manifest rewrite data file '{}': {}",
1788 data_path,
1789 err
1790 );
1791 }
1792 }
1793 self.cleanup_uncommitted_manifest_index_dirs(object_store, index_uuids.iter().copied())
1794 .await;
1795 }
1796
1797 async fn cleanup_uncommitted_manifest_index_dirs(
1798 &self,
1799 object_store: &ObjectStore,
1800 index_uuids: impl IntoIterator<Item = Uuid>,
1801 ) {
1802 for index_uuid in index_uuids {
1803 let index_dir = self
1804 .base_path
1805 .clone()
1806 .join(MANIFEST_TABLE_NAME)
1807 .join(LANCE_INDICES_DIR)
1808 .join(index_uuid.to_string());
1809 if let Err(err) = object_store.remove_dir_all(index_dir.clone()).await
1810 && !matches!(err, LanceError::NotFound { .. })
1811 {
1812 log::warn!(
1813 "Failed to clean up uncommitted manifest rewrite index directory '{}': {}",
1814 index_dir,
1815 err
1816 );
1817 }
1818 }
1819 }
1820
1821 async fn manifest_commit_handler(&self) -> Result<Arc<dyn CommitHandler>> {
1823 commit_handler_from_url(&self.root, &None)
1824 .await
1825 .map_err(|e| {
1826 lance_core::Error::from(NamespaceError::Internal {
1827 message: format!("Failed to resolve manifest commit handler: {:?}", e),
1828 })
1829 })
1830 }
1831
1832 async fn commit_manifest_overwrite(
1836 &self,
1837 dataset: &Dataset,
1838 commit_handler: &dyn CommitHandler,
1839 manifest: &mut Manifest,
1840 indices: Option<Vec<IndexMetadata>>,
1841 transaction: Transaction,
1842 ) -> std::result::Result<(), CommitError> {
1843 apply_feature_flags(manifest, false, false).map_err(CommitError::from)?;
1844 let timestamp_nanos = SystemTime::now()
1845 .duration_since(UNIX_EPOCH)
1846 .map(|d| d.as_nanos())
1847 .unwrap_or(0);
1848 manifest.set_timestamp(timestamp_nanos);
1849 manifest.update_max_fragment_id();
1850
1851 let object_store = dataset
1855 .object_store(None)
1856 .await
1857 .map_err(CommitError::from)?;
1858 let base_path = self.base_path.clone().join(MANIFEST_TABLE_NAME);
1859 let naming_scheme = dataset.manifest_location().naming_scheme;
1860 commit_handler
1861 .commit(
1862 manifest,
1863 indices,
1864 &base_path,
1865 &object_store,
1866 write_manifest_file_to_path,
1867 naming_scheme,
1868 Some((&transaction).into()),
1869 )
1870 .await
1871 .map(|_location| ())
1872 }
1873
1874 async fn manifest_commit_landed(
1879 &self,
1880 dataset: &Dataset,
1881 target_version: u64,
1882 data_files: &HashSet<String>,
1883 ) -> bool {
1884 let Ok(committed) = dataset.checkout_version(target_version).await else {
1885 return false;
1886 };
1887 committed.manifest().fragments.iter().any(|fragment| {
1888 fragment
1889 .files
1890 .iter()
1891 .any(|file| data_files.contains(file.path.as_str()))
1892 })
1893 }
1894
1895 async fn resolve_manifest_conflict<O: Clone>(
1899 &self,
1900 resolution: &ConflictResolution<O>,
1901 ) -> Result<Option<O>> {
1902 match resolution {
1903 ConflictResolution::Retry => Ok(None),
1904 ConflictResolution::FailIfExists(object_ids) => {
1905 for object_id in object_ids {
1906 if self.manifest_contains_object(object_id).await? {
1907 return Err(NamespaceError::ConcurrentModification {
1908 message: format!(
1909 "Object '{}' was concurrently created by another operation",
1910 object_id
1911 ),
1912 }
1913 .into());
1914 }
1915 }
1916 Ok(None)
1917 }
1918 ConflictResolution::SucceedIfAbsent { object_id, output } => {
1919 if self.manifest_contains_object(object_id).await? {
1920 Ok(None)
1921 } else {
1922 Ok(Some(output.clone()))
1923 }
1924 }
1925 }
1926 }
1927
1928 async fn ensure_manifest_writable(&self) -> Result<()> {
1934 let dataset_guard = self.manifest_dataset.get().await?;
1935 ensure_writable(dataset_guard.metadata())
1936 }
1937
1938 async fn rewrite_manifest<M, F>(
1939 &self,
1940 operation: &str,
1941 mut make_mutation: F,
1942 ) -> Result<M::Output>
1943 where
1944 M: ManifestStreamMutation + 'static,
1945 F: FnMut() -> M,
1946 {
1947 let _mutation_guard = self.manifest_mutation_lock.lock().await;
1948 let max_retries = self.manifest_rewrite_commit_retries();
1949 let mut retries = 0;
1950 let build_indices = self.inline_optimization_enabled;
1951 let commit_handler = self.manifest_commit_handler().await?;
1952
1953 loop {
1954 let dataset_guard = self.manifest_dataset.get_refreshed().await?;
1955 let dataset = Arc::new(dataset_guard.clone());
1956 drop(dataset_guard);
1957 ensure_writable(dataset.metadata())?;
1960 let object_store = dataset.object_store(None).await?;
1963
1964 let source = Self::manifest_projected_stream(&dataset).await?;
1965 let resolution = make_mutation().conflict_resolution();
1966 let shared = Arc::new(StdMutex::new(ManifestRewriteShared::new(make_mutation())));
1967 let output_stream = Self::manifest_rewrite_output_stream(source, shared.clone());
1968 let write_params = WriteParams {
1973 mode: WriteMode::Overwrite,
1974 session: self.session.clone(),
1975 max_rows_per_file: u32::MAX as usize,
1976 max_bytes_per_file: usize::MAX,
1977 skip_auto_cleanup: true,
1978 ..WriteParams::default()
1979 };
1980
1981 let transaction = match InsertBuilder::new(dataset.clone())
1982 .with_params(&write_params)
1983 .execute_uncommitted_stream(output_stream)
1984 .await
1985 {
1986 Ok(transaction) => transaction,
1987 Err(err) => {
1988 if let Some(stream_err) = Self::take_manifest_rewrite_error(&shared)? {
1989 return Err(stream_err);
1990 }
1991 return Err(convert_lance_commit_error(&err, operation, None));
1992 }
1993 };
1994
1995 let (mutation, index_data) = Self::take_manifest_rewrite_result(&shared)?;
1996
1997 let Operation::Overwrite {
1998 fragments, schema, ..
1999 } = &transaction.operation
2000 else {
2001 return Err(NamespaceError::Internal {
2002 message: "Manifest rewrite transaction is not an overwrite".to_string(),
2003 }
2004 .into());
2005 };
2006 let staged_data_files = fragments
2009 .iter()
2010 .flat_map(|fragment| fragment.files.iter())
2011 .filter(|file| file.base_id.is_none())
2012 .map(|file| file.path.clone())
2013 .collect::<HashSet<_>>();
2014
2015 if !mutation.has_changes {
2016 self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[])
2017 .await;
2018 return Ok(mutation.result);
2019 }
2020
2021 let mut manifest = Self::manifest_from_overwrite_transaction(
2022 dataset.manifest(),
2023 schema.clone(),
2024 fragments,
2025 );
2026 let target_version = manifest.version;
2027
2028 let index_uuids = [Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4()];
2029 let indices = if build_indices {
2030 match Self::build_manifest_indices(&dataset, &manifest, index_data, index_uuids)
2031 .await
2032 {
2033 Ok(indices) => Some(indices),
2034 Err(err) => {
2035 self.cleanup_staged_manifest_files(
2036 &object_store,
2037 &staged_data_files,
2038 &index_uuids,
2039 )
2040 .await;
2041 return Err(err);
2042 }
2043 }
2044 } else {
2045 None
2046 };
2047 let staged_index_uuids: &[Uuid] = if build_indices { &index_uuids } else { &[] };
2048
2049 let commit_result = self
2050 .commit_manifest_overwrite(
2051 &dataset,
2052 commit_handler.as_ref(),
2053 &mut manifest,
2054 indices,
2055 transaction,
2056 )
2057 .await;
2058
2059 match commit_result {
2060 Ok(()) => {
2061 let _ = self.manifest_dataset.get_refreshed().await;
2062 return Ok(mutation.result);
2063 }
2064 Err(err) => {
2065 if self
2069 .manifest_commit_landed(&dataset, target_version, &staged_data_files)
2070 .await
2071 {
2072 let _ = self.manifest_dataset.get_refreshed().await;
2073 return Ok(mutation.result);
2074 }
2075 self.cleanup_staged_manifest_files(
2076 &object_store,
2077 &staged_data_files,
2078 staged_index_uuids,
2079 )
2080 .await;
2081 match err {
2082 CommitError::CommitConflict => {
2083 if let Some(output) =
2084 self.resolve_manifest_conflict(&resolution).await?
2085 {
2086 return Ok(output);
2087 }
2088 if retries >= max_retries {
2089 return Err(NamespaceError::ConcurrentModification {
2090 message: format!(
2091 "{}: still conflicting after {} retries",
2092 operation, max_retries
2093 ),
2094 }
2095 .into());
2096 }
2097 retries += 1;
2098 tokio::time::sleep(std::time::Duration::from_millis(
2099 10 * u64::from(retries),
2100 ))
2101 .await;
2102 }
2103 CommitError::OtherError(err) => {
2104 return Err(convert_lance_commit_error(&err, operation, None));
2105 }
2106 }
2107 }
2108 }
2109 }
2110 }
2111
2112 async fn manifest_contains_object(&self, object_id: &str) -> Result<bool> {
2114 let escaped_id = object_id.replace('\'', "''");
2115 let filter = format!("object_id = '{}'", escaped_id);
2116
2117 let dataset_guard = self.manifest_dataset.get().await?;
2118 let mut scanner = dataset_guard.scan();
2119
2120 scanner.filter(&filter).map_err(|e| {
2121 lance_core::Error::from(NamespaceError::Internal {
2122 message: format!("Failed to filter: {:?}", e),
2123 })
2124 })?;
2125
2126 scanner.project::<&str>(&[]).map_err(|e| {
2128 lance_core::Error::from(NamespaceError::Internal {
2129 message: format!("Failed to project: {:?}", e),
2130 })
2131 })?;
2132
2133 scanner.with_row_id();
2134
2135 let count = scanner.count_rows().await.map_err(|e| {
2136 lance_core::Error::from(NamespaceError::Internal {
2137 message: format!("Failed to count rows: {:?}", e),
2138 })
2139 })?;
2140
2141 Ok(count > 0)
2142 }
2143
2144 async fn query_manifest_for_table(&self, object_id: &str) -> Result<Option<TableInfo>> {
2146 let escaped_id = object_id.replace('\'', "''");
2147 let filter = format!("object_id = '{}' AND object_type = 'table'", escaped_id);
2148 let mut scanner = self.manifest_scanner().await?;
2149 scanner.filter(&filter).map_err(|e| {
2150 lance_core::Error::from(NamespaceError::Internal {
2151 message: format!("Failed to filter: {:?}", e),
2152 })
2153 })?;
2154 scanner
2155 .project(&["object_id", "location", "metadata"])
2156 .map_err(|e| {
2157 lance_core::Error::from(NamespaceError::Internal {
2158 message: format!("Failed to project: {:?}", e),
2159 })
2160 })?;
2161 let batches = Self::execute_scanner(scanner).await?;
2162
2163 let mut found_result: Option<TableInfo> = None;
2164 let mut total_rows = 0;
2165
2166 for batch in batches {
2167 if batch.num_rows() == 0 {
2168 continue;
2169 }
2170
2171 total_rows += batch.num_rows();
2172 if total_rows > 1 {
2173 return Err(NamespaceError::Internal {
2174 message: format!(
2175 "Expected exactly 1 table with id '{}', found {}",
2176 object_id, total_rows
2177 ),
2178 }
2179 .into());
2180 }
2181
2182 let object_id_array = Self::get_string_column(&batch, "object_id")?;
2183 let location_array = Self::get_string_column(&batch, "location")?;
2184 let metadata_array = Self::get_string_column(&batch, "metadata")?;
2185 let location = location_array.value(0).to_string();
2186 let metadata = if !metadata_array.is_null(0) {
2187 let metadata_str = metadata_array.value(0);
2188 match serde_json::from_str::<HashMap<String, String>>(metadata_str) {
2189 Ok(map) => Some(map),
2190 Err(e) => {
2191 return Err(NamespaceError::Internal {
2192 message: format!(
2193 "Failed to deserialize metadata for table '{}': {}",
2194 object_id, e
2195 ),
2196 }
2197 .into());
2198 }
2199 }
2200 } else {
2201 None
2202 };
2203 let (namespace, name) = Self::parse_object_id(object_id_array.value(0));
2204 found_result = Some(TableInfo {
2205 namespace,
2206 name,
2207 location,
2208 metadata,
2209 });
2210 }
2211
2212 Ok(found_result)
2213 }
2214
2215 fn serialize_metadata(
2216 properties: Option<&HashMap<String, String>>,
2217 object_type: &str,
2218 object_id: &str,
2219 ) -> Result<Option<String>> {
2220 match properties {
2221 Some(properties) if !properties.is_empty() => {
2222 serde_json::to_string(properties).map(Some).map_err(|e| {
2223 LanceError::from(NamespaceError::Internal {
2224 message: format!(
2225 "Failed to serialize {} metadata for '{}': {}",
2226 object_type, object_id, e
2227 ),
2228 })
2229 })
2230 }
2231 _ => Ok(None),
2232 }
2233 }
2234
2235 pub(crate) async fn path_has_actual_manifests(
2236 object_store: &ObjectStore,
2237 table_path: &Path,
2238 ) -> Result<bool> {
2239 let versions_path = table_path
2240 .clone()
2241 .join(lance_table::io::commit::VERSIONS_DIR);
2242 Ok(object_store
2245 .list(Some(versions_path))
2246 .try_next()
2247 .await?
2248 .is_some())
2249 }
2250
2251 async fn location_has_actual_manifests(&self, location: &str) -> Result<bool> {
2252 Self::path_has_actual_manifests(&self.object_store, &self.base_path.clone().join(location))
2253 .await
2254 }
2255
2256 pub(crate) fn is_not_found_load_error(err: &LanceError) -> bool {
2257 match err {
2258 LanceError::NotFound { .. } => true,
2259 LanceError::IO { source, .. } => source
2260 .downcast_ref::<ObjectStoreError>()
2261 .is_some_and(|source| matches!(source, ObjectStoreError::NotFound { .. })),
2262 LanceError::DatasetNotFound { source, .. } => {
2263 source
2264 .downcast_ref::<LanceError>()
2265 .is_some_and(|source| matches!(source, LanceError::NotFound { .. }))
2266 || source
2267 .downcast_ref::<ObjectStoreError>()
2268 .is_some_and(|source| matches!(source, ObjectStoreError::NotFound { .. }))
2269 }
2270 _ => false,
2271 }
2272 }
2273
2274 pub async fn list_manifest_table_locations(&self) -> Result<std::collections::HashSet<String>> {
2277 let filter = "object_type = 'table' AND NOT contains(object_id, '$')";
2278 let mut scanner = self.manifest_scanner().await?;
2279 scanner.filter(filter).map_err(|e| {
2280 lance_core::Error::from(NamespaceError::Internal {
2281 message: format!("Failed to filter: {:?}", e),
2282 })
2283 })?;
2284 scanner.project(&["location"]).map_err(|e| {
2285 lance_core::Error::from(NamespaceError::Internal {
2286 message: format!("Failed to project: {:?}", e),
2287 })
2288 })?;
2289
2290 let batches = Self::execute_scanner(scanner).await?;
2291 let mut locations = std::collections::HashSet::new();
2292
2293 for batch in batches {
2294 if batch.num_rows() == 0 {
2295 continue;
2296 }
2297 let location_array = Self::get_string_column(&batch, "location")?;
2298 for i in 0..location_array.len() {
2299 locations.insert(location_array.value(i).to_string());
2300 }
2301 }
2302
2303 Ok(locations)
2304 }
2305
2306 async fn insert_into_manifest(
2308 &self,
2309 object_id: String,
2310 object_type: ObjectType,
2311 location: Option<String>,
2312 ) -> Result<()> {
2313 self.insert_into_manifest_with_metadata(
2314 vec![ManifestEntry {
2315 object_id,
2316 object_type,
2317 location,
2318 metadata: None,
2319 }],
2320 None,
2321 )
2322 .await
2323 }
2324
2325 pub async fn insert_into_manifest_with_metadata(
2330 &self,
2331 entries: Vec<ManifestEntry>,
2332 base_objects: Option<Vec<String>>,
2333 ) -> Result<()> {
2334 self.merge_into_manifest_with_metadata(entries, base_objects, WhenMatched::Fail)
2335 .await
2336 }
2337
2338 async fn upsert_into_manifest_with_metadata(
2339 &self,
2340 entries: Vec<ManifestEntry>,
2341 base_objects: Option<Vec<String>>,
2342 ) -> Result<()> {
2343 self.merge_into_manifest_with_metadata(entries, base_objects, WhenMatched::UpdateAll)
2344 .await
2345 }
2346
2347 async fn merge_into_manifest_with_metadata(
2348 &self,
2349 entries: Vec<ManifestEntry>,
2350 base_objects: Option<Vec<String>>,
2351 when_matched: WhenMatched,
2352 ) -> Result<()> {
2353 if entries.is_empty() {
2354 return Ok(());
2355 }
2356
2357 self.rewrite_manifest("Failed to overwrite manifest", || {
2358 UpsertManifestMutation::new(entries.clone(), base_objects.clone(), when_matched.clone())
2359 })
2360 .await
2361 }
2362
2363 pub async fn delete_from_manifest(&self, object_id: &str) -> Result<()> {
2365 let object_id = object_id.to_string();
2366 self.rewrite_manifest("Failed to delete from manifest", || DeleteObjectMutation {
2367 object_id: object_id.clone(),
2368 deleted: false,
2369 })
2370 .await
2371 }
2372
2373 pub async fn register_table(&self, name: &str, location: String) -> Result<()> {
2375 let object_id = Self::build_object_id(&[], name);
2376 if self.manifest_contains_object(&object_id).await? {
2377 return Err(NamespaceError::Internal {
2378 message: format!("Table '{}' already exists", name),
2379 }
2380 .into());
2381 }
2382
2383 self.insert_into_manifest(object_id, ObjectType::Table, Some(location))
2384 .await
2385 }
2386
2387 async fn validate_namespace_levels_exist(&self, namespace_path: &[String]) -> Result<()> {
2389 for i in 1..=namespace_path.len() {
2390 let partial_path = &namespace_path[..i];
2391 let object_id = partial_path.join(DELIMITER);
2392 if !self.manifest_contains_object(&object_id).await? {
2393 return Err(NamespaceError::NamespaceNotFound {
2394 message: format!("parent namespace '{}'", object_id),
2395 }
2396 .into());
2397 }
2398 }
2399 Ok(())
2400 }
2401
2402 async fn query_manifest_for_namespace(&self, object_id: &str) -> Result<Option<NamespaceInfo>> {
2404 let escaped_id = object_id.replace('\'', "''");
2405 let filter = format!("object_id = '{}' AND object_type = 'namespace'", escaped_id);
2406 let mut scanner = self.manifest_scanner().await?;
2407 scanner.filter(&filter).map_err(|e| {
2408 lance_core::Error::from(NamespaceError::Internal {
2409 message: format!("Failed to filter: {:?}", e),
2410 })
2411 })?;
2412 scanner.project(&["object_id", "metadata"]).map_err(|e| {
2413 lance_core::Error::from(NamespaceError::Internal {
2414 message: format!("Failed to project: {:?}", e),
2415 })
2416 })?;
2417 let batches = Self::execute_scanner(scanner).await?;
2418
2419 let mut found_result: Option<NamespaceInfo> = None;
2420 let mut total_rows = 0;
2421
2422 for batch in batches {
2423 if batch.num_rows() == 0 {
2424 continue;
2425 }
2426
2427 total_rows += batch.num_rows();
2428 if total_rows > 1 {
2429 return Err(NamespaceError::Internal {
2430 message: format!(
2431 "Expected exactly 1 namespace with id '{}', found {}",
2432 object_id, total_rows
2433 ),
2434 }
2435 .into());
2436 }
2437
2438 let object_id_array = Self::get_string_column(&batch, "object_id")?;
2439 let metadata_array = Self::get_string_column(&batch, "metadata")?;
2440
2441 let object_id_str = object_id_array.value(0);
2442 let metadata = if !metadata_array.is_null(0) {
2443 let metadata_str = metadata_array.value(0);
2444 match serde_json::from_str::<HashMap<String, String>>(metadata_str) {
2445 Ok(map) => Some(map),
2446 Err(e) => {
2447 return Err(NamespaceError::Internal {
2448 message: format!(
2449 "Failed to deserialize metadata for namespace '{}': {}",
2450 object_id, e
2451 ),
2452 }
2453 .into());
2454 }
2455 }
2456 } else {
2457 None
2458 };
2459
2460 let (namespace, name) = Self::parse_object_id(object_id_str);
2461 found_result = Some(NamespaceInfo {
2462 namespace,
2463 name,
2464 metadata,
2465 });
2466 }
2467
2468 Ok(found_result)
2469 }
2470
2471 async fn open_manifest_table(
2473 root: &str,
2474 storage_options: &Option<HashMap<String, String>>,
2475 session: Option<Arc<Session>>,
2476 ) -> Result<DatasetConsistencyWrapper> {
2477 let manifest_path = format!("{}/{}", root, MANIFEST_TABLE_NAME);
2478 log::debug!("Attempting to load manifest from {}", manifest_path);
2479 let store_options = ObjectStoreParams {
2480 storage_options_accessor: storage_options.as_ref().map(|opts| {
2481 Arc::new(
2482 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2483 opts.clone(),
2484 ),
2485 )
2486 }),
2487 ..Default::default()
2488 };
2489 let read_params = ReadParams {
2490 session,
2491 store_options: Some(store_options),
2492 ..Default::default()
2493 };
2494 let dataset = DatasetBuilder::from_uri(&manifest_path)
2495 .with_read_params(read_params)
2496 .load()
2497 .await?;
2498 ensure_readable(dataset.metadata())?;
2499 Ok(DatasetConsistencyWrapper::new(dataset))
2500 }
2501
2502 async fn ensure_manifest_table_up_to_date(
2509 root: &str,
2510 storage_options: &Option<HashMap<String, String>>,
2511 session: Option<Arc<Session>>,
2512 ) -> Result<DatasetConsistencyWrapper> {
2513 let manifest_path = format!("{}/{}", root, MANIFEST_TABLE_NAME);
2514 log::debug!("Attempting to load manifest from {}", manifest_path);
2515 let store_options = ObjectStoreParams {
2516 storage_options_accessor: storage_options.as_ref().map(|opts| {
2517 Arc::new(
2518 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2519 opts.clone(),
2520 ),
2521 )
2522 }),
2523 ..Default::default()
2524 };
2525 let read_params = ReadParams {
2526 session: session.clone(),
2527 store_options: Some(store_options.clone()),
2528 ..Default::default()
2529 };
2530 let dataset_result = DatasetBuilder::from_uri(&manifest_path)
2531 .with_read_params(read_params)
2532 .load()
2533 .await;
2534 match dataset_result {
2535 Ok(mut dataset) => {
2536 ensure_readable(dataset.metadata())?;
2539
2540 let needs_pk_migration = dataset
2542 .schema()
2543 .field("object_id")
2544 .map(|f| {
2545 !f.metadata
2546 .contains_key(LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
2547 })
2548 .unwrap_or(false);
2549
2550 if needs_pk_migration {
2551 ensure_writable(dataset.metadata())?;
2554 log::info!(
2555 "Migrating __manifest table to add primary key metadata on object_id"
2556 );
2557 dataset
2558 .update_field_metadata()
2559 .update("object_id", [(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "0")])
2560 .map_err(|e| {
2561 lance_core::Error::from(NamespaceError::Internal {
2562 message: format!(
2563 "Failed to find object_id field for migration: {:?}",
2564 e
2565 ),
2566 })
2567 })?
2568 .await
2569 .map_err(|e| {
2570 lance_core::Error::from(NamespaceError::Internal {
2571 message: format!("Failed to migrate primary key metadata: {:?}", e),
2572 })
2573 })?;
2574 }
2575
2576 Ok(DatasetConsistencyWrapper::new(dataset))
2577 }
2578 Err(err) if Self::is_not_found_load_error(&err) => {
2579 log::info!("Creating new manifest table at {}", manifest_path);
2580 let schema = Self::manifest_schema();
2581 let empty_batch = RecordBatch::new_empty(schema.clone());
2582 let reader = RecordBatchIterator::new(vec![Ok(empty_batch)], schema.clone());
2583
2584 let store_params = ObjectStoreParams {
2585 storage_options_accessor: storage_options.as_ref().map(|opts| {
2586 Arc::new(
2587 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2588 opts.clone(),
2589 ),
2590 )
2591 }),
2592 ..Default::default()
2593 };
2594 let write_params = WriteParams {
2595 session: session.clone(),
2596 store_params: Some(store_params),
2597 ..Default::default()
2598 };
2599
2600 let dataset =
2601 Dataset::write(Box::new(reader), &manifest_path, Some(write_params)).await;
2602
2603 match dataset {
2605 Ok(dataset) => {
2606 log::info!(
2607 "Successfully created manifest table at {}, version={}, uri={}",
2608 manifest_path,
2609 dataset.version().version,
2610 dataset.uri()
2611 );
2612 Ok(DatasetConsistencyWrapper::new(dataset))
2613 }
2614 Err(ref e)
2615 if matches!(
2616 e,
2617 LanceError::DatasetAlreadyExists { .. }
2618 | LanceError::CommitConflict { .. }
2619 | LanceError::IncompatibleTransaction { .. }
2620 | LanceError::RetryableCommitConflict { .. }
2621 ) =>
2622 {
2623 log::info!(
2625 "Manifest table was created by another process, loading it: {}",
2626 manifest_path
2627 );
2628 let recovery_store_options = ObjectStoreParams {
2629 storage_options_accessor: storage_options.as_ref().map(|opts| {
2630 Arc::new(
2631 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2632 opts.clone(),
2633 ),
2634 )
2635 }),
2636 ..Default::default()
2637 };
2638 let recovery_read_params = ReadParams {
2639 session,
2640 store_options: Some(recovery_store_options),
2641 ..Default::default()
2642 };
2643 let dataset = DatasetBuilder::from_uri(&manifest_path)
2644 .with_read_params(recovery_read_params)
2645 .load()
2646 .await
2647 .map_err(|e| {
2648 lance_core::Error::from(NamespaceError::Internal {
2649 message: format!(
2650 "Failed to load manifest dataset after creation conflict: {}",
2651 e
2652 ),
2653 })
2654 })?;
2655 Ok(DatasetConsistencyWrapper::new(dataset))
2656 }
2657 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2658 message: format!("Failed to create manifest dataset: {:?}", e),
2659 })),
2660 }
2661 }
2662 Err(err) => Err(err),
2663 }
2664 }
2665
2666 fn apply_pagination(
2671 names: &mut Vec<String>,
2672 page_token: Option<String>,
2673 limit: Option<i32>,
2674 ) -> Option<String> {
2675 names.sort();
2676
2677 if let Some(start_after) = page_token {
2678 if let Some(index) = names
2679 .iter()
2680 .position(|name| name.as_str() > start_after.as_str())
2681 {
2682 names.drain(0..index);
2683 } else {
2684 names.clear();
2685 }
2686 }
2687
2688 if let Some(limit) = limit
2689 && limit >= 0
2690 {
2691 let limit = limit as usize;
2692 if names.len() > limit {
2693 let next_page_token = if limit > 0 {
2694 Some(names[limit - 1].clone())
2695 } else {
2696 None
2697 };
2698 names.truncate(limit);
2699 return next_page_token;
2700 }
2701 }
2702
2703 None
2704 }
2705}
2706
2707#[async_trait]
2708impl LanceNamespace for ManifestNamespace {
2709 fn namespace_id(&self) -> String {
2710 self.root.clone()
2711 }
2712
2713 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
2714 let namespace_id = request.id.as_ref().ok_or_else(|| {
2715 lance_core::Error::from(NamespaceError::InvalidInput {
2716 message: "Namespace ID is required".to_string(),
2717 })
2718 })?;
2719
2720 let filter = if namespace_id.is_empty() {
2722 "object_type = 'table' AND NOT contains(object_id, '$')".to_string()
2724 } else {
2725 let prefix = namespace_id.join(DELIMITER);
2727 format!(
2728 "object_type = 'table' AND starts_with(object_id, '{}{}') AND NOT contains(substring(object_id, {}), '$')",
2729 prefix,
2730 DELIMITER,
2731 prefix.len() + 2
2732 )
2733 };
2734
2735 let mut scanner = self.manifest_scanner().await?;
2736 scanner.filter(&filter).map_err(|e| {
2737 lance_core::Error::from(NamespaceError::Internal {
2738 message: format!("Failed to filter: {:?}", e),
2739 })
2740 })?;
2741 scanner.project(&["object_id", "location"]).map_err(|e| {
2742 lance_core::Error::from(NamespaceError::Internal {
2743 message: format!("Failed to project: {:?}", e),
2744 })
2745 })?;
2746
2747 let batches = Self::execute_scanner(scanner).await?;
2748
2749 let mut table_entries = Vec::new();
2750 for batch in batches {
2751 if batch.num_rows() == 0 {
2752 continue;
2753 }
2754
2755 let object_id_array = Self::get_string_column(&batch, "object_id")?;
2756 let location_array = Self::get_string_column(&batch, "location")?;
2757 for i in 0..batch.num_rows() {
2758 let object_id = object_id_array.value(i);
2759 let location = location_array.value(i);
2760 let (_namespace, name) = Self::parse_object_id(object_id);
2761 table_entries.push((name, location.to_string()));
2762 }
2763 }
2764
2765 let mut tables: Vec<String> = if request.include_declared.unwrap_or(true) {
2766 table_entries.into_iter().map(|(name, _)| name).collect()
2767 } else {
2768 let mut stream = futures::stream::iter(table_entries.into_iter().map(
2769 |(name, location)| async move {
2770 if self.location_has_actual_manifests(&location).await? {
2775 Ok::<Option<String>, Error>(Some(name))
2776 } else {
2777 Ok::<Option<String>, Error>(None)
2778 }
2779 },
2780 ))
2781 .buffered(DECLARED_FILTER_CONCURRENCY);
2782
2783 let mut filtered = Vec::new();
2784 while let Some(result) = stream.next().await {
2785 if let Some(name) = result? {
2786 filtered.push(name);
2787 }
2788 }
2789 filtered
2790 };
2791
2792 let next_page_token =
2793 Self::apply_pagination(&mut tables, request.page_token, request.limit);
2794 let mut response = ListTablesResponse::new(tables);
2795 response.page_token = next_page_token;
2796 Ok(response)
2797 }
2798
2799 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
2800 let table_id = request.id.as_ref().ok_or_else(|| {
2801 lance_core::Error::from(NamespaceError::InvalidInput {
2802 message: "Table ID is required".to_string(),
2803 })
2804 })?;
2805
2806 if table_id.is_empty() {
2807 return Err(NamespaceError::InvalidInput {
2808 message: "Table ID cannot be empty".to_string(),
2809 }
2810 .into());
2811 }
2812
2813 let object_id = Self::str_object_id(table_id);
2814 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
2815
2816 let table_name = table_id.last().cloned().unwrap_or_default();
2818 let namespace_id: Vec<String> = if table_id.len() > 1 {
2819 table_id[..table_id.len() - 1].to_vec()
2820 } else {
2821 vec![]
2822 };
2823
2824 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
2825 let should_check_declared =
2826 load_detailed_metadata || request.check_declared.unwrap_or(false);
2827 let vend_credentials = request.vend_credentials.unwrap_or(true);
2829
2830 match table_info {
2831 Some(info) => {
2832 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
2834
2835 let storage_options = if vend_credentials {
2836 self.storage_options.clone()
2837 } else {
2838 None
2839 };
2840 let is_only_declared = if should_check_declared {
2841 Some(!self.location_has_actual_manifests(&info.location).await?)
2842 } else {
2843 None
2844 };
2845
2846 if !load_detailed_metadata {
2847 return Ok(DescribeTableResponse {
2848 table: Some(table_name),
2849 namespace: Some(namespace_id),
2850 location: Some(table_uri.clone()),
2851 table_uri: Some(table_uri),
2852 storage_options,
2853 properties: info.metadata,
2854 is_only_declared,
2855 ..Default::default()
2856 });
2857 }
2858
2859 if is_only_declared == Some(true) {
2860 return Ok(DescribeTableResponse {
2861 table: Some(table_name),
2862 namespace: Some(namespace_id),
2863 location: Some(table_uri.clone()),
2864 table_uri: Some(table_uri),
2865 storage_options,
2866 properties: info.metadata,
2867 is_only_declared,
2868 ..Default::default()
2869 });
2870 }
2871
2872 let mut builder = DatasetBuilder::from_uri(&table_uri);
2873 if let Some(opts) = &self.storage_options {
2874 builder = builder.with_storage_options(opts.clone());
2875 }
2876 if let Some(session) = &self.session {
2877 builder = builder.with_session(session.clone());
2878 }
2879
2880 match builder.load().await {
2881 Ok(mut dataset) => {
2882 if let Some(requested_version) = request.version {
2884 dataset = dataset.checkout_version(requested_version as u64).await?;
2885 }
2886
2887 let version = dataset.version().version;
2888 let lance_schema = dataset.schema();
2889 let arrow_schema: arrow_schema::Schema = lance_schema.into();
2890 let json_schema = arrow_schema_to_json(&arrow_schema)?;
2891
2892 Ok(DescribeTableResponse {
2893 table: Some(table_name.clone()),
2894 namespace: Some(namespace_id.clone()),
2895 version: Some(version as i64),
2896 location: Some(table_uri.clone()),
2897 table_uri: Some(table_uri),
2898 schema: Some(Box::new(json_schema)),
2899 storage_options,
2900 properties: info.metadata.clone(),
2901 is_only_declared,
2902 ..Default::default()
2903 })
2904 }
2905 Err(err) => Err(NamespaceError::Internal {
2906 message: format!(
2907 "Table exists in manifest but failed to load dataset '{}': {}",
2908 object_id, err
2909 ),
2910 }
2911 .into()),
2912 }
2913 }
2914 None => Err(NamespaceError::TableNotFound {
2915 message: Self::format_table_id(table_id),
2916 }
2917 .into()),
2918 }
2919 }
2920
2921 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
2922 let table_id = request.id.as_ref().ok_or_else(|| {
2923 lance_core::Error::from(NamespaceError::InvalidInput {
2924 message: "Table ID is required".to_string(),
2925 })
2926 })?;
2927
2928 if table_id.is_empty() {
2929 return Err(NamespaceError::InvalidInput {
2930 message: "Table ID cannot be empty".to_string(),
2931 }
2932 .into());
2933 }
2934
2935 let object_id = Self::str_object_id(table_id);
2936 let exists = self.manifest_contains_object(&object_id).await?;
2937 if exists {
2938 Ok(())
2939 } else {
2940 Err(NamespaceError::TableNotFound {
2941 message: Self::format_table_id(table_id),
2942 }
2943 .into())
2944 }
2945 }
2946
2947 async fn create_table(
2948 &self,
2949 request: CreateTableRequest,
2950 data: Bytes,
2951 ) -> Result<CreateTableResponse> {
2952 let table_id = request.id.as_ref().ok_or_else(|| {
2953 lance_core::Error::from(NamespaceError::InvalidInput {
2954 message: "Table ID is required".to_string(),
2955 })
2956 })?;
2957
2958 if table_id.is_empty() {
2959 return Err(NamespaceError::InvalidInput {
2960 message: "Table ID cannot be empty".to_string(),
2961 }
2962 .into());
2963 }
2964
2965 let (namespace, table_name) = Self::split_object_id(table_id);
2966 let object_id = Self::build_object_id(&namespace, &table_name);
2967
2968 self.ensure_manifest_writable().await?;
2971
2972 let existing_table = self.query_manifest_for_table(&object_id).await?;
2973 let existing_has_manifests = if let Some(existing_table) = &existing_table {
2974 Some(
2975 self.location_has_actual_manifests(&existing_table.location)
2976 .await?,
2977 )
2978 } else {
2979 None
2980 };
2981
2982 if existing_has_manifests == Some(false)
2983 && request
2984 .properties
2985 .as_ref()
2986 .is_some_and(|properties| !properties.is_empty())
2987 {
2988 return Err(NamespaceError::InvalidInput {
2989 message: format!(
2990 "create_table cannot set properties for already declared table '{}'",
2991 object_id
2992 ),
2993 }
2994 .into());
2995 }
2996
2997 let create_mode = if existing_has_manifests == Some(false) {
2998 CreateTableMode::Create
2999 } else {
3000 CreateTableMode::parse(request.mode.as_deref())?
3001 };
3002 let dir_name = if let Some(existing_table) = &existing_table {
3003 existing_table.location.clone()
3004 } else if namespace.is_empty() && self.dir_listing_enabled {
3005 format!("{}.lance", table_name)
3006 } else {
3007 Self::generate_dir_name(&object_id)
3008 };
3009 let table_uri = Self::construct_full_uri(&self.root, &dir_name)?;
3010 let overwriting_existing_table =
3011 existing_has_manifests == Some(true) && create_mode == CreateTableMode::Overwrite;
3012
3013 if existing_has_manifests == Some(true) {
3014 match create_mode {
3015 CreateTableMode::Create => {
3016 return Err(NamespaceError::TableAlreadyExists {
3017 message: table_name.clone(),
3018 }
3019 .into());
3020 }
3021 CreateTableMode::ExistOk => {
3022 let properties = existing_table
3023 .as_ref()
3024 .and_then(|table| table.metadata.clone());
3025 return Ok(CreateTableResponse {
3026 location: Some(table_uri),
3027 storage_options: self.storage_options.clone(),
3028 properties,
3029 ..Default::default()
3030 });
3031 }
3032 CreateTableMode::Overwrite => {}
3033 }
3034 }
3035
3036 if data.is_empty() {
3038 return Err(NamespaceError::InvalidInput {
3039 message: "Request data (Arrow IPC stream) is required for create_table".to_string(),
3040 }
3041 .into());
3042 }
3043
3044 let cursor = Cursor::new(data.to_vec());
3046 let stream_reader = StreamReader::try_new(cursor, None).map_err(|e| {
3047 lance_core::Error::from(NamespaceError::Internal {
3048 message: format!("Failed to read IPC stream: {:?}", e),
3049 })
3050 })?;
3051
3052 let batches: Vec<RecordBatch> = stream_reader
3053 .collect::<std::result::Result<Vec<_>, _>>()
3054 .map_err(|e| {
3055 lance_core::Error::from(NamespaceError::Internal {
3056 message: format!("Failed to collect batches: {:?}", e),
3057 })
3058 })?;
3059
3060 if batches.is_empty() {
3061 return Err(NamespaceError::Internal {
3062 message: "No data provided for table creation".to_string(),
3063 }
3064 .into());
3065 }
3066
3067 let schema = batches[0].schema();
3068 let batch_results: Vec<std::result::Result<RecordBatch, arrow_schema::ArrowError>> =
3069 batches.into_iter().map(Ok).collect();
3070 let reader = RecordBatchIterator::new(batch_results, schema);
3071
3072 let mut write_storage_options = self.storage_options.clone().unwrap_or_default();
3073 if let Some(request_storage_options) = request.storage_options.as_ref() {
3074 write_storage_options.extend(request_storage_options.clone());
3075 }
3076
3077 let store_params = ObjectStoreParams {
3078 storage_options_accessor: (!write_storage_options.is_empty()).then(|| {
3079 Arc::new(
3080 lance_io::object_store::StorageOptionsAccessor::with_static_options(
3081 write_storage_options,
3082 ),
3083 )
3084 }),
3085 ..Default::default()
3086 };
3087 let write_params = WriteParams {
3088 mode: create_mode.write_mode(),
3089 session: self.session.clone(),
3090 store_params: Some(store_params),
3091 ..Default::default()
3092 };
3093 let dataset = Dataset::write(Box::new(reader), &table_uri, Some(write_params))
3094 .await
3095 .map_err(|e| {
3096 lance_core::Error::from(NamespaceError::Internal {
3097 message: format!("Failed to write dataset: {:?}", e),
3098 })
3099 })?;
3100 let version = dataset.version().version as i64;
3101
3102 if overwriting_existing_table {
3103 let metadata =
3104 Self::serialize_metadata(request.properties.as_ref(), "table", &object_id)?;
3105 self.upsert_into_manifest_with_metadata(
3106 vec![ManifestEntry {
3107 object_id,
3108 object_type: ObjectType::Table,
3109 location: Some(dir_name),
3110 metadata,
3111 }],
3112 None,
3113 )
3114 .await?;
3115
3116 Ok(CreateTableResponse {
3117 version: Some(version),
3118 location: Some(table_uri),
3119 storage_options: self.storage_options.clone(),
3120 properties: request.properties,
3121 ..Default::default()
3122 })
3123 } else {
3124 match existing_table {
3125 Some(existing_table) => Ok(CreateTableResponse {
3126 version: Some(version),
3127 location: Some(table_uri),
3128 storage_options: self.storage_options.clone(),
3129 properties: existing_table.metadata,
3130 ..Default::default()
3131 }),
3132 None => {
3133 let metadata =
3134 Self::serialize_metadata(request.properties.as_ref(), "table", &object_id)?;
3135 self.insert_into_manifest_with_metadata(
3137 vec![ManifestEntry {
3138 object_id,
3139 object_type: ObjectType::Table,
3140 location: Some(dir_name.clone()),
3141 metadata,
3142 }],
3143 None,
3144 )
3145 .await?;
3146
3147 Ok(CreateTableResponse {
3148 version: Some(version),
3149 location: Some(table_uri),
3150 storage_options: self.storage_options.clone(),
3151 properties: request.properties,
3152 ..Default::default()
3153 })
3154 }
3155 }
3156 }
3157 }
3158
3159 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3160 let table_id = request.id.as_ref().ok_or_else(|| {
3161 lance_core::Error::from(NamespaceError::InvalidInput {
3162 message: "Table ID is required".to_string(),
3163 })
3164 })?;
3165
3166 if table_id.is_empty() {
3167 return Err(NamespaceError::InvalidInput {
3168 message: "Table ID cannot be empty".to_string(),
3169 }
3170 .into());
3171 }
3172
3173 let (namespace, table_name) = Self::split_object_id(table_id);
3174 let object_id = Self::build_object_id(&namespace, &table_name);
3175
3176 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3178
3179 match table_info {
3180 Some(info) => {
3181 self.delete_from_manifest(&object_id).boxed().await?;
3183
3184 let table_path = self.base_path.clone().join(info.location.as_str());
3186 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3187
3188 self.object_store
3190 .remove_dir_all(table_path)
3191 .boxed()
3192 .await
3193 .map_err(|e| {
3194 lance_core::Error::from(NamespaceError::Internal {
3195 message: format!("Failed to delete table directory: {:?}", e),
3196 })
3197 })?;
3198
3199 Ok(DropTableResponse {
3200 id: request.id.clone(),
3201 location: Some(table_uri),
3202 ..Default::default()
3203 })
3204 }
3205 None => Err(NamespaceError::TableNotFound {
3206 message: table_name.to_string(),
3207 }
3208 .into()),
3209 }
3210 }
3211
3212 async fn list_namespaces(
3213 &self,
3214 request: ListNamespacesRequest,
3215 ) -> Result<ListNamespacesResponse> {
3216 let parent_namespace = request.id.as_ref().ok_or_else(|| {
3217 lance_core::Error::from(NamespaceError::InvalidInput {
3218 message: "Namespace ID is required".to_string(),
3219 })
3220 })?;
3221
3222 let filter = if parent_namespace.is_empty() {
3224 "object_type = 'namespace' AND NOT contains(object_id, '$')".to_string()
3226 } else {
3227 let prefix = parent_namespace.join(DELIMITER);
3229 format!(
3230 "object_type = 'namespace' AND starts_with(object_id, '{}{}') AND NOT contains(substring(object_id, {}), '$')",
3231 prefix,
3232 DELIMITER,
3233 prefix.len() + 2
3234 )
3235 };
3236
3237 let mut scanner = self.manifest_scanner().await?;
3238 scanner.filter(&filter).map_err(|e| {
3239 lance_core::Error::from(NamespaceError::Internal {
3240 message: format!("Failed to filter: {:?}", e),
3241 })
3242 })?;
3243 scanner.project(&["object_id"]).map_err(|e| {
3244 lance_core::Error::from(NamespaceError::Internal {
3245 message: format!("Failed to project: {:?}", e),
3246 })
3247 })?;
3248
3249 let batches = Self::execute_scanner(scanner).await?;
3250 let mut namespaces = Vec::new();
3251
3252 for batch in batches {
3253 if batch.num_rows() == 0 {
3254 continue;
3255 }
3256
3257 let object_id_array = Self::get_string_column(&batch, "object_id")?;
3258 for i in 0..batch.num_rows() {
3259 let object_id = object_id_array.value(i);
3260 let (_namespace, name) = Self::parse_object_id(object_id);
3261 namespaces.push(name);
3262 }
3263 }
3264
3265 let next_page_token =
3266 Self::apply_pagination(&mut namespaces, request.page_token, request.limit);
3267 let mut response = ListNamespacesResponse::new(namespaces);
3268 response.page_token = next_page_token;
3269 Ok(response)
3270 }
3271
3272 async fn describe_namespace(
3273 &self,
3274 request: DescribeNamespaceRequest,
3275 ) -> Result<DescribeNamespaceResponse> {
3276 let namespace_id = request.id.as_ref().ok_or_else(|| {
3277 lance_core::Error::from(NamespaceError::InvalidInput {
3278 message: "Namespace ID is required".to_string(),
3279 })
3280 })?;
3281
3282 if namespace_id.is_empty() {
3284 #[allow(clippy::needless_update)]
3285 return Ok(DescribeNamespaceResponse {
3286 properties: Some(HashMap::new()),
3287 ..Default::default()
3288 });
3289 }
3290
3291 let object_id = namespace_id.join(DELIMITER);
3293 let namespace_info = self.query_manifest_for_namespace(&object_id).await?;
3294
3295 match namespace_info {
3296 #[allow(clippy::needless_update)]
3297 Some(info) => Ok(DescribeNamespaceResponse {
3298 properties: info.metadata,
3299 ..Default::default()
3300 }),
3301 None => Err(NamespaceError::NamespaceNotFound {
3302 message: object_id.to_string(),
3303 }
3304 .into()),
3305 }
3306 }
3307
3308 async fn create_namespace(
3309 &self,
3310 request: CreateNamespaceRequest,
3311 ) -> Result<CreateNamespaceResponse> {
3312 let namespace_id = request.id.as_ref().ok_or_else(|| {
3313 lance_core::Error::from(NamespaceError::InvalidInput {
3314 message: "Namespace ID is required".to_string(),
3315 })
3316 })?;
3317
3318 if namespace_id.is_empty() {
3320 return Err(NamespaceError::NamespaceAlreadyExists {
3321 message: "root namespace".to_string(),
3322 }
3323 .into());
3324 }
3325
3326 if namespace_id.len() > 1 {
3328 self.validate_namespace_levels_exist(&namespace_id[..namespace_id.len() - 1])
3329 .await?;
3330 }
3331
3332 let object_id = namespace_id.join(DELIMITER);
3333 if self.manifest_contains_object(&object_id).await? {
3334 return Err(NamespaceError::NamespaceAlreadyExists {
3335 message: object_id.to_string(),
3336 }
3337 .into());
3338 }
3339
3340 let metadata =
3341 Self::serialize_metadata(request.properties.as_ref(), "namespace", &object_id)?;
3342
3343 self.insert_into_manifest_with_metadata(
3344 vec![ManifestEntry {
3345 object_id,
3346 object_type: ObjectType::Namespace,
3347 location: None,
3348 metadata,
3349 }],
3350 None,
3351 )
3352 .await?;
3353
3354 Ok(CreateNamespaceResponse {
3355 properties: request.properties,
3356 ..Default::default()
3357 })
3358 }
3359
3360 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
3361 let namespace_id = request.id.as_ref().ok_or_else(|| {
3362 lance_core::Error::from(NamespaceError::InvalidInput {
3363 message: "Namespace ID is required".to_string(),
3364 })
3365 })?;
3366
3367 if namespace_id.is_empty() {
3369 return Err(NamespaceError::InvalidInput {
3370 message: "Root namespace cannot be dropped".to_string(),
3371 }
3372 .into());
3373 }
3374
3375 let object_id = namespace_id.join(DELIMITER);
3376
3377 if !self.manifest_contains_object(&object_id).boxed().await? {
3379 return Err(NamespaceError::NamespaceNotFound {
3380 message: object_id.to_string(),
3381 }
3382 .into());
3383 }
3384
3385 let escaped_id = object_id.replace('\'', "''");
3387 let prefix = format!("{}{}", escaped_id, DELIMITER);
3388 let filter = format!("starts_with(object_id, '{}')", prefix);
3389 let mut scanner = self.manifest_scanner().boxed().await?;
3390 scanner.filter(&filter).map_err(|e| {
3391 lance_core::Error::from(NamespaceError::Internal {
3392 message: format!("Failed to filter: {:?}", e),
3393 })
3394 })?;
3395 scanner.project::<&str>(&[]).map_err(|e| {
3396 lance_core::Error::from(NamespaceError::Internal {
3397 message: format!("Failed to project: {:?}", e),
3398 })
3399 })?;
3400 scanner.with_row_id();
3401 let count = scanner.count_rows().boxed().await.map_err(|e| {
3402 lance_core::Error::from(NamespaceError::Internal {
3403 message: format!("Failed to count rows: {:?}", e),
3404 })
3405 })?;
3406
3407 if count > 0 {
3408 return Err(NamespaceError::NamespaceNotEmpty {
3409 message: format!("'{}' (contains {} child objects)", object_id, count),
3410 }
3411 .into());
3412 }
3413
3414 self.delete_from_manifest(&object_id).boxed().await?;
3415
3416 Ok(DropNamespaceResponse::default())
3417 }
3418
3419 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
3420 let namespace_id = request.id.as_ref().ok_or_else(|| {
3421 lance_core::Error::from(NamespaceError::InvalidInput {
3422 message: "Namespace ID is required".to_string(),
3423 })
3424 })?;
3425
3426 if namespace_id.is_empty() {
3428 return Ok(());
3429 }
3430
3431 let object_id = namespace_id.join(DELIMITER);
3432 if self.manifest_contains_object(&object_id).await? {
3433 Ok(())
3434 } else {
3435 Err(NamespaceError::NamespaceNotFound {
3436 message: object_id.to_string(),
3437 }
3438 .into())
3439 }
3440 }
3441
3442 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3443 let table_id = request.id.as_ref().ok_or_else(|| {
3444 lance_core::Error::from(NamespaceError::InvalidInput {
3445 message: "Table ID is required".to_string(),
3446 })
3447 })?;
3448
3449 if table_id.is_empty() {
3450 return Err(NamespaceError::InvalidInput {
3451 message: "Table ID cannot be empty".to_string(),
3452 }
3453 .into());
3454 }
3455
3456 let (namespace, table_name) = Self::split_object_id(table_id);
3457 let object_id = Self::build_object_id(&namespace, &table_name);
3458
3459 let existing = self.query_manifest_for_table(&object_id).await?;
3461 if existing.is_some() {
3462 return Err(NamespaceError::TableAlreadyExists {
3463 message: table_name.to_string(),
3464 }
3465 .into());
3466 }
3467
3468 let dir_name = if namespace.is_empty() && self.dir_listing_enabled {
3472 format!("{}.lance", table_name)
3474 } else {
3475 Self::generate_dir_name(&object_id)
3477 };
3478 let table_path = self.base_path.clone().join(dir_name.as_str());
3479 let table_uri = Self::construct_full_uri(&self.root, &dir_name)?;
3480
3481 if let Some(req_location) = &request.location {
3483 let req_location = req_location.trim_end_matches('/');
3484 if req_location != table_uri {
3485 return Err(NamespaceError::InvalidInput {
3486 message: format!(
3487 "Cannot declare table {} at location {}, must be at location {}",
3488 table_name, req_location, table_uri
3489 ),
3490 }
3491 .into());
3492 }
3493 }
3494
3495 let reserved_file_path = table_path.clone().join(".lance-reserved");
3499 super::put_marker_file_atomic(
3500 &self.object_store,
3501 &reserved_file_path,
3502 &format!("table {}", table_name),
3503 )
3504 .await
3505 .map_err(|e| match e {
3506 super::MarkerFileError::AlreadyExists { .. } => {
3507 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3508 message: table_name.to_string(),
3509 })
3510 }
3511 super::MarkerFileError::Other { message } => {
3512 lance_core::Error::from(NamespaceError::Internal { message })
3513 }
3514 })?;
3515
3516 let metadata = Self::serialize_metadata(request.properties.as_ref(), "table", &object_id)?;
3517
3518 self.insert_into_manifest_with_metadata(
3520 vec![ManifestEntry {
3521 object_id,
3522 object_type: ObjectType::Table,
3523 location: Some(dir_name),
3524 metadata,
3525 }],
3526 None,
3527 )
3528 .await?;
3529
3530 log::info!(
3531 "Declared table '{}' in manifest at {}",
3532 table_name,
3533 table_uri
3534 );
3535
3536 let vend_credentials = request.vend_credentials.unwrap_or(true);
3538 let storage_options = if vend_credentials {
3539 self.storage_options.clone()
3540 } else {
3541 None
3542 };
3543
3544 Ok(DeclareTableResponse {
3545 location: Some(table_uri),
3546 storage_options,
3547 properties: request.properties,
3548 ..Default::default()
3549 })
3550 }
3551
3552 async fn register_table(&self, request: RegisterTableRequest) -> Result<RegisterTableResponse> {
3553 let table_id = request.id.as_ref().ok_or_else(|| {
3554 lance_core::Error::from(NamespaceError::InvalidInput {
3555 message: "Table ID is required".to_string(),
3556 })
3557 })?;
3558
3559 if table_id.is_empty() {
3560 return Err(NamespaceError::InvalidInput {
3561 message: "Table ID cannot be empty".to_string(),
3562 }
3563 .into());
3564 }
3565
3566 let location = request.location.clone();
3567
3568 if location.contains("://") {
3571 return Err(NamespaceError::InvalidInput {
3572 message: format!(
3573 "Absolute URIs are not allowed for register_table. Location must be a relative path within the root directory: {}",
3574 location
3575 ),
3576 }
3577 .into());
3578 }
3579
3580 if location.starts_with('/') {
3581 return Err(NamespaceError::InvalidInput {
3582 message: format!(
3583 "Absolute paths are not allowed for register_table. Location must be a relative path within the root directory: {}",
3584 location
3585 ),
3586 }
3587 .into());
3588 }
3589
3590 if location.contains("..") {
3592 return Err(NamespaceError::InvalidInput {
3593 message: format!(
3594 "Path traversal is not allowed. Location must be a relative path within the root directory: {}",
3595 location
3596 ),
3597 }
3598 .into());
3599 }
3600
3601 let (namespace, table_name) = Self::split_object_id(table_id);
3602 let object_id = Self::build_object_id(&namespace, &table_name);
3603
3604 if !namespace.is_empty() {
3606 self.validate_namespace_levels_exist(&namespace).await?;
3607 }
3608
3609 if self.manifest_contains_object(&object_id).await? {
3611 return Err(NamespaceError::TableAlreadyExists {
3612 message: object_id.to_string(),
3613 }
3614 .into());
3615 }
3616
3617 self.insert_into_manifest(object_id, ObjectType::Table, Some(location.clone()))
3619 .await?;
3620
3621 Ok(RegisterTableResponse {
3622 location: Some(location),
3623 ..Default::default()
3624 })
3625 }
3626
3627 async fn deregister_table(
3628 &self,
3629 request: DeregisterTableRequest,
3630 ) -> Result<DeregisterTableResponse> {
3631 let table_id = request.id.as_ref().ok_or_else(|| {
3632 lance_core::Error::from(NamespaceError::InvalidInput {
3633 message: "Table ID is required".to_string(),
3634 })
3635 })?;
3636
3637 if table_id.is_empty() {
3638 return Err(NamespaceError::InvalidInput {
3639 message: "Table ID cannot be empty".to_string(),
3640 }
3641 .into());
3642 }
3643
3644 let (namespace, table_name) = Self::split_object_id(table_id);
3645 let object_id = Self::build_object_id(&namespace, &table_name);
3646
3647 let table_info = self.query_manifest_for_table(&object_id).await?;
3649
3650 let table_uri = match table_info {
3651 Some(info) => {
3652 self.delete_from_manifest(&object_id).boxed().await?;
3654 Self::construct_full_uri(&self.root, &info.location)?
3655 }
3656 None => {
3657 return Err(NamespaceError::TableNotFound {
3658 message: object_id.to_string(),
3659 }
3660 .into());
3661 }
3662 };
3663
3664 Ok(DeregisterTableResponse {
3665 id: request.id.clone(),
3666 location: Some(table_uri),
3667 ..Default::default()
3668 })
3669 }
3670
3671 async fn alter_table_add_columns(
3676 &self,
3677 request: AlterTableAddColumnsRequest,
3678 ) -> Result<AlterTableAddColumnsResponse> {
3679 let table_id = request
3680 .id
3681 .as_ref()
3682 .ok_or_else(|| Error::invalid_input_source("Table ID is required".into()))?;
3683
3684 if table_id.is_empty() {
3685 return Err(Error::invalid_input_source(
3686 "Table ID cannot be empty".into(),
3687 ));
3688 }
3689
3690 let object_id = Self::str_object_id(table_id);
3691 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3692
3693 match table_info {
3694 Some(info) => {
3695 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3696 let mut builder = DatasetBuilder::from_uri(&table_uri);
3699 if let Some(opts) = &self.storage_options {
3700 builder = builder.with_storage_options(opts.clone());
3701 }
3702 if let Some(session) = &self.session {
3703 builder = builder.with_session(session.clone());
3704 }
3705 let mut dataset = builder.load().await.map_err(|e| {
3706 Error::io_source(box_error(std::io::Error::other(format!(
3707 "Failed to open dataset: {}",
3708 e
3709 ))))
3710 })?;
3711
3712 let sql_expressions = super::build_sql_expressions(&request.new_columns)?;
3714
3715 dataset
3716 .add_columns(
3717 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3718 None,
3719 None,
3720 )
3721 .await
3722 .map_err(|e| {
3723 convert_lance_commit_error(&e, "add_columns", Some(&object_id))
3727 })?;
3728
3729 let version = dataset.version().version as i64;
3730 Ok(AlterTableAddColumnsResponse::new(version))
3731 }
3732 None => Err(NamespaceError::TableNotFound { message: object_id }.into()),
3733 }
3734 }
3735
3736 async fn alter_table_alter_columns(
3741 &self,
3742 request: AlterTableAlterColumnsRequest,
3743 ) -> Result<AlterTableAlterColumnsResponse> {
3744 let table_id = request
3745 .id
3746 .as_ref()
3747 .ok_or_else(|| Error::invalid_input_source("Table ID is required".into()))?;
3748
3749 if table_id.is_empty() {
3750 return Err(Error::invalid_input_source(
3751 "Table ID cannot be empty".into(),
3752 ));
3753 }
3754
3755 let object_id = Self::str_object_id(table_id);
3756 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3757
3758 match table_info {
3759 Some(info) => {
3760 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3761 let mut builder = DatasetBuilder::from_uri(&table_uri);
3762 if let Some(opts) = &self.storage_options {
3763 builder = builder.with_storage_options(opts.clone());
3764 }
3765 if let Some(session) = &self.session {
3766 builder = builder.with_session(session.clone());
3767 }
3768 let mut dataset = builder.load().await.map_err(|e| {
3769 Error::io_source(box_error(std::io::Error::other(format!(
3770 "Failed to open dataset: {}",
3771 e
3772 ))))
3773 })?;
3774
3775 let alterations = super::build_column_alterations(&request.alterations)?;
3777
3778 dataset.alter_columns(&alterations).await.map_err(|e| {
3779 convert_lance_commit_error(&e, "alter_columns", Some(&object_id))
3780 })?;
3781
3782 let version = dataset.version().version as i64;
3783 Ok(AlterTableAlterColumnsResponse::new(version))
3784 }
3785 None => Err(NamespaceError::TableNotFound { message: object_id }.into()),
3786 }
3787 }
3788
3789 async fn alter_table_drop_columns(
3793 &self,
3794 request: AlterTableDropColumnsRequest,
3795 ) -> Result<AlterTableDropColumnsResponse> {
3796 let table_id = request
3797 .id
3798 .as_ref()
3799 .ok_or_else(|| Error::invalid_input_source("Table ID is required".into()))?;
3800
3801 if table_id.is_empty() {
3802 return Err(Error::invalid_input_source(
3803 "Table ID cannot be empty".into(),
3804 ));
3805 }
3806
3807 let object_id = Self::str_object_id(table_id);
3808 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3809
3810 match table_info {
3811 Some(info) => {
3812 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3813 let mut builder = DatasetBuilder::from_uri(&table_uri);
3814 if let Some(opts) = &self.storage_options {
3815 builder = builder.with_storage_options(opts.clone());
3816 }
3817 if let Some(session) = &self.session {
3818 builder = builder.with_session(session.clone());
3819 }
3820 let mut dataset = builder.load().await.map_err(|e| {
3821 Error::io_source(box_error(std::io::Error::other(format!(
3822 "Failed to open dataset: {}",
3823 e
3824 ))))
3825 })?;
3826
3827 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
3828 dataset.drop_columns(&columns).await.map_err(|e| {
3829 convert_lance_commit_error(&e, "drop_columns", Some(&object_id))
3830 })?;
3831
3832 let version = dataset.version().version as i64;
3833 Ok(AlterTableDropColumnsResponse::new(version))
3834 }
3835 None => Err(NamespaceError::TableNotFound { message: object_id }.into()),
3836 }
3837 }
3838}
3839
3840#[cfg(test)]
3841mod tests {
3842 use super::{
3843 BASE_OBJECTS_INDEX_NAME, ConflictResolution, CopyOnWriteMutation, DeleteObjectMutation,
3844 LANCE_DATA_DIR, LANCE_INDICES_DIR, MANIFEST_TABLE_NAME, ManifestBatchBuilder,
3845 ManifestEntry, ManifestIndexAccumulator, ManifestNamespace, ManifestOutputRow,
3846 ManifestRowValue, ManifestStreamMutation, OBJECT_ID_INDEX_NAME, OBJECT_TYPE_INDEX_NAME,
3847 ObjectType,
3848 };
3849 use crate::DirectoryNamespaceBuilder;
3850 use arrow::datatypes::DataType;
3851 use bytes::Bytes;
3852 use futures::StreamExt;
3853 use lance::index::DatasetIndexExt;
3854 use lance_core::utils::tempfile::TempStdDir;
3855 use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
3856 use lance_namespace::LanceNamespace;
3857 use lance_namespace::models::{
3858 CreateNamespaceRequest, CreateTableRequest, DescribeTableRequest, DropTableRequest,
3859 ListTablesRequest, TableExistsRequest,
3860 };
3861 use lance_table::format::Fragment;
3862 use rstest::rstest;
3863 use std::collections::{HashMap, HashSet};
3864 use std::sync::Arc;
3865
3866 async fn create_manifest_namespace(
3867 root: &str,
3868 inline_optimization_enabled: bool,
3869 ) -> ManifestNamespace {
3870 create_manifest_namespace_with_retries(root, inline_optimization_enabled, None).await
3871 }
3872
3873 async fn create_manifest_namespace_with_retries(
3874 root: &str,
3875 inline_optimization_enabled: bool,
3876 commit_retries: Option<u32>,
3877 ) -> ManifestNamespace {
3878 let (object_store, base_path) = ObjectStore::from_uri_and_params(
3879 Arc::new(ObjectStoreRegistry::default()),
3880 root,
3881 &ObjectStoreParams::default(),
3882 )
3883 .await
3884 .unwrap();
3885 ManifestNamespace::from_directory(
3886 root.to_string(),
3887 None,
3888 None,
3889 object_store,
3890 base_path,
3891 true,
3892 inline_optimization_enabled,
3893 commit_retries,
3894 )
3895 .await
3896 .unwrap()
3897 }
3898
3899 struct CommitConflictAfterRewriteMutation {
3900 root: String,
3901 conflict_object_id: String,
3902 }
3903
3904 impl ManifestStreamMutation for CommitConflictAfterRewriteMutation {
3905 type Output = ();
3906
3907 fn process_existing_row(
3908 &mut self,
3909 row: ManifestRowValue,
3910 output: &mut ManifestBatchBuilder,
3911 index_data: &mut ManifestIndexAccumulator,
3912 ) -> lance_core::Result<()> {
3913 output.append(
3914 index_data,
3915 ManifestOutputRow {
3916 object_id: &row.object_id,
3917 object_type: row.object_type,
3918 location: row.location.as_deref(),
3919 metadata: row.metadata.as_deref(),
3920 base_objects: row.base_objects.as_deref(),
3921 },
3922 )
3923 }
3924
3925 fn append_rows(
3926 &mut self,
3927 output: &mut ManifestBatchBuilder,
3928 index_data: &mut ManifestIndexAccumulator,
3929 ) -> lance_core::Result<()> {
3930 output.append(
3931 index_data,
3932 ManifestOutputRow {
3933 object_id: "attempted_table",
3934 object_type: ObjectType::Table,
3935 location: Some("attempted_table.lance"),
3936 metadata: None,
3937 base_objects: None,
3938 },
3939 )
3940 }
3941
3942 fn finish(&self) -> CopyOnWriteMutation<Self::Output> {
3943 let root = self.root.clone();
3944 let object_id = self.conflict_object_id.clone();
3945 std::thread::spawn(move || {
3946 let runtime = tokio::runtime::Runtime::new().unwrap();
3947 runtime.block_on(async move {
3948 let writer = create_manifest_namespace(&root, false).await;
3949 writer
3950 .insert_into_manifest_with_metadata(
3951 vec![ManifestEntry {
3952 object_id,
3953 object_type: ObjectType::Table,
3954 location: Some("conflicting_table.lance".to_string()),
3955 metadata: None,
3956 }],
3957 None,
3958 )
3959 .await
3960 .unwrap();
3961 });
3962 })
3963 .join()
3964 .unwrap();
3965 CopyOnWriteMutation::updated(())
3966 }
3967 }
3968
3969 struct ConcurrentDeleteBeforeCommitMutation {
3973 inner: DeleteObjectMutation,
3974 root: String,
3975 target: String,
3976 }
3977
3978 impl ManifestStreamMutation for ConcurrentDeleteBeforeCommitMutation {
3979 type Output = ();
3980
3981 fn process_existing_row(
3982 &mut self,
3983 row: ManifestRowValue,
3984 output: &mut ManifestBatchBuilder,
3985 index_data: &mut ManifestIndexAccumulator,
3986 ) -> lance_core::Result<()> {
3987 self.inner.process_existing_row(row, output, index_data)
3988 }
3989
3990 fn append_rows(
3991 &mut self,
3992 output: &mut ManifestBatchBuilder,
3993 index_data: &mut ManifestIndexAccumulator,
3994 ) -> lance_core::Result<()> {
3995 self.inner.append_rows(output, index_data)
3996 }
3997
3998 fn finish(&self) -> CopyOnWriteMutation<Self::Output> {
3999 let root = self.root.clone();
4000 let target = self.target.clone();
4001 std::thread::spawn(move || {
4002 let runtime = tokio::runtime::Runtime::new().unwrap();
4003 runtime.block_on(async move {
4004 let writer = create_manifest_namespace(&root, false).await;
4005 writer.delete_from_manifest(&target).await.unwrap();
4006 });
4007 })
4008 .join()
4009 .unwrap();
4010 self.inner.finish()
4011 }
4012
4013 fn conflict_resolution(&self) -> ConflictResolution<Self::Output> {
4014 ConflictResolution::SucceedIfAbsent {
4015 object_id: self.target.clone(),
4016 output: (),
4017 }
4018 }
4019 }
4020
4021 async fn manifest_base_objects(
4022 manifest_ns: &ManifestNamespace,
4023 ) -> HashMap<String, Option<Vec<String>>> {
4024 let mut scanner = manifest_ns.manifest_scanner().await.unwrap();
4025 scanner.project(&["object_id", "base_objects"]).unwrap();
4026 let batches = ManifestNamespace::execute_scanner(scanner).await.unwrap();
4027 let mut rows = HashMap::new();
4028 for batch in batches {
4029 let object_ids = ManifestNamespace::get_string_column(&batch, "object_id").unwrap();
4030 let base_objects = ManifestNamespace::base_objects_column_values(&batch).unwrap();
4031 for (row, value) in base_objects.into_iter().enumerate() {
4032 rows.insert(object_ids.value(row).to_string(), value);
4033 }
4034 }
4035 rows
4036 }
4037
4038 async fn manifest_data_paths(manifest_ns: &ManifestNamespace) -> HashSet<String> {
4039 let data_dir = manifest_ns
4040 .base_path
4041 .clone()
4042 .join(MANIFEST_TABLE_NAME)
4043 .join(LANCE_DATA_DIR);
4044 let mut stream = manifest_ns.object_store.read_dir_all(&data_dir, None);
4045 let mut paths = HashSet::new();
4046 while let Some(meta) = stream.next().await.transpose().unwrap() {
4047 paths.insert(meta.location.to_string());
4048 }
4049 paths
4050 }
4051
4052 async fn manifest_index_paths(manifest_ns: &ManifestNamespace) -> HashSet<String> {
4053 let index_dir = manifest_ns
4054 .base_path
4055 .clone()
4056 .join(MANIFEST_TABLE_NAME)
4057 .join(LANCE_INDICES_DIR);
4058 let mut stream = manifest_ns.object_store.read_dir_all(&index_dir, None);
4059 let mut paths = HashSet::new();
4060 while let Some(meta) = stream.next().await.transpose().unwrap() {
4061 paths.insert(meta.location.to_string());
4062 }
4063 paths
4064 }
4065
4066 fn create_test_ipc_data() -> Vec<u8> {
4067 use arrow::array::{Int32Array, StringArray};
4068 use arrow::datatypes::{DataType, Field, Schema};
4069 use arrow::ipc::writer::StreamWriter;
4070 use arrow::record_batch::RecordBatch;
4071 use std::sync::Arc;
4072
4073 let schema = Arc::new(Schema::new(vec![
4074 Field::new("id", DataType::Int32, false),
4075 Field::new("name", DataType::Utf8, false),
4076 ]));
4077
4078 let batch = RecordBatch::try_new(
4079 schema.clone(),
4080 vec![
4081 Arc::new(Int32Array::from(vec![1, 2, 3])),
4082 Arc::new(StringArray::from(vec!["a", "b", "c"])),
4083 ],
4084 )
4085 .unwrap();
4086
4087 let mut buffer = Vec::new();
4088 {
4089 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
4090 writer.write(&batch).unwrap();
4091 writer.finish().unwrap();
4092 }
4093 buffer
4094 }
4095
4096 async fn set_manifest_table_metadata(temp_path: &str, key: &str, value: &str) {
4099 use lance::dataset::builder::DatasetBuilder;
4100 let mut ds = DatasetBuilder::from_uri(format!("{}/{}", temp_path, MANIFEST_TABLE_NAME))
4101 .load()
4102 .await
4103 .unwrap();
4104 ds.update_metadata([(key, value)]).await.unwrap();
4105 }
4106
4107 async fn create_namespace_with_one_table(temp_path: &str) {
4108 let ns = DirectoryNamespaceBuilder::new(temp_path)
4109 .build()
4110 .await
4111 .unwrap();
4112 let mut create_request = CreateTableRequest::new();
4113 create_request.id = Some(vec!["t1".to_string()]);
4114 ns.create_table(create_request, Bytes::from(create_test_ipc_data()))
4115 .await
4116 .unwrap();
4117 }
4118
4119 #[tokio::test]
4122 async fn test_manifest_has_no_feature_flags_by_default() {
4123 use lance::dataset::builder::DatasetBuilder;
4124 let temp_dir = TempStdDir::default();
4125 let temp_path = temp_dir.to_str().unwrap();
4126 create_namespace_with_one_table(temp_path).await;
4127
4128 let ds = DatasetBuilder::from_uri(format!("{}/{}", temp_path, MANIFEST_TABLE_NAME))
4129 .load()
4130 .await
4131 .unwrap();
4132 assert!(
4133 !ds.metadata()
4134 .contains_key(crate::dir::manifest_feature_flags::READER_FEATURE_FLAGS_KEY)
4135 );
4136 assert!(
4137 !ds.metadata()
4138 .contains_key(crate::dir::manifest_feature_flags::WRITER_FEATURE_FLAGS_KEY)
4139 );
4140 }
4141
4142 #[tokio::test]
4145 async fn test_unknown_reader_flag_blocks_access() {
4146 let temp_dir = TempStdDir::default();
4147 let temp_path = temp_dir.to_str().unwrap();
4148 create_namespace_with_one_table(temp_path).await;
4149 set_manifest_table_metadata(
4150 temp_path,
4151 crate::dir::manifest_feature_flags::READER_FEATURE_FLAGS_KEY,
4152 "1",
4153 )
4154 .await;
4155
4156 let err = DirectoryNamespaceBuilder::new(temp_path)
4157 .build()
4158 .await
4159 .expect_err("opening a manifest with an unknown reader flag should fail");
4160 assert!(
4161 err.to_string().to_lowercase().contains("upgrade"),
4162 "expected an upgrade error, got: {err}"
4163 );
4164 }
4165
4166 #[tokio::test]
4168 async fn test_unknown_writer_flag_blocks_writes_but_allows_reads() {
4169 let temp_dir = TempStdDir::default();
4170 let temp_path = temp_dir.to_str().unwrap();
4171 create_namespace_with_one_table(temp_path).await;
4172 set_manifest_table_metadata(
4173 temp_path,
4174 crate::dir::manifest_feature_flags::WRITER_FEATURE_FLAGS_KEY,
4175 "1",
4176 )
4177 .await;
4178
4179 let ns = DirectoryNamespaceBuilder::new(temp_path)
4180 .build()
4181 .await
4182 .expect("reads should still be allowed with only a writer flag set");
4183 let mut list_request = ListTablesRequest::new();
4184 list_request.id = Some(vec![]);
4185 assert_eq!(ns.list_tables(list_request).await.unwrap().tables.len(), 1);
4186
4187 let entries_before = dir_entry_names(temp_path);
4189 let mut create_request = CreateTableRequest::new();
4190 create_request.id = Some(vec!["t2".to_string()]);
4191 let err = ns
4192 .create_table(create_request, Bytes::from(create_test_ipc_data()))
4193 .await
4194 .expect_err("writing through an unknown writer flag should fail");
4195 assert!(
4196 err.to_string().to_lowercase().contains("upgrade"),
4197 "expected an upgrade error, got: {err}"
4198 );
4199 assert_eq!(
4200 entries_before,
4201 dir_entry_names(temp_path),
4202 "a refused create_table must not create an orphaned table directory"
4203 );
4204
4205 let mut create_ns = CreateNamespaceRequest::new();
4210 create_ns.id = Some(vec!["ns1".to_string()]);
4211 let err = ns
4212 .create_namespace(create_ns)
4213 .await
4214 .expect_err("create_namespace through an unknown writer flag should fail");
4215 assert!(
4216 err.to_string().to_lowercase().contains("upgrade"),
4217 "expected an upgrade error, got: {err}"
4218 );
4219
4220 let mut drop_request = DropTableRequest::new();
4221 drop_request.id = Some(vec!["t1".to_string()]);
4222 let err = ns
4223 .drop_table(drop_request)
4224 .await
4225 .expect_err("drop_table through an unknown writer flag should fail");
4226 assert!(
4227 err.to_string().to_lowercase().contains("upgrade"),
4228 "expected an upgrade error, got: {err}"
4229 );
4230 }
4231
4232 fn dir_entry_names(path: &str) -> std::collections::BTreeSet<String> {
4233 std::fs::read_dir(path)
4234 .unwrap()
4235 .map(|e| e.unwrap().file_name().to_string_lossy().into_owned())
4236 .collect()
4237 }
4238
4239 #[tokio::test]
4240 async fn test_manifest_rewrite_preserves_utf8_metadata_and_base_objects() {
4241 let temp_dir = TempStdDir::default();
4242 let temp_path = temp_dir.to_str().unwrap();
4243 let manifest_ns = create_manifest_namespace(temp_path, true).await;
4244
4245 manifest_ns
4246 .insert_into_manifest_with_metadata(
4247 vec![ManifestEntry {
4248 object_id: "view".to_string(),
4249 object_type: ObjectType::Table,
4250 location: Some("view.lance".to_string()),
4251 metadata: Some(r#"{"kind":"view"}"#.to_string()),
4252 }],
4253 Some(vec!["base_a".to_string(), "base_b".to_string()]),
4254 )
4255 .await
4256 .unwrap();
4257 manifest_ns
4258 .insert_into_manifest_with_metadata(
4259 vec![ManifestEntry {
4260 object_id: "other".to_string(),
4261 object_type: ObjectType::Namespace,
4262 location: None,
4263 metadata: Some(r#"{"kind":"namespace"}"#.to_string()),
4264 }],
4265 None,
4266 )
4267 .await
4268 .unwrap();
4269
4270 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4271 let metadata_field = dataset_guard.schema().field("metadata").unwrap();
4272 assert_eq!(metadata_field.data_type(), DataType::Utf8);
4273 drop(dataset_guard);
4274
4275 let base_objects = manifest_base_objects(&manifest_ns).await;
4276 assert_eq!(
4277 base_objects.get("view").cloned().unwrap(),
4278 Some(vec!["base_a".to_string(), "base_b".to_string()])
4279 );
4280 assert_eq!(base_objects.get("other").cloned().unwrap(), None);
4281 }
4282
4283 #[tokio::test]
4284 async fn test_manifest_rewrite_replacement_indices_are_versioned() {
4285 let temp_dir = TempStdDir::default();
4286 let temp_path = temp_dir.to_str().unwrap();
4287 let manifest_ns = create_manifest_namespace(temp_path, true).await;
4288
4289 manifest_ns
4290 .insert_into_manifest_with_metadata(
4291 vec![ManifestEntry {
4292 object_id: "table".to_string(),
4293 object_type: ObjectType::Table,
4294 location: Some("table.lance".to_string()),
4295 metadata: None,
4296 }],
4297 Some(vec!["base".to_string()]),
4298 )
4299 .await
4300 .unwrap();
4301
4302 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4303 let dataset_version = dataset_guard.version().version;
4304 let indices = dataset_guard.load_indices().await.unwrap();
4305 let names = indices
4306 .iter()
4307 .map(|index| index.name.as_str())
4308 .collect::<HashSet<_>>();
4309 assert!(names.contains(OBJECT_ID_INDEX_NAME));
4310 assert!(names.contains(OBJECT_TYPE_INDEX_NAME));
4311 assert!(names.contains(BASE_OBJECTS_INDEX_NAME));
4312 for index in indices.iter() {
4313 assert_eq!(index.dataset_version, dataset_version);
4314 assert!(!index.fragment_bitmap.as_ref().unwrap().is_empty());
4315 }
4316 }
4317
4318 #[tokio::test]
4319 async fn test_manifest_rewrite_empty_manifest_keeps_replacement_indices_valid() {
4320 let temp_dir = TempStdDir::default();
4321 let temp_path = temp_dir.to_str().unwrap();
4322 let manifest_ns = create_manifest_namespace(temp_path, true).await;
4323
4324 manifest_ns
4325 .insert_into_manifest_with_metadata(
4326 vec![ManifestEntry {
4327 object_id: "table".to_string(),
4328 object_type: ObjectType::Table,
4329 location: Some("table.lance".to_string()),
4330 metadata: None,
4331 }],
4332 None,
4333 )
4334 .await
4335 .unwrap();
4336 manifest_ns.delete_from_manifest("table").await.unwrap();
4337
4338 assert!(!manifest_ns.manifest_contains_object("table").await.unwrap());
4339 let mut scanner = manifest_ns.manifest_scanner().await.unwrap();
4340 scanner.project(&["object_id"]).unwrap();
4341 let rows = ManifestNamespace::execute_scanner(scanner)
4342 .await
4343 .unwrap()
4344 .into_iter()
4345 .map(|batch| batch.num_rows())
4346 .sum::<usize>();
4347 assert_eq!(rows, 0);
4348
4349 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4350 let dataset_version = dataset_guard.version().version;
4351 let indices = dataset_guard.load_indices().await.unwrap();
4352 let names = indices
4353 .iter()
4354 .map(|index| index.name.as_str())
4355 .collect::<HashSet<_>>();
4356 assert!(names.contains(OBJECT_ID_INDEX_NAME));
4357 assert!(names.contains(OBJECT_TYPE_INDEX_NAME));
4358 assert!(names.contains(BASE_OBJECTS_INDEX_NAME));
4359 for index in indices.iter() {
4360 assert_eq!(index.dataset_version, dataset_version);
4361 }
4362 }
4363
4364 #[tokio::test]
4365 async fn test_manifest_rewrite_fragment_bitmap_uses_overwrite_fragment_ids() {
4366 let temp_dir = TempStdDir::default();
4367 let temp_path = temp_dir.to_str().unwrap();
4368 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4369 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4370 let fragments = vec![Fragment::new(0), Fragment::new(0), Fragment::new(7)];
4371
4372 let manifest = ManifestNamespace::manifest_from_overwrite_transaction(
4373 dataset_guard.manifest(),
4374 dataset_guard.manifest().schema.clone(),
4375 &fragments,
4376 );
4377
4378 let fragment_ids = manifest
4379 .fragments
4380 .iter()
4381 .map(|fragment| fragment.id)
4382 .collect::<Vec<_>>();
4383 assert_eq!(fragment_ids, vec![0, 1, 7]);
4384 assert_eq!(
4385 ManifestNamespace::manifest_fragment_bitmap(&manifest)
4386 .unwrap()
4387 .into_iter()
4388 .collect::<Vec<_>>(),
4389 vec![0, 1, 7]
4390 );
4391 }
4392
4393 #[tokio::test]
4394 async fn test_manifest_noop_delete_uses_latest_snapshot() {
4395 let temp_dir = TempStdDir::default();
4396 let temp_path = temp_dir.to_str().unwrap();
4397 let stale_ns = create_manifest_namespace(temp_path, false).await;
4398 let writer_ns = create_manifest_namespace(temp_path, false).await;
4399
4400 writer_ns
4401 .insert_into_manifest_with_metadata(
4402 vec![ManifestEntry {
4403 object_id: "late_table".to_string(),
4404 object_type: ObjectType::Table,
4405 location: Some("late_table.lance".to_string()),
4406 metadata: None,
4407 }],
4408 None,
4409 )
4410 .await
4411 .unwrap();
4412
4413 stale_ns.delete_from_manifest("late_table").await.unwrap();
4414
4415 let check_ns = create_manifest_namespace(temp_path, false).await;
4416 assert!(
4417 !check_ns
4418 .manifest_contains_object("late_table")
4419 .await
4420 .unwrap()
4421 );
4422 }
4423
4424 #[tokio::test]
4425 async fn test_manifest_noop_delete_cleans_uncommitted_data_file() {
4426 let temp_dir = TempStdDir::default();
4427 let temp_path = temp_dir.to_str().unwrap();
4428 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4429
4430 manifest_ns
4431 .insert_into_manifest_with_metadata(
4432 vec![ManifestEntry {
4433 object_id: "table".to_string(),
4434 object_type: ObjectType::Table,
4435 location: Some("table.lance".to_string()),
4436 metadata: None,
4437 }],
4438 None,
4439 )
4440 .await
4441 .unwrap();
4442
4443 let before = manifest_data_paths(&manifest_ns).await;
4444 assert!(!before.is_empty());
4445
4446 manifest_ns
4447 .delete_from_manifest("missing_table")
4448 .await
4449 .unwrap();
4450
4451 let after = manifest_data_paths(&manifest_ns).await;
4452 assert_eq!(after, before);
4453 }
4454
4455 #[tokio::test]
4456 async fn test_manifest_final_commit_failure_cleans_uncommitted_rewrite_files() {
4457 let temp_dir = TempStdDir::default();
4458 let temp_path = temp_dir.to_str().unwrap();
4459 let manifest_ns = create_manifest_namespace_with_retries(temp_path, true, Some(0)).await;
4460
4461 manifest_ns
4462 .insert_into_manifest_with_metadata(
4463 vec![ManifestEntry {
4464 object_id: "table".to_string(),
4465 object_type: ObjectType::Table,
4466 location: Some("table.lance".to_string()),
4467 metadata: None,
4468 }],
4469 None,
4470 )
4471 .await
4472 .unwrap();
4473
4474 let before_data_paths = manifest_data_paths(&manifest_ns).await;
4475 let before_index_paths = manifest_index_paths(&manifest_ns).await;
4476
4477 let result = manifest_ns
4478 .rewrite_manifest("Failed to test manifest cleanup", || {
4479 CommitConflictAfterRewriteMutation {
4480 root: temp_path.to_string(),
4481 conflict_object_id: "conflicting_table".to_string(),
4482 }
4483 })
4484 .await;
4485 assert!(result.is_err());
4486
4487 let after_data_paths = manifest_data_paths(&manifest_ns).await;
4488 assert!(before_data_paths.is_subset(&after_data_paths));
4489 assert_eq!(after_data_paths.len(), before_data_paths.len() + 1);
4490 assert_eq!(manifest_index_paths(&manifest_ns).await, before_index_paths);
4491 assert!(
4492 manifest_ns
4493 .manifest_contains_object("conflicting_table")
4494 .await
4495 .unwrap()
4496 );
4497 assert!(
4498 !manifest_ns
4499 .manifest_contains_object("attempted_table")
4500 .await
4501 .unwrap()
4502 );
4503 }
4504
4505 #[tokio::test]
4506 async fn test_manifest_commit_visible_on_memory_store() {
4507 let manifest_ns = create_manifest_namespace("memory://test_commit_visible", false).await;
4512 manifest_ns
4513 .insert_into_manifest_with_metadata(
4514 vec![ManifestEntry {
4515 object_id: "table".to_string(),
4516 object_type: ObjectType::Table,
4517 location: Some("table.lance".to_string()),
4518 metadata: None,
4519 }],
4520 None,
4521 )
4522 .await
4523 .unwrap();
4524 assert!(manifest_ns.manifest_contains_object("table").await.unwrap());
4525 manifest_ns
4527 .insert_into_manifest_with_metadata(
4528 vec![ManifestEntry {
4529 object_id: "table2".to_string(),
4530 object_type: ObjectType::Table,
4531 location: Some("table2.lance".to_string()),
4532 metadata: None,
4533 }],
4534 None,
4535 )
4536 .await
4537 .unwrap();
4538 assert!(
4539 manifest_ns
4540 .manifest_contains_object("table2")
4541 .await
4542 .unwrap()
4543 );
4544 }
4545
4546 #[tokio::test]
4547 async fn test_manifest_commit_uses_inline_transaction() {
4548 let temp_dir = TempStdDir::default();
4549 let temp_path = temp_dir.to_str().unwrap();
4550 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4551
4552 manifest_ns
4553 .insert_into_manifest_with_metadata(
4554 vec![ManifestEntry {
4555 object_id: "table".to_string(),
4556 object_type: ObjectType::Table,
4557 location: Some("table.lance".to_string()),
4558 metadata: None,
4559 }],
4560 None,
4561 )
4562 .await
4563 .unwrap();
4564
4565 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4566 let manifest = dataset_guard.manifest();
4567 assert!(manifest.transaction_section.is_some());
4570 assert!(manifest.transaction_file.is_none());
4571 }
4572
4573 #[tokio::test]
4574 async fn test_manifest_commit_landed_attributes_data_file() {
4575 let temp_dir = TempStdDir::default();
4576 let temp_path = temp_dir.to_str().unwrap();
4577 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4578
4579 manifest_ns
4580 .insert_into_manifest_with_metadata(
4581 vec![ManifestEntry {
4582 object_id: "table".to_string(),
4583 object_type: ObjectType::Table,
4584 location: Some("table.lance".to_string()),
4585 metadata: None,
4586 }],
4587 None,
4588 )
4589 .await
4590 .unwrap();
4591
4592 let dataset = Arc::new(manifest_ns.manifest_dataset.get().await.unwrap().clone());
4593 let version = dataset.manifest().version;
4594 let our_files = dataset
4595 .manifest()
4596 .fragments
4597 .iter()
4598 .flat_map(|fragment| fragment.files.iter())
4599 .map(|file| file.path.clone())
4600 .collect::<HashSet<_>>();
4601 assert!(!our_files.is_empty());
4602
4603 assert!(
4606 manifest_ns
4607 .manifest_commit_landed(&dataset, version, &our_files)
4608 .await
4609 );
4610 let other = HashSet::from(["missing.lance".to_string()]);
4612 assert!(
4613 !manifest_ns
4614 .manifest_commit_landed(&dataset, version, &other)
4615 .await
4616 );
4617 assert!(
4619 !manifest_ns
4620 .manifest_commit_landed(&dataset, version + 100, &our_files)
4621 .await
4622 );
4623 }
4624
4625 #[tokio::test]
4626 async fn test_manifest_delete_conflict_with_concurrent_delete_succeeds() {
4627 let temp_dir = TempStdDir::default();
4628 let temp_path = temp_dir.to_str().unwrap();
4629 let manifest_ns = create_manifest_namespace_with_retries(temp_path, false, Some(0)).await;
4630
4631 manifest_ns
4632 .insert_into_manifest_with_metadata(
4633 vec![ManifestEntry {
4634 object_id: "table".to_string(),
4635 object_type: ObjectType::Table,
4636 location: Some("table.lance".to_string()),
4637 metadata: None,
4638 }],
4639 None,
4640 )
4641 .await
4642 .unwrap();
4643 assert!(manifest_ns.manifest_contains_object("table").await.unwrap());
4644
4645 let result = manifest_ns
4649 .rewrite_manifest("Failed to delete from manifest", || {
4650 ConcurrentDeleteBeforeCommitMutation {
4651 inner: DeleteObjectMutation {
4652 object_id: "table".to_string(),
4653 deleted: false,
4654 },
4655 root: temp_path.to_string(),
4656 target: "table".to_string(),
4657 }
4658 })
4659 .await;
4660
4661 assert!(result.is_ok(), "delete should succeed: {result:?}");
4662 assert!(!manifest_ns.manifest_contains_object("table").await.unwrap());
4663 }
4664
4665 #[rstest]
4666 #[case::with_optimization(true)]
4667 #[case::without_optimization(false)]
4668 #[tokio::test]
4669 async fn test_manifest_namespace_basic_create_and_list(#[case] inline_optimization: bool) {
4670 let temp_dir = TempStdDir::default();
4671 let temp_path = temp_dir.to_str().unwrap();
4672
4673 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4675 .inline_optimization_enabled(inline_optimization)
4676 .build()
4677 .await
4678 .unwrap();
4679
4680 let mut request = ListTablesRequest::new();
4682 request.id = Some(vec![]);
4683 let response = dir_namespace.list_tables(request).await.unwrap();
4684 assert_eq!(response.tables.len(), 0);
4685
4686 let buffer = create_test_ipc_data();
4688 let mut create_request = CreateTableRequest::new();
4689 create_request.id = Some(vec!["test_table".to_string()]);
4690
4691 let _response = dir_namespace
4692 .create_table(create_request, Bytes::from(buffer))
4693 .await
4694 .unwrap();
4695
4696 let mut request = ListTablesRequest::new();
4698 request.id = Some(vec![]);
4699 let response = dir_namespace.list_tables(request).await.unwrap();
4700 assert_eq!(response.tables.len(), 1);
4701 assert_eq!(response.tables[0], "test_table");
4702 }
4703
4704 #[rstest]
4705 #[case::with_optimization(true)]
4706 #[case::without_optimization(false)]
4707 #[tokio::test]
4708 async fn test_manifest_namespace_table_exists(#[case] inline_optimization: bool) {
4709 let temp_dir = TempStdDir::default();
4710 let temp_path = temp_dir.to_str().unwrap();
4711
4712 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4713 .inline_optimization_enabled(inline_optimization)
4714 .build()
4715 .await
4716 .unwrap();
4717
4718 let mut request = TableExistsRequest::new();
4720 request.id = Some(vec!["nonexistent".to_string()]);
4721 let result = dir_namespace.table_exists(request).await;
4722 assert!(result.is_err());
4723
4724 let buffer = create_test_ipc_data();
4726 let mut create_request = CreateTableRequest::new();
4727 create_request.id = Some(vec!["test_table".to_string()]);
4728 dir_namespace
4729 .create_table(create_request, Bytes::from(buffer))
4730 .await
4731 .unwrap();
4732
4733 let mut request = TableExistsRequest::new();
4735 request.id = Some(vec!["test_table".to_string()]);
4736 let result = dir_namespace.table_exists(request).await;
4737 assert!(result.is_ok());
4738 }
4739
4740 #[rstest]
4741 #[case::with_optimization(true)]
4742 #[case::without_optimization(false)]
4743 #[tokio::test]
4744 async fn test_manifest_namespace_describe_table(#[case] inline_optimization: bool) {
4745 let temp_dir = TempStdDir::default();
4746 let temp_path = temp_dir.to_str().unwrap();
4747
4748 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4749 .inline_optimization_enabled(inline_optimization)
4750 .build()
4751 .await
4752 .unwrap();
4753
4754 let mut request = DescribeTableRequest::new();
4756 request.id = Some(vec!["nonexistent".to_string()]);
4757 let result = dir_namespace.describe_table(request).await;
4758 assert!(result.is_err());
4759
4760 let buffer = create_test_ipc_data();
4762 let mut create_request = CreateTableRequest::new();
4763 create_request.id = Some(vec!["test_table".to_string()]);
4764 dir_namespace
4765 .create_table(create_request, Bytes::from(buffer))
4766 .await
4767 .unwrap();
4768
4769 let mut request = DescribeTableRequest::new();
4771 request.id = Some(vec!["test_table".to_string()]);
4772 let response = dir_namespace.describe_table(request).await.unwrap();
4773 assert!(response.location.is_some());
4774 assert!(response.location.unwrap().contains("test_table"));
4775 }
4776
4777 #[rstest]
4778 #[case::with_optimization(true)]
4779 #[case::without_optimization(false)]
4780 #[tokio::test]
4781 async fn test_manifest_namespace_drop_table(#[case] inline_optimization: bool) {
4782 let temp_dir = TempStdDir::default();
4783 let temp_path = temp_dir.to_str().unwrap();
4784
4785 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4786 .inline_optimization_enabled(inline_optimization)
4787 .build()
4788 .await
4789 .unwrap();
4790
4791 let buffer = create_test_ipc_data();
4793 let mut create_request = CreateTableRequest::new();
4794 create_request.id = Some(vec!["test_table".to_string()]);
4795 dir_namespace
4796 .create_table(create_request, Bytes::from(buffer))
4797 .await
4798 .unwrap();
4799
4800 let mut request = ListTablesRequest::new();
4802 request.id = Some(vec![]);
4803 let response = dir_namespace.list_tables(request).await.unwrap();
4804 assert_eq!(response.tables.len(), 1);
4805
4806 let mut drop_request = DropTableRequest::new();
4808 drop_request.id = Some(vec!["test_table".to_string()]);
4809 let _response = dir_namespace.drop_table(drop_request).await.unwrap();
4810
4811 let mut request = ListTablesRequest::new();
4813 request.id = Some(vec![]);
4814 let response = dir_namespace.list_tables(request).await.unwrap();
4815 assert_eq!(response.tables.len(), 0);
4816 }
4817
4818 #[tokio::test]
4819 async fn test_list_tables_pagination_limit_zero() {
4820 let temp_dir = TempStdDir::default();
4821 let temp_path = temp_dir.to_str().unwrap();
4822
4823 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4824 .build()
4825 .await
4826 .unwrap();
4827
4828 let buffer = create_test_ipc_data();
4829 let mut create_request = CreateTableRequest::new();
4830 create_request.id = Some(vec!["alpha".to_string()]);
4831 dir_namespace
4832 .create_table(create_request, Bytes::from(buffer))
4833 .await
4834 .unwrap();
4835
4836 let response = dir_namespace
4837 .list_tables(ListTablesRequest {
4838 id: Some(vec![]),
4839 limit: Some(0),
4840 ..Default::default()
4841 })
4842 .await
4843 .unwrap();
4844
4845 assert!(response.tables.is_empty());
4846 assert!(response.page_token.is_none());
4847 }
4848
4849 #[rstest]
4850 #[case::with_optimization(true)]
4851 #[case::without_optimization(false)]
4852 #[tokio::test]
4853 async fn test_manifest_namespace_multiple_tables(#[case] inline_optimization: bool) {
4854 let temp_dir = TempStdDir::default();
4855 let temp_path = temp_dir.to_str().unwrap();
4856
4857 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4858 .inline_optimization_enabled(inline_optimization)
4859 .build()
4860 .await
4861 .unwrap();
4862
4863 let buffer = create_test_ipc_data();
4865 for i in 1..=3 {
4866 let mut create_request = CreateTableRequest::new();
4867 create_request.id = Some(vec![format!("table{}", i)]);
4868 dir_namespace
4869 .create_table(create_request, Bytes::from(buffer.clone()))
4870 .await
4871 .unwrap();
4872 }
4873
4874 let mut request = ListTablesRequest::new();
4876 request.id = Some(vec![]);
4877 let response = dir_namespace.list_tables(request).await.unwrap();
4878 assert_eq!(response.tables.len(), 3);
4879 assert!(response.tables.contains(&"table1".to_string()));
4880 assert!(response.tables.contains(&"table2".to_string()));
4881 assert!(response.tables.contains(&"table3".to_string()));
4882 }
4883
4884 #[rstest]
4885 #[case::with_optimization(true)]
4886 #[case::without_optimization(false)]
4887 #[tokio::test]
4888 async fn test_directory_only_mode(#[case] inline_optimization: bool) {
4889 let temp_dir = TempStdDir::default();
4890 let temp_path = temp_dir.to_str().unwrap();
4891
4892 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4894 .manifest_enabled(false)
4895 .inline_optimization_enabled(inline_optimization)
4896 .build()
4897 .await
4898 .unwrap();
4899
4900 let mut request = ListTablesRequest::new();
4902 request.id = Some(vec![]);
4903 let response = dir_namespace.list_tables(request).await.unwrap();
4904 assert_eq!(response.tables.len(), 0);
4905
4906 let buffer = create_test_ipc_data();
4908 let mut create_request = CreateTableRequest::new();
4909 create_request.id = Some(vec!["test_table".to_string()]);
4910
4911 let _response = dir_namespace
4913 .create_table(create_request, Bytes::from(buffer))
4914 .await
4915 .unwrap();
4916
4917 let mut request = ListTablesRequest::new();
4919 request.id = Some(vec![]);
4920 let response = dir_namespace.list_tables(request).await.unwrap();
4921 assert_eq!(response.tables.len(), 1);
4922 assert_eq!(response.tables[0], "test_table");
4923 }
4924
4925 #[rstest]
4926 #[case::with_optimization(true)]
4927 #[case::without_optimization(false)]
4928 #[tokio::test]
4929 async fn test_dual_mode_merge(#[case] inline_optimization: bool) {
4930 let temp_dir = TempStdDir::default();
4931 let temp_path = temp_dir.to_str().unwrap();
4932
4933 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4935 .manifest_enabled(true)
4936 .dir_listing_enabled(true)
4937 .inline_optimization_enabled(inline_optimization)
4938 .build()
4939 .await
4940 .unwrap();
4941
4942 let buffer = create_test_ipc_data();
4944 let mut create_request = CreateTableRequest::new();
4945 create_request.id = Some(vec!["table1".to_string()]);
4946 dir_namespace
4947 .create_table(create_request, Bytes::from(buffer))
4948 .await
4949 .unwrap();
4950
4951 let mut request = ListTablesRequest::new();
4953 request.id = Some(vec![]);
4954 let response = dir_namespace.list_tables(request).await.unwrap();
4955 assert_eq!(response.tables.len(), 1);
4956 assert_eq!(response.tables[0], "table1");
4957 }
4958
4959 #[rstest]
4960 #[case::with_optimization(true)]
4961 #[case::without_optimization(false)]
4962 #[tokio::test]
4963 async fn test_manifest_only_mode(#[case] inline_optimization: bool) {
4964 let temp_dir = TempStdDir::default();
4965 let temp_path = temp_dir.to_str().unwrap();
4966
4967 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4969 .manifest_enabled(true)
4970 .dir_listing_enabled(false)
4971 .inline_optimization_enabled(inline_optimization)
4972 .build()
4973 .await
4974 .unwrap();
4975
4976 let buffer = create_test_ipc_data();
4978 let mut create_request = CreateTableRequest::new();
4979 create_request.id = Some(vec!["test_table".to_string()]);
4980 dir_namespace
4981 .create_table(create_request, Bytes::from(buffer))
4982 .await
4983 .unwrap();
4984
4985 let mut request = ListTablesRequest::new();
4987 request.id = Some(vec![]);
4988 let response = dir_namespace.list_tables(request).await.unwrap();
4989 assert_eq!(response.tables.len(), 1);
4990 assert_eq!(response.tables[0], "test_table");
4991 }
4992
4993 #[rstest]
4994 #[case::with_optimization(true)]
4995 #[case::without_optimization(false)]
4996 #[tokio::test]
4997 async fn test_drop_nonexistent_table(#[case] inline_optimization: bool) {
4998 let temp_dir = TempStdDir::default();
4999 let temp_path = temp_dir.to_str().unwrap();
5000
5001 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5002 .inline_optimization_enabled(inline_optimization)
5003 .build()
5004 .await
5005 .unwrap();
5006
5007 let mut drop_request = DropTableRequest::new();
5009 drop_request.id = Some(vec!["nonexistent".to_string()]);
5010 let result = dir_namespace.drop_table(drop_request).await;
5011 assert!(result.is_err());
5012 }
5013
5014 #[rstest]
5015 #[case::with_optimization(true)]
5016 #[case::without_optimization(false)]
5017 #[tokio::test]
5018 async fn test_create_duplicate_table_fails(#[case] inline_optimization: bool) {
5019 let temp_dir = TempStdDir::default();
5020 let temp_path = temp_dir.to_str().unwrap();
5021
5022 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5023 .inline_optimization_enabled(inline_optimization)
5024 .build()
5025 .await
5026 .unwrap();
5027
5028 let buffer = create_test_ipc_data();
5030 let mut create_request = CreateTableRequest::new();
5031 create_request.id = Some(vec!["test_table".to_string()]);
5032 dir_namespace
5033 .create_table(create_request, Bytes::from(buffer.clone()))
5034 .await
5035 .unwrap();
5036
5037 let mut create_request = CreateTableRequest::new();
5039 create_request.id = Some(vec!["test_table".to_string()]);
5040 let result = dir_namespace
5041 .create_table(create_request, Bytes::from(buffer))
5042 .await;
5043 assert!(result.is_err());
5044 }
5045
5046 #[rstest]
5047 #[case::with_optimization(true)]
5048 #[case::without_optimization(false)]
5049 #[tokio::test]
5050 async fn test_create_child_namespace(#[case] inline_optimization: bool) {
5051 use lance_namespace::models::{
5052 CreateNamespaceRequest, ListNamespacesRequest, NamespaceExistsRequest,
5053 };
5054
5055 let temp_dir = TempStdDir::default();
5056 let temp_path = temp_dir.to_str().unwrap();
5057
5058 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5059 .inline_optimization_enabled(inline_optimization)
5060 .build()
5061 .await
5062 .unwrap();
5063
5064 let mut create_req = CreateNamespaceRequest::new();
5066 create_req.id = Some(vec!["ns1".to_string()]);
5067 let result = dir_namespace.create_namespace(create_req).await;
5068 assert!(
5069 result.is_ok(),
5070 "Failed to create child namespace: {:?}",
5071 result.err()
5072 );
5073
5074 let exists_req = NamespaceExistsRequest {
5076 id: Some(vec!["ns1".to_string()]),
5077 ..Default::default()
5078 };
5079 let result = dir_namespace.namespace_exists(exists_req).await;
5080 assert!(result.is_ok(), "Namespace should exist");
5081
5082 let list_req = ListNamespacesRequest {
5084 id: Some(vec![]),
5085 page_token: None,
5086 limit: None,
5087 ..Default::default()
5088 };
5089 let result = dir_namespace.list_namespaces(list_req).await;
5090 assert!(result.is_ok());
5091 let namespaces = result.unwrap();
5092 assert_eq!(namespaces.namespaces.len(), 1);
5093 assert_eq!(namespaces.namespaces[0], "ns1");
5094 }
5095
5096 #[rstest]
5097 #[case::with_optimization(true)]
5098 #[case::without_optimization(false)]
5099 #[tokio::test]
5100 async fn test_create_nested_namespace(#[case] inline_optimization: bool) {
5101 use lance_namespace::models::{
5102 CreateNamespaceRequest, ListNamespacesRequest, NamespaceExistsRequest,
5103 };
5104
5105 let temp_dir = TempStdDir::default();
5106 let temp_path = temp_dir.to_str().unwrap();
5107
5108 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5109 .inline_optimization_enabled(inline_optimization)
5110 .build()
5111 .await
5112 .unwrap();
5113
5114 let mut create_req = CreateNamespaceRequest::new();
5116 create_req.id = Some(vec!["parent".to_string()]);
5117 dir_namespace.create_namespace(create_req).await.unwrap();
5118
5119 let mut create_req = CreateNamespaceRequest::new();
5121 create_req.id = Some(vec!["parent".to_string(), "child".to_string()]);
5122 let result = dir_namespace.create_namespace(create_req).await;
5123 assert!(
5124 result.is_ok(),
5125 "Failed to create nested namespace: {:?}",
5126 result.err()
5127 );
5128
5129 let exists_req = NamespaceExistsRequest {
5131 id: Some(vec!["parent".to_string(), "child".to_string()]),
5132 ..Default::default()
5133 };
5134 let result = dir_namespace.namespace_exists(exists_req).await;
5135 assert!(result.is_ok(), "Nested namespace should exist");
5136
5137 let list_req = ListNamespacesRequest {
5139 id: Some(vec!["parent".to_string()]),
5140 page_token: None,
5141 limit: None,
5142 ..Default::default()
5143 };
5144 let result = dir_namespace.list_namespaces(list_req).await;
5145 assert!(result.is_ok());
5146 let namespaces = result.unwrap();
5147 assert_eq!(namespaces.namespaces.len(), 1);
5148 assert_eq!(namespaces.namespaces[0], "child");
5149 }
5150
5151 #[rstest]
5152 #[case::with_optimization(true)]
5153 #[case::without_optimization(false)]
5154 #[tokio::test]
5155 async fn test_create_namespace_without_parent_fails(#[case] inline_optimization: bool) {
5156 use lance_namespace::models::CreateNamespaceRequest;
5157
5158 let temp_dir = TempStdDir::default();
5159 let temp_path = temp_dir.to_str().unwrap();
5160
5161 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5162 .inline_optimization_enabled(inline_optimization)
5163 .build()
5164 .await
5165 .unwrap();
5166
5167 let mut create_req = CreateNamespaceRequest::new();
5169 create_req.id = Some(vec!["nonexistent_parent".to_string(), "child".to_string()]);
5170 let result = dir_namespace.create_namespace(create_req).await;
5171 assert!(result.is_err(), "Should fail when parent doesn't exist");
5172 }
5173
5174 #[rstest]
5175 #[case::with_optimization(true)]
5176 #[case::without_optimization(false)]
5177 #[tokio::test]
5178 async fn test_drop_child_namespace(#[case] inline_optimization: bool) {
5179 use lance_namespace::models::{
5180 CreateNamespaceRequest, DropNamespaceRequest, NamespaceExistsRequest,
5181 };
5182
5183 let temp_dir = TempStdDir::default();
5184 let temp_path = temp_dir.to_str().unwrap();
5185
5186 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5187 .inline_optimization_enabled(inline_optimization)
5188 .build()
5189 .await
5190 .unwrap();
5191
5192 let mut create_req = CreateNamespaceRequest::new();
5194 create_req.id = Some(vec!["ns1".to_string()]);
5195 dir_namespace.create_namespace(create_req).await.unwrap();
5196
5197 let mut drop_req = DropNamespaceRequest::new();
5199 drop_req.id = Some(vec!["ns1".to_string()]);
5200 let result = dir_namespace.drop_namespace(drop_req).await;
5201 assert!(
5202 result.is_ok(),
5203 "Failed to drop namespace: {:?}",
5204 result.err()
5205 );
5206
5207 let exists_req = NamespaceExistsRequest {
5209 id: Some(vec!["ns1".to_string()]),
5210 ..Default::default()
5211 };
5212 let result = dir_namespace.namespace_exists(exists_req).await;
5213 assert!(result.is_err(), "Namespace should not exist after drop");
5214 }
5215
5216 #[rstest]
5217 #[case::with_optimization(true)]
5218 #[case::without_optimization(false)]
5219 #[tokio::test]
5220 async fn test_drop_namespace_with_children_fails(#[case] inline_optimization: bool) {
5221 use lance_namespace::models::{CreateNamespaceRequest, DropNamespaceRequest};
5222
5223 let temp_dir = TempStdDir::default();
5224 let temp_path = temp_dir.to_str().unwrap();
5225
5226 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5227 .inline_optimization_enabled(inline_optimization)
5228 .build()
5229 .await
5230 .unwrap();
5231
5232 let mut create_req = CreateNamespaceRequest::new();
5234 create_req.id = Some(vec!["parent".to_string()]);
5235 dir_namespace.create_namespace(create_req).await.unwrap();
5236
5237 let mut create_req = CreateNamespaceRequest::new();
5238 create_req.id = Some(vec!["parent".to_string(), "child".to_string()]);
5239 dir_namespace.create_namespace(create_req).await.unwrap();
5240
5241 let mut drop_req = DropNamespaceRequest::new();
5243 drop_req.id = Some(vec!["parent".to_string()]);
5244 let result = dir_namespace.drop_namespace(drop_req).await;
5245 assert!(result.is_err(), "Should fail when namespace has children");
5246 }
5247
5248 #[rstest]
5249 #[case::with_optimization(true)]
5250 #[case::without_optimization(false)]
5251 #[tokio::test]
5252 async fn test_create_table_in_child_namespace(#[case] inline_optimization: bool) {
5253 use lance_namespace::models::{
5254 CreateNamespaceRequest, CreateTableRequest, ListTablesRequest,
5255 };
5256
5257 let temp_dir = TempStdDir::default();
5258 let temp_path = temp_dir.to_str().unwrap();
5259
5260 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5261 .inline_optimization_enabled(inline_optimization)
5262 .build()
5263 .await
5264 .unwrap();
5265
5266 let mut create_ns_req = CreateNamespaceRequest::new();
5268 create_ns_req.id = Some(vec!["ns1".to_string()]);
5269 dir_namespace.create_namespace(create_ns_req).await.unwrap();
5270
5271 let buffer = create_test_ipc_data();
5273 let mut create_table_req = CreateTableRequest::new();
5274 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
5275 let result = dir_namespace
5276 .create_table(create_table_req, Bytes::from(buffer))
5277 .await;
5278 assert!(
5279 result.is_ok(),
5280 "Failed to create table in child namespace: {:?}",
5281 result.err()
5282 );
5283
5284 let list_req = ListTablesRequest {
5286 id: Some(vec!["ns1".to_string()]),
5287 page_token: None,
5288 limit: None,
5289 ..Default::default()
5290 };
5291 let result = dir_namespace.list_tables(list_req).await;
5292 assert!(result.is_ok());
5293 let tables = result.unwrap();
5294 assert_eq!(tables.tables.len(), 1);
5295 assert_eq!(tables.tables[0], "table1");
5296 }
5297
5298 #[rstest]
5299 #[case::with_optimization(true)]
5300 #[case::without_optimization(false)]
5301 #[tokio::test]
5302 async fn test_describe_child_namespace(#[case] inline_optimization: bool) {
5303 use lance_namespace::models::{CreateNamespaceRequest, DescribeNamespaceRequest};
5304
5305 let temp_dir = TempStdDir::default();
5306 let temp_path = temp_dir.to_str().unwrap();
5307
5308 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5309 .inline_optimization_enabled(inline_optimization)
5310 .build()
5311 .await
5312 .unwrap();
5313
5314 let mut properties = std::collections::HashMap::new();
5316 properties.insert("key1".to_string(), "value1".to_string());
5317
5318 let mut create_req = CreateNamespaceRequest::new();
5319 create_req.id = Some(vec!["ns1".to_string()]);
5320 create_req.properties = Some(properties.clone());
5321 dir_namespace.create_namespace(create_req).await.unwrap();
5322
5323 let describe_req = DescribeNamespaceRequest {
5325 id: Some(vec!["ns1".to_string()]),
5326 ..Default::default()
5327 };
5328 let result = dir_namespace.describe_namespace(describe_req).await;
5329 assert!(
5330 result.is_ok(),
5331 "Failed to describe namespace: {:?}",
5332 result.err()
5333 );
5334 let response = result.unwrap();
5335 assert!(response.properties.is_some());
5336 assert_eq!(
5337 response.properties.unwrap().get("key1"),
5338 Some(&"value1".to_string())
5339 );
5340 }
5341
5342 #[rstest]
5343 #[case::with_optimization(true)]
5344 #[case::without_optimization(false)]
5345 #[tokio::test]
5346 async fn test_concurrent_create_and_drop_single_instance(#[case] inline_optimization: bool) {
5347 use futures::future::join_all;
5348 use std::sync::Arc;
5349
5350 let temp_dir = TempStdDir::default();
5351 let temp_path = temp_dir.to_str().unwrap();
5352
5353 let dir_namespace = Arc::new(
5354 DirectoryNamespaceBuilder::new(temp_path)
5355 .inline_optimization_enabled(inline_optimization)
5356 .build()
5357 .await
5358 .unwrap(),
5359 );
5360
5361 let mut create_ns_request = CreateNamespaceRequest::new();
5364 create_ns_request.id = Some(vec!["test_ns".to_string()]);
5365 dir_namespace
5366 .create_namespace(create_ns_request)
5367 .await
5368 .unwrap();
5369
5370 let num_tables = 10;
5371 let mut handles = Vec::new();
5372
5373 for i in 0..num_tables {
5374 let ns = dir_namespace.clone();
5375 let handle = async move {
5376 let table_name = format!("concurrent_table_{}", i);
5377 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5378 let buffer = create_test_ipc_data();
5379
5380 let mut create_request = CreateTableRequest::new();
5382 create_request.id = Some(table_id.clone());
5383 ns.create_table(create_request, Bytes::from(buffer))
5384 .await
5385 .unwrap_or_else(|e| panic!("Failed to create table {}: {}", table_name, e));
5386
5387 let mut drop_request = DropTableRequest::new();
5389 drop_request.id = Some(table_id);
5390 ns.drop_table(drop_request)
5391 .await
5392 .unwrap_or_else(|e| panic!("Failed to drop table {}: {}", table_name, e));
5393
5394 Ok::<_, lance_core::Error>(())
5395 };
5396 handles.push(handle);
5397 }
5398
5399 let results = join_all(handles).await;
5400 for result in results {
5401 assert!(result.is_ok(), "All concurrent operations should succeed");
5402 }
5403
5404 let mut request = ListTablesRequest::new();
5406 request.id = Some(vec!["test_ns".to_string()]);
5407 let response = dir_namespace.list_tables(request).await.unwrap();
5408 assert_eq!(response.tables.len(), 0, "All tables should be dropped");
5409 }
5410
5411 #[rstest]
5412 #[case::with_optimization(true)]
5413 #[case::without_optimization(false)]
5414 #[tokio::test]
5415 async fn test_concurrent_create_and_drop_multiple_instances(#[case] inline_optimization: bool) {
5416 use futures::future::join_all;
5417
5418 let temp_dir = TempStdDir::default();
5419 let temp_path = temp_dir.to_str().unwrap().to_string();
5420
5421 let init_ns = DirectoryNamespaceBuilder::new(&temp_path)
5424 .inline_optimization_enabled(inline_optimization)
5425 .build()
5426 .await
5427 .unwrap();
5428 let mut create_ns_request = CreateNamespaceRequest::new();
5429 create_ns_request.id = Some(vec!["test_ns".to_string()]);
5430 init_ns.create_namespace(create_ns_request).await.unwrap();
5431
5432 let num_tables = 10;
5433 let mut handles = Vec::new();
5434
5435 for i in 0..num_tables {
5436 let path = temp_path.clone();
5437 let handle = async move {
5438 let ns = DirectoryNamespaceBuilder::new(&path)
5440 .inline_optimization_enabled(inline_optimization)
5441 .build()
5442 .await
5443 .unwrap();
5444
5445 let table_name = format!("multi_ns_table_{}", i);
5446 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5447 let buffer = create_test_ipc_data();
5448
5449 let mut create_request = CreateTableRequest::new();
5451 create_request.id = Some(table_id.clone());
5452 ns.create_table(create_request, Bytes::from(buffer))
5453 .await
5454 .unwrap_or_else(|e| panic!("Failed to create table {}: {}", table_name, e));
5455
5456 let mut drop_request = DropTableRequest::new();
5458 drop_request.id = Some(table_id);
5459 ns.drop_table(drop_request)
5460 .await
5461 .unwrap_or_else(|e| panic!("Failed to drop table {}: {}", table_name, e));
5462
5463 Ok::<_, lance_core::Error>(())
5464 };
5465 handles.push(handle);
5466 }
5467
5468 let results = join_all(handles).await;
5469 for result in results {
5470 assert!(result.is_ok(), "All concurrent operations should succeed");
5471 }
5472
5473 let verify_ns = DirectoryNamespaceBuilder::new(&temp_path)
5475 .inline_optimization_enabled(inline_optimization)
5476 .build()
5477 .await
5478 .unwrap();
5479
5480 let mut request = ListTablesRequest::new();
5481 request.id = Some(vec!["test_ns".to_string()]);
5482 let response = verify_ns.list_tables(request).await.unwrap();
5483 assert_eq!(response.tables.len(), 0, "All tables should be dropped");
5484 }
5485
5486 #[rstest]
5487 #[case::with_optimization(true)]
5488 #[case::without_optimization(false)]
5489 #[tokio::test]
5490 async fn test_concurrent_create_then_drop_from_different_instance(
5491 #[case] inline_optimization: bool,
5492 ) {
5493 use futures::future::join_all;
5494
5495 let temp_dir = TempStdDir::default();
5496 let temp_path = temp_dir.to_str().unwrap().to_string();
5497
5498 let init_ns = DirectoryNamespaceBuilder::new(&temp_path)
5501 .inline_optimization_enabled(inline_optimization)
5502 .build()
5503 .await
5504 .unwrap();
5505 let mut create_ns_request = CreateNamespaceRequest::new();
5506 create_ns_request.id = Some(vec!["test_ns".to_string()]);
5507 init_ns.create_namespace(create_ns_request).await.unwrap();
5508
5509 let num_tables = 10;
5510
5511 let mut create_handles = Vec::new();
5513 for i in 0..num_tables {
5514 let path = temp_path.clone();
5515 let handle = async move {
5516 let ns = DirectoryNamespaceBuilder::new(&path)
5517 .inline_optimization_enabled(inline_optimization)
5518 .build()
5519 .await
5520 .unwrap();
5521
5522 let table_name = format!("cross_instance_table_{}", i);
5523 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5524 let buffer = create_test_ipc_data();
5525
5526 let mut create_request = CreateTableRequest::new();
5527 create_request.id = Some(table_id);
5528 ns.create_table(create_request, Bytes::from(buffer))
5529 .await
5530 .unwrap_or_else(|e| panic!("Failed to create table {}: {}", table_name, e));
5531
5532 Ok::<_, lance_core::Error>(())
5533 };
5534 create_handles.push(handle);
5535 }
5536
5537 let create_results = join_all(create_handles).await;
5538 for result in create_results {
5539 assert!(result.is_ok(), "All create operations should succeed");
5540 }
5541
5542 let mut drop_handles = Vec::new();
5544 for i in 0..num_tables {
5545 let path = temp_path.clone();
5546 let handle = async move {
5547 let ns = DirectoryNamespaceBuilder::new(&path)
5548 .inline_optimization_enabled(inline_optimization)
5549 .build()
5550 .await
5551 .unwrap();
5552
5553 let table_name = format!("cross_instance_table_{}", i);
5554 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5555
5556 let mut drop_request = DropTableRequest::new();
5557 drop_request.id = Some(table_id);
5558 ns.drop_table(drop_request)
5559 .await
5560 .unwrap_or_else(|e| panic!("Failed to drop table {}: {}", table_name, e));
5561
5562 Ok::<_, lance_core::Error>(())
5563 };
5564 drop_handles.push(handle);
5565 }
5566
5567 let drop_results = join_all(drop_handles).await;
5568 for result in drop_results {
5569 assert!(result.is_ok(), "All drop operations should succeed");
5570 }
5571
5572 let verify_ns = DirectoryNamespaceBuilder::new(&temp_path)
5574 .inline_optimization_enabled(inline_optimization)
5575 .build()
5576 .await
5577 .unwrap();
5578
5579 let mut request = ListTablesRequest::new();
5580 request.id = Some(vec!["test_ns".to_string()]);
5581 let response = verify_ns.list_tables(request).await.unwrap();
5582 assert_eq!(response.tables.len(), 0, "All tables should be dropped");
5583 }
5584
5585 #[test]
5586 fn test_construct_full_uri_with_cloud_urls() {
5587 let s3_result =
5589 ManifestNamespace::construct_full_uri("s3://bucket/path/subdir", "table.lance")
5590 .unwrap();
5591 assert_eq!(
5592 s3_result, "s3://bucket/path/subdir/table.lance",
5593 "S3 URL should correctly append table name to nested path"
5594 );
5595
5596 let az_result =
5598 ManifestNamespace::construct_full_uri("az://container/path/subdir", "table.lance")
5599 .unwrap();
5600 assert_eq!(
5601 az_result, "az://container/path/subdir/table.lance",
5602 "Azure URL should correctly append table name to nested path"
5603 );
5604
5605 let gs_result =
5607 ManifestNamespace::construct_full_uri("gs://bucket/path/subdir", "table.lance")
5608 .unwrap();
5609 assert_eq!(
5610 gs_result, "gs://bucket/path/subdir/table.lance",
5611 "GCS URL should correctly append table name to nested path"
5612 );
5613
5614 let deep_result =
5616 ManifestNamespace::construct_full_uri("s3://bucket/a/b/c/d", "my_table.lance").unwrap();
5617 assert_eq!(
5618 deep_result, "s3://bucket/a/b/c/d/my_table.lance",
5619 "Deeply nested path should work correctly"
5620 );
5621
5622 let shallow_result =
5624 ManifestNamespace::construct_full_uri("s3://bucket", "table.lance").unwrap();
5625 assert_eq!(
5626 shallow_result, "s3://bucket/table.lance",
5627 "Single-level nested path should work correctly"
5628 );
5629
5630 let trailing_slash_result =
5632 ManifestNamespace::construct_full_uri("s3://bucket/path/subdir/", "table.lance")
5633 .unwrap();
5634 assert_eq!(
5635 trailing_slash_result, "s3://bucket/path/subdir/table.lance",
5636 "URL with existing trailing slash should still work"
5637 );
5638
5639 let empty_query_result =
5642 ManifestNamespace::construct_full_uri("s3://bucket/path?", "table.lance").unwrap();
5643 assert_eq!(
5644 empty_query_result, "s3://bucket/path/table.lance",
5645 "URL with empty query string should not include trailing '?'"
5646 );
5647
5648 let query_param_result =
5651 ManifestNamespace::construct_full_uri("s3://bucket/path?param=value", "table.lance")
5652 .unwrap();
5653 assert_eq!(
5654 query_param_result, "s3://bucket/path/table.lance",
5655 "URL with query parameters should have them stripped"
5656 );
5657 }
5658
5659 #[test]
5660 fn test_construct_full_uri_with_dollar_sign() {
5661 let result =
5662 ManifestNamespace::construct_full_uri("/tmp/root", "hash_workspace$test_table")
5663 .unwrap();
5664
5665 assert!(
5666 result.ends_with("/tmp/root/hash_workspace$test_table"),
5667 "local file URI should preserve dollar signs without adding empty path segments: {}",
5668 result
5669 );
5670 assert!(
5671 !result.contains("//hash_workspace$test_table"),
5672 "local file URI should not add a double slash before table directory: {}",
5673 result
5674 );
5675 }
5676
5677 #[test]
5678 fn test_construct_full_uri_with_nested_relative_location() {
5679 let result =
5680 ManifestNamespace::construct_full_uri("/tmp/root", "workspace/physical_table.lance")
5681 .unwrap();
5682
5683 assert!(
5684 result.ends_with("/tmp/root/workspace/physical_table.lance"),
5685 "nested relative location should preserve path separators: {}",
5686 result
5687 );
5688 assert!(
5689 !result.contains("%2Fphysical_table.lance"),
5690 "nested relative location should not encode path separators: {}",
5691 result
5692 );
5693 }
5694
5695 #[tokio::test]
5702 async fn test_concurrent_create_table_no_duplicates() {
5703 let temp_dir = TempStdDir::default();
5704 let temp_path = temp_dir.to_str().unwrap();
5705
5706 let ns1 = DirectoryNamespaceBuilder::new(temp_path)
5709 .inline_optimization_enabled(false)
5710 .build()
5711 .await
5712 .unwrap();
5713 let ns2 = DirectoryNamespaceBuilder::new(temp_path)
5714 .inline_optimization_enabled(false)
5715 .build()
5716 .await
5717 .unwrap();
5718
5719 let buffer = create_test_ipc_data();
5720
5721 let mut req1 = CreateTableRequest::new();
5722 req1.id = Some(vec!["race_table".to_string()]);
5723 let mut req2 = CreateTableRequest::new();
5724 req2.id = Some(vec!["race_table".to_string()]);
5725
5726 let (result1, result2) = tokio::join!(
5728 ns1.create_table(req1, Bytes::from(buffer.clone())),
5729 ns2.create_table(req2, Bytes::from(buffer.clone())),
5730 );
5731
5732 let success_count = [&result1, &result2].iter().filter(|r| r.is_ok()).count();
5734 let failure_count = [&result1, &result2].iter().filter(|r| r.is_err()).count();
5735 assert_eq!(
5736 success_count, 1,
5737 "Exactly one create should succeed, got: result1={:?}, result2={:?}",
5738 result1, result2
5739 );
5740 assert_eq!(
5741 failure_count, 1,
5742 "Exactly one create should fail, got: result1={:?}, result2={:?}",
5743 result1, result2
5744 );
5745
5746 let ns_check = DirectoryNamespaceBuilder::new(temp_path)
5748 .inline_optimization_enabled(false)
5749 .build()
5750 .await
5751 .unwrap();
5752 let mut list_request = ListTablesRequest::new();
5753 list_request.id = Some(vec![]);
5754 let response = ns_check.list_tables(list_request).await.unwrap();
5755 assert_eq!(
5756 response.tables.len(),
5757 1,
5758 "Should have exactly 1 table, found: {:?}",
5759 response.tables
5760 );
5761 assert_eq!(response.tables[0], "race_table");
5762
5763 let mut describe_request = DescribeTableRequest::new();
5765 describe_request.id = Some(vec!["race_table".to_string()]);
5766 let describe_result = ns_check.describe_table(describe_request).await;
5767 assert!(
5768 describe_result.is_ok(),
5769 "describe_table should not fail with duplicate entries: {:?}",
5770 describe_result
5771 );
5772 }
5773
5774 fn names(v: &[&str]) -> Vec<String> {
5777 v.iter().map(|s| s.to_string()).collect()
5778 }
5779
5780 #[test]
5781 fn test_apply_pagination_no_token_no_limit() {
5782 let mut n = names(&["b", "a", "c"]);
5783 let next = ManifestNamespace::apply_pagination(&mut n, None, None);
5784 assert_eq!(n, names(&["a", "b", "c"]));
5785 assert_eq!(next, None);
5786 }
5787
5788 #[test]
5789 fn test_apply_pagination_limit_truncates_and_returns_token() {
5790 let mut n = names(&["c", "a", "b"]);
5791 let next = ManifestNamespace::apply_pagination(&mut n, None, Some(2));
5792 assert_eq!(n, names(&["a", "b"]));
5793 assert_eq!(next, Some("b".to_string()));
5794 }
5795
5796 #[test]
5797 fn test_apply_pagination_limit_zero_returns_empty_no_token() {
5798 let mut n = names(&["a", "b", "c"]);
5799 let next = ManifestNamespace::apply_pagination(&mut n, None, Some(0));
5800 assert!(n.is_empty());
5801 assert_eq!(next, None);
5802 }
5803
5804 #[test]
5805 fn test_apply_pagination_page_token_in_list() {
5806 let mut n = names(&["a", "b", "c", "d"]);
5808 let next = ManifestNamespace::apply_pagination(&mut n, Some("b".to_string()), None);
5809 assert_eq!(n, names(&["c", "d"]));
5810 assert_eq!(next, None);
5811 }
5812
5813 #[test]
5814 fn test_apply_pagination_page_token_past_all_items() {
5815 let mut n = names(&["a", "b", "c"]);
5816 let next = ManifestNamespace::apply_pagination(&mut n, Some("z".to_string()), None);
5817 assert!(n.is_empty());
5818 assert_eq!(next, None);
5819 }
5820
5821 #[test]
5822 fn test_apply_pagination_token_and_limit_combined() {
5823 let mut n = names(&["a", "b", "c", "d", "e"]);
5824 let next = ManifestNamespace::apply_pagination(&mut n, Some("b".to_string()), Some(2));
5825 assert_eq!(n, names(&["c", "d"]));
5826 assert_eq!(next, Some("d".to_string()));
5827 }
5828
5829 #[rstest]
5830 #[case::with_optimization(true)]
5831 #[case::without_optimization(false)]
5832 #[tokio::test]
5833 async fn test_alter_table_add_columns(#[case] inline_optimization: bool) {
5834 use lance_namespace::models::{
5835 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
5836 };
5837
5838 let temp_dir = TempStdDir::default();
5839 let temp_path = temp_dir.to_str().unwrap();
5840
5841 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5842 .inline_optimization_enabled(inline_optimization)
5843 .build()
5844 .await
5845 .unwrap();
5846
5847 let buffer = create_test_ipc_data();
5849 let mut create_request = CreateTableRequest::new();
5850 create_request.id = Some(vec!["test_table".to_string()]);
5851 dir_namespace
5852 .create_table(create_request, Bytes::from(buffer))
5853 .await
5854 .unwrap();
5855
5856 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
5858 new_col.expression = Some(Some("id * 2".to_string()));
5859 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
5860 add_request.id = Some(vec!["test_table".to_string()]);
5861
5862 let response = dir_namespace
5863 .alter_table_add_columns(add_request)
5864 .await
5865 .unwrap();
5866 assert!(response.version > 1);
5868
5869 let mut describe_request = DescribeTableRequest::new();
5871 describe_request.id = Some(vec!["test_table".to_string()]);
5872 describe_request.load_detailed_metadata = Some(true);
5873 let describe_response = dir_namespace
5874 .describe_table(describe_request)
5875 .await
5876 .unwrap();
5877 assert!(describe_response.schema.is_some());
5878
5879 let schema = describe_response.schema.unwrap();
5880 let field_names: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect();
5881 assert!(
5882 field_names.contains(&"doubled_id"),
5883 "Column 'doubled_id' should exist after add_columns, got: {:?}",
5884 field_names
5885 );
5886 }
5887
5888 #[rstest]
5889 #[case::with_optimization(true)]
5890 #[case::without_optimization(false)]
5891 #[tokio::test]
5892 async fn test_alter_table_add_columns_missing_id(#[case] inline_optimization: bool) {
5893 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
5894
5895 let temp_dir = TempStdDir::default();
5896 let temp_path = temp_dir.to_str().unwrap();
5897
5898 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5899 .inline_optimization_enabled(inline_optimization)
5900 .build()
5901 .await
5902 .unwrap();
5903
5904 let new_col = AddColumnsEntry::new("col".to_string());
5906 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
5907 let result = dir_namespace.alter_table_add_columns(request).await;
5908 assert!(result.is_err(), "Should fail when table ID is missing");
5909 }
5910
5911 #[rstest]
5912 #[case::with_optimization(true)]
5913 #[case::without_optimization(false)]
5914 #[tokio::test]
5915 async fn test_alter_table_add_columns_nonexistent_table(#[case] inline_optimization: bool) {
5916 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
5917
5918 let temp_dir = TempStdDir::default();
5919 let temp_path = temp_dir.to_str().unwrap();
5920
5921 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5922 .inline_optimization_enabled(inline_optimization)
5923 .build()
5924 .await
5925 .unwrap();
5926
5927 let new_col = AddColumnsEntry::new("col".to_string());
5929 let mut request = AlterTableAddColumnsRequest::new(vec![new_col]);
5930 request.id = Some(vec!["nonexistent".to_string()]);
5931 let result = dir_namespace.alter_table_add_columns(request).await;
5932 assert!(result.is_err(), "Should fail when table does not exist");
5933 }
5934
5935 #[rstest]
5936 #[case::with_optimization(true)]
5937 #[case::without_optimization(false)]
5938 #[tokio::test]
5939 async fn test_alter_table_alter_columns_rename(#[case] inline_optimization: bool) {
5940 use lance_namespace::models::{
5941 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
5942 };
5943
5944 let temp_dir = TempStdDir::default();
5945 let temp_path = temp_dir.to_str().unwrap();
5946
5947 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5948 .inline_optimization_enabled(inline_optimization)
5949 .build()
5950 .await
5951 .unwrap();
5952
5953 let buffer = create_test_ipc_data();
5955 let mut create_request = CreateTableRequest::new();
5956 create_request.id = Some(vec!["test_table".to_string()]);
5957 dir_namespace
5958 .create_table(create_request, Bytes::from(buffer))
5959 .await
5960 .unwrap();
5961
5962 let mut entry = AlterColumnsEntry::new("name".to_string());
5964 entry.rename = Some(Some("full_name".to_string()));
5965 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
5966 alter_request.id = Some(vec!["test_table".to_string()]);
5967
5968 let response = dir_namespace
5969 .alter_table_alter_columns(alter_request)
5970 .await
5971 .unwrap();
5972 assert!(response.version > 1);
5973
5974 let mut describe_request = DescribeTableRequest::new();
5976 describe_request.id = Some(vec!["test_table".to_string()]);
5977 describe_request.load_detailed_metadata = Some(true);
5978 let describe_response = dir_namespace
5979 .describe_table(describe_request)
5980 .await
5981 .unwrap();
5982 assert!(describe_response.schema.is_some());
5983
5984 let schema = describe_response.schema.unwrap();
5985 let field_names: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect();
5986 assert!(
5987 field_names.contains(&"full_name"),
5988 "Column should be renamed to 'full_name', got: {:?}",
5989 field_names
5990 );
5991 assert!(
5992 !field_names.contains(&"name"),
5993 "Old column name 'name' should no longer exist, got: {:?}",
5994 field_names
5995 );
5996 }
5997
5998 #[rstest]
5999 #[case::with_optimization(true)]
6000 #[case::without_optimization(false)]
6001 #[tokio::test]
6002 async fn test_alter_table_alter_columns_missing_id(#[case] inline_optimization: bool) {
6003 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
6004
6005 let temp_dir = TempStdDir::default();
6006 let temp_path = temp_dir.to_str().unwrap();
6007
6008 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6009 .inline_optimization_enabled(inline_optimization)
6010 .build()
6011 .await
6012 .unwrap();
6013
6014 let entry = AlterColumnsEntry::new("name".to_string());
6015 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
6016 let result = dir_namespace.alter_table_alter_columns(request).await;
6017 assert!(result.is_err(), "Should fail when table ID is missing");
6018 }
6019
6020 #[rstest]
6021 #[case::with_optimization(true)]
6022 #[case::without_optimization(false)]
6023 #[tokio::test]
6024 async fn test_alter_table_drop_columns(#[case] inline_optimization: bool) {
6025 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
6026
6027 let temp_dir = TempStdDir::default();
6028 let temp_path = temp_dir.to_str().unwrap();
6029
6030 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6031 .inline_optimization_enabled(inline_optimization)
6032 .build()
6033 .await
6034 .unwrap();
6035
6036 let buffer = create_test_ipc_data();
6038 let mut create_request = CreateTableRequest::new();
6039 create_request.id = Some(vec!["test_table".to_string()]);
6040 dir_namespace
6041 .create_table(create_request, Bytes::from(buffer))
6042 .await
6043 .unwrap();
6044
6045 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
6047 drop_request.id = Some(vec!["test_table".to_string()]);
6048
6049 let response = dir_namespace
6050 .alter_table_drop_columns(drop_request)
6051 .await
6052 .unwrap();
6053 assert!(response.version > 1);
6054
6055 let mut describe_request = DescribeTableRequest::new();
6057 describe_request.id = Some(vec!["test_table".to_string()]);
6058 describe_request.load_detailed_metadata = Some(true);
6059 let describe_response = dir_namespace
6060 .describe_table(describe_request)
6061 .await
6062 .unwrap();
6063 assert!(describe_response.schema.is_some());
6064
6065 let schema = describe_response.schema.unwrap();
6066 let field_names: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect();
6067 assert!(
6068 !field_names.contains(&"name"),
6069 "Column 'name' should have been dropped, got: {:?}",
6070 field_names
6071 );
6072 assert!(
6073 field_names.contains(&"id"),
6074 "Column 'id' should still exist, got: {:?}",
6075 field_names
6076 );
6077 }
6078
6079 #[rstest]
6080 #[case::with_optimization(true)]
6081 #[case::without_optimization(false)]
6082 #[tokio::test]
6083 async fn test_alter_table_drop_columns_missing_id(#[case] inline_optimization: bool) {
6084 use lance_namespace::models::AlterTableDropColumnsRequest;
6085
6086 let temp_dir = TempStdDir::default();
6087 let temp_path = temp_dir.to_str().unwrap();
6088
6089 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6090 .inline_optimization_enabled(inline_optimization)
6091 .build()
6092 .await
6093 .unwrap();
6094
6095 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
6096 let result = dir_namespace.alter_table_drop_columns(request).await;
6097 assert!(result.is_err(), "Should fail when table ID is missing");
6098 }
6099
6100 #[rstest]
6101 #[case::with_optimization(true)]
6102 #[case::without_optimization(false)]
6103 #[tokio::test]
6104 async fn test_alter_table_drop_columns_nonexistent_table(#[case] inline_optimization: bool) {
6105 use lance_namespace::models::AlterTableDropColumnsRequest;
6106
6107 let temp_dir = TempStdDir::default();
6108 let temp_path = temp_dir.to_str().unwrap();
6109
6110 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6111 .inline_optimization_enabled(inline_optimization)
6112 .build()
6113 .await
6114 .unwrap();
6115
6116 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
6117 request.id = Some(vec!["nonexistent".to_string()]);
6118 let result = dir_namespace.alter_table_drop_columns(request).await;
6119 assert!(result.is_err(), "Should fail when table does not exist");
6120 }
6121}