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, ensure_can_write_manifest};
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 ensure_can_write_manifest(manifest).map_err(CommitError::from)?;
1844 apply_feature_flags(manifest, false, false).map_err(CommitError::from)?;
1845 let timestamp_nanos = SystemTime::now()
1846 .duration_since(UNIX_EPOCH)
1847 .map(|d| d.as_nanos())
1848 .unwrap_or(0);
1849 manifest.set_timestamp(timestamp_nanos);
1850 manifest.update_max_fragment_id();
1851
1852 let object_store = dataset
1856 .object_store(None)
1857 .await
1858 .map_err(CommitError::from)?;
1859 let base_path = self.base_path.clone().join(MANIFEST_TABLE_NAME);
1860 let naming_scheme = dataset.manifest_location().naming_scheme;
1861 commit_handler
1862 .commit(
1863 manifest,
1864 indices,
1865 &base_path,
1866 &object_store,
1867 write_manifest_file_to_path,
1868 naming_scheme,
1869 Some((&transaction).into()),
1870 )
1871 .await
1872 .map(|_location| ())
1873 }
1874
1875 async fn manifest_commit_landed(
1880 &self,
1881 dataset: &Dataset,
1882 target_version: u64,
1883 data_files: &HashSet<String>,
1884 ) -> bool {
1885 let Ok(committed) = dataset.checkout_version(target_version).await else {
1886 return false;
1887 };
1888 committed.manifest().fragments.iter().any(|fragment| {
1889 fragment
1890 .files
1891 .iter()
1892 .any(|file| data_files.contains(file.path.as_str()))
1893 })
1894 }
1895
1896 async fn resolve_manifest_conflict<O: Clone>(
1900 &self,
1901 resolution: &ConflictResolution<O>,
1902 ) -> Result<Option<O>> {
1903 match resolution {
1904 ConflictResolution::Retry => Ok(None),
1905 ConflictResolution::FailIfExists(object_ids) => {
1906 for object_id in object_ids {
1907 if self.manifest_contains_object(object_id).await? {
1908 return Err(NamespaceError::ConcurrentModification {
1909 message: format!(
1910 "Object '{}' was concurrently created by another operation",
1911 object_id
1912 ),
1913 }
1914 .into());
1915 }
1916 }
1917 Ok(None)
1918 }
1919 ConflictResolution::SucceedIfAbsent { object_id, output } => {
1920 if self.manifest_contains_object(object_id).await? {
1921 Ok(None)
1922 } else {
1923 Ok(Some(output.clone()))
1924 }
1925 }
1926 }
1927 }
1928
1929 async fn ensure_manifest_writable(&self) -> Result<()> {
1935 let dataset_guard = self.manifest_dataset.get().await?;
1936 ensure_can_write_manifest(dataset_guard.manifest())?;
1937 ensure_writable(dataset_guard.metadata())
1938 }
1939
1940 async fn rewrite_manifest<M, F>(
1941 &self,
1942 operation: &str,
1943 mut make_mutation: F,
1944 ) -> Result<M::Output>
1945 where
1946 M: ManifestStreamMutation + 'static,
1947 F: FnMut() -> M,
1948 {
1949 let _mutation_guard = self.manifest_mutation_lock.lock().await;
1950 let max_retries = self.manifest_rewrite_commit_retries();
1951 let mut retries = 0;
1952 let build_indices = self.inline_optimization_enabled;
1953 let commit_handler = self.manifest_commit_handler().await?;
1954
1955 loop {
1956 let dataset_guard = self.manifest_dataset.get_refreshed().await?;
1957 ensure_can_write_manifest(dataset_guard.manifest())?;
1958 let dataset = Arc::new(dataset_guard.clone());
1959 drop(dataset_guard);
1960 ensure_writable(dataset.metadata())?;
1963 let object_store = dataset.object_store(None).await?;
1966
1967 let source = Self::manifest_projected_stream(&dataset).await?;
1968 let resolution = make_mutation().conflict_resolution();
1969 let shared = Arc::new(StdMutex::new(ManifestRewriteShared::new(make_mutation())));
1970 let output_stream = Self::manifest_rewrite_output_stream(source, shared.clone());
1971 let write_params = WriteParams {
1976 mode: WriteMode::Overwrite,
1977 session: self.session.clone(),
1978 max_rows_per_file: u32::MAX as usize,
1979 max_bytes_per_file: usize::MAX,
1980 skip_auto_cleanup: true,
1981 ..WriteParams::default()
1982 };
1983
1984 let transaction = match InsertBuilder::new(dataset.clone())
1985 .with_params(&write_params)
1986 .execute_uncommitted_stream(output_stream)
1987 .await
1988 {
1989 Ok(transaction) => transaction,
1990 Err(err) => {
1991 if let Some(stream_err) = Self::take_manifest_rewrite_error(&shared)? {
1992 return Err(stream_err);
1993 }
1994 return Err(convert_lance_commit_error(&err, operation, None));
1995 }
1996 };
1997
1998 let (mutation, index_data) = Self::take_manifest_rewrite_result(&shared)?;
1999
2000 let Operation::Overwrite {
2001 fragments, schema, ..
2002 } = &transaction.operation
2003 else {
2004 return Err(NamespaceError::Internal {
2005 message: "Manifest rewrite transaction is not an overwrite".to_string(),
2006 }
2007 .into());
2008 };
2009 let staged_data_files = fragments
2012 .iter()
2013 .flat_map(|fragment| fragment.files.iter())
2014 .filter(|file| file.base_id.is_none())
2015 .map(|file| file.path.clone())
2016 .collect::<HashSet<_>>();
2017
2018 if !mutation.has_changes {
2019 self.cleanup_staged_manifest_files(&object_store, &staged_data_files, &[])
2020 .await;
2021 return Ok(mutation.result);
2022 }
2023
2024 let mut manifest = Self::manifest_from_overwrite_transaction(
2025 dataset.manifest(),
2026 schema.clone(),
2027 fragments,
2028 );
2029 let target_version = manifest.version;
2030
2031 let index_uuids = [Uuid::new_v4(), Uuid::new_v4(), Uuid::new_v4()];
2032 let indices = if build_indices {
2033 match Self::build_manifest_indices(&dataset, &manifest, index_data, index_uuids)
2034 .await
2035 {
2036 Ok(indices) => Some(indices),
2037 Err(err) => {
2038 self.cleanup_staged_manifest_files(
2039 &object_store,
2040 &staged_data_files,
2041 &index_uuids,
2042 )
2043 .await;
2044 return Err(err);
2045 }
2046 }
2047 } else {
2048 None
2049 };
2050 let staged_index_uuids: &[Uuid] = if build_indices { &index_uuids } else { &[] };
2051
2052 let commit_result = self
2053 .commit_manifest_overwrite(
2054 &dataset,
2055 commit_handler.as_ref(),
2056 &mut manifest,
2057 indices,
2058 transaction,
2059 )
2060 .await;
2061
2062 match commit_result {
2063 Ok(()) => {
2064 let _ = self.manifest_dataset.get_refreshed().await;
2065 return Ok(mutation.result);
2066 }
2067 Err(err) => {
2068 if self
2072 .manifest_commit_landed(&dataset, target_version, &staged_data_files)
2073 .await
2074 {
2075 let _ = self.manifest_dataset.get_refreshed().await;
2076 return Ok(mutation.result);
2077 }
2078 self.cleanup_staged_manifest_files(
2079 &object_store,
2080 &staged_data_files,
2081 staged_index_uuids,
2082 )
2083 .await;
2084 match err {
2085 CommitError::CommitConflict => {
2086 if let Some(output) =
2087 self.resolve_manifest_conflict(&resolution).await?
2088 {
2089 return Ok(output);
2090 }
2091 if retries >= max_retries {
2092 return Err(NamespaceError::ConcurrentModification {
2093 message: format!(
2094 "{}: still conflicting after {} retries",
2095 operation, max_retries
2096 ),
2097 }
2098 .into());
2099 }
2100 retries += 1;
2101 tokio::time::sleep(std::time::Duration::from_millis(
2102 10 * u64::from(retries),
2103 ))
2104 .await;
2105 }
2106 CommitError::OtherError(err) => {
2107 return Err(convert_lance_commit_error(&err, operation, None));
2108 }
2109 }
2110 }
2111 }
2112 }
2113 }
2114
2115 async fn manifest_contains_object(&self, object_id: &str) -> Result<bool> {
2117 let escaped_id = object_id.replace('\'', "''");
2118 let filter = format!("object_id = '{}'", escaped_id);
2119
2120 let dataset_guard = self.manifest_dataset.get().await?;
2121 let mut scanner = dataset_guard.scan();
2122
2123 scanner.filter(&filter).map_err(|e| {
2124 lance_core::Error::from(NamespaceError::Internal {
2125 message: format!("Failed to filter: {:?}", e),
2126 })
2127 })?;
2128
2129 scanner.project::<&str>(&[]).map_err(|e| {
2131 lance_core::Error::from(NamespaceError::Internal {
2132 message: format!("Failed to project: {:?}", e),
2133 })
2134 })?;
2135
2136 scanner.with_row_id();
2137
2138 let count = scanner.count_rows().await.map_err(|e| {
2139 lance_core::Error::from(NamespaceError::Internal {
2140 message: format!("Failed to count rows: {:?}", e),
2141 })
2142 })?;
2143
2144 Ok(count > 0)
2145 }
2146
2147 async fn query_manifest_for_table(&self, object_id: &str) -> Result<Option<TableInfo>> {
2149 let escaped_id = object_id.replace('\'', "''");
2150 let filter = format!("object_id = '{}' AND object_type = 'table'", escaped_id);
2151 let mut scanner = self.manifest_scanner().await?;
2152 scanner.filter(&filter).map_err(|e| {
2153 lance_core::Error::from(NamespaceError::Internal {
2154 message: format!("Failed to filter: {:?}", e),
2155 })
2156 })?;
2157 scanner
2158 .project(&["object_id", "location", "metadata"])
2159 .map_err(|e| {
2160 lance_core::Error::from(NamespaceError::Internal {
2161 message: format!("Failed to project: {:?}", e),
2162 })
2163 })?;
2164 let batches = Self::execute_scanner(scanner).await?;
2165
2166 let mut found_result: Option<TableInfo> = None;
2167 let mut total_rows = 0;
2168
2169 for batch in batches {
2170 if batch.num_rows() == 0 {
2171 continue;
2172 }
2173
2174 total_rows += batch.num_rows();
2175 if total_rows > 1 {
2176 return Err(NamespaceError::Internal {
2177 message: format!(
2178 "Expected exactly 1 table with id '{}', found {}",
2179 object_id, total_rows
2180 ),
2181 }
2182 .into());
2183 }
2184
2185 let object_id_array = Self::get_string_column(&batch, "object_id")?;
2186 let location_array = Self::get_string_column(&batch, "location")?;
2187 let metadata_array = Self::get_string_column(&batch, "metadata")?;
2188 let location = location_array.value(0).to_string();
2189 let metadata = if !metadata_array.is_null(0) {
2190 let metadata_str = metadata_array.value(0);
2191 match serde_json::from_str::<HashMap<String, String>>(metadata_str) {
2192 Ok(map) => Some(map),
2193 Err(e) => {
2194 return Err(NamespaceError::Internal {
2195 message: format!(
2196 "Failed to deserialize metadata for table '{}': {}",
2197 object_id, e
2198 ),
2199 }
2200 .into());
2201 }
2202 }
2203 } else {
2204 None
2205 };
2206 let (namespace, name) = Self::parse_object_id(object_id_array.value(0));
2207 found_result = Some(TableInfo {
2208 namespace,
2209 name,
2210 location,
2211 metadata,
2212 });
2213 }
2214
2215 Ok(found_result)
2216 }
2217
2218 fn serialize_metadata(
2219 properties: Option<&HashMap<String, String>>,
2220 object_type: &str,
2221 object_id: &str,
2222 ) -> Result<Option<String>> {
2223 match properties {
2224 Some(properties) if !properties.is_empty() => {
2225 serde_json::to_string(properties).map(Some).map_err(|e| {
2226 LanceError::from(NamespaceError::Internal {
2227 message: format!(
2228 "Failed to serialize {} metadata for '{}': {}",
2229 object_type, object_id, e
2230 ),
2231 })
2232 })
2233 }
2234 _ => Ok(None),
2235 }
2236 }
2237
2238 pub(crate) async fn path_has_actual_manifests(
2239 object_store: &ObjectStore,
2240 table_path: &Path,
2241 ) -> Result<bool> {
2242 let versions_path = table_path
2243 .clone()
2244 .join(lance_table::io::commit::VERSIONS_DIR);
2245 Ok(object_store
2248 .list(Some(versions_path))
2249 .try_next()
2250 .await?
2251 .is_some())
2252 }
2253
2254 async fn location_has_actual_manifests(&self, location: &str) -> Result<bool> {
2255 Self::path_has_actual_manifests(&self.object_store, &self.base_path.clone().join(location))
2256 .await
2257 }
2258
2259 pub(crate) fn is_not_found_load_error(err: &LanceError) -> bool {
2260 match err {
2261 LanceError::NotFound { .. } => true,
2262 LanceError::IO { source, .. } => source
2263 .downcast_ref::<ObjectStoreError>()
2264 .is_some_and(|source| matches!(source, ObjectStoreError::NotFound { .. })),
2265 LanceError::DatasetNotFound { source, .. } => {
2266 source
2267 .downcast_ref::<LanceError>()
2268 .is_some_and(|source| matches!(source, LanceError::NotFound { .. }))
2269 || source
2270 .downcast_ref::<ObjectStoreError>()
2271 .is_some_and(|source| matches!(source, ObjectStoreError::NotFound { .. }))
2272 }
2273 _ => false,
2274 }
2275 }
2276
2277 pub async fn list_manifest_table_locations(&self) -> Result<std::collections::HashSet<String>> {
2280 let filter = "object_type = 'table' AND NOT contains(object_id, '$')";
2281 let mut scanner = self.manifest_scanner().await?;
2282 scanner.filter(filter).map_err(|e| {
2283 lance_core::Error::from(NamespaceError::Internal {
2284 message: format!("Failed to filter: {:?}", e),
2285 })
2286 })?;
2287 scanner.project(&["location"]).map_err(|e| {
2288 lance_core::Error::from(NamespaceError::Internal {
2289 message: format!("Failed to project: {:?}", e),
2290 })
2291 })?;
2292
2293 let batches = Self::execute_scanner(scanner).await?;
2294 let mut locations = std::collections::HashSet::new();
2295
2296 for batch in batches {
2297 if batch.num_rows() == 0 {
2298 continue;
2299 }
2300 let location_array = Self::get_string_column(&batch, "location")?;
2301 for i in 0..location_array.len() {
2302 locations.insert(location_array.value(i).to_string());
2303 }
2304 }
2305
2306 Ok(locations)
2307 }
2308
2309 async fn insert_into_manifest(
2311 &self,
2312 object_id: String,
2313 object_type: ObjectType,
2314 location: Option<String>,
2315 ) -> Result<()> {
2316 self.insert_into_manifest_with_metadata(
2317 vec![ManifestEntry {
2318 object_id,
2319 object_type,
2320 location,
2321 metadata: None,
2322 }],
2323 None,
2324 )
2325 .await
2326 }
2327
2328 pub async fn insert_into_manifest_with_metadata(
2333 &self,
2334 entries: Vec<ManifestEntry>,
2335 base_objects: Option<Vec<String>>,
2336 ) -> Result<()> {
2337 self.merge_into_manifest_with_metadata(entries, base_objects, WhenMatched::Fail)
2338 .await
2339 }
2340
2341 async fn upsert_into_manifest_with_metadata(
2342 &self,
2343 entries: Vec<ManifestEntry>,
2344 base_objects: Option<Vec<String>>,
2345 ) -> Result<()> {
2346 self.merge_into_manifest_with_metadata(entries, base_objects, WhenMatched::UpdateAll)
2347 .await
2348 }
2349
2350 async fn merge_into_manifest_with_metadata(
2351 &self,
2352 entries: Vec<ManifestEntry>,
2353 base_objects: Option<Vec<String>>,
2354 when_matched: WhenMatched,
2355 ) -> Result<()> {
2356 if entries.is_empty() {
2357 return Ok(());
2358 }
2359
2360 self.rewrite_manifest("Failed to overwrite manifest", || {
2361 UpsertManifestMutation::new(entries.clone(), base_objects.clone(), when_matched.clone())
2362 })
2363 .await
2364 }
2365
2366 pub async fn delete_from_manifest(&self, object_id: &str) -> Result<()> {
2368 let object_id = object_id.to_string();
2369 self.rewrite_manifest("Failed to delete from manifest", || DeleteObjectMutation {
2370 object_id: object_id.clone(),
2371 deleted: false,
2372 })
2373 .await
2374 }
2375
2376 pub async fn register_table(&self, name: &str, location: String) -> Result<()> {
2378 let object_id = Self::build_object_id(&[], name);
2379 if self.manifest_contains_object(&object_id).await? {
2380 return Err(NamespaceError::Internal {
2381 message: format!("Table '{}' already exists", name),
2382 }
2383 .into());
2384 }
2385
2386 self.insert_into_manifest(object_id, ObjectType::Table, Some(location))
2387 .await
2388 }
2389
2390 async fn validate_namespace_levels_exist(&self, namespace_path: &[String]) -> Result<()> {
2392 for i in 1..=namespace_path.len() {
2393 let partial_path = &namespace_path[..i];
2394 let object_id = partial_path.join(DELIMITER);
2395 if !self.manifest_contains_object(&object_id).await? {
2396 return Err(NamespaceError::NamespaceNotFound {
2397 message: format!("parent namespace '{}'", object_id),
2398 }
2399 .into());
2400 }
2401 }
2402 Ok(())
2403 }
2404
2405 async fn query_manifest_for_namespace(&self, object_id: &str) -> Result<Option<NamespaceInfo>> {
2407 let escaped_id = object_id.replace('\'', "''");
2408 let filter = format!("object_id = '{}' AND object_type = 'namespace'", escaped_id);
2409 let mut scanner = self.manifest_scanner().await?;
2410 scanner.filter(&filter).map_err(|e| {
2411 lance_core::Error::from(NamespaceError::Internal {
2412 message: format!("Failed to filter: {:?}", e),
2413 })
2414 })?;
2415 scanner.project(&["object_id", "metadata"]).map_err(|e| {
2416 lance_core::Error::from(NamespaceError::Internal {
2417 message: format!("Failed to project: {:?}", e),
2418 })
2419 })?;
2420 let batches = Self::execute_scanner(scanner).await?;
2421
2422 let mut found_result: Option<NamespaceInfo> = None;
2423 let mut total_rows = 0;
2424
2425 for batch in batches {
2426 if batch.num_rows() == 0 {
2427 continue;
2428 }
2429
2430 total_rows += batch.num_rows();
2431 if total_rows > 1 {
2432 return Err(NamespaceError::Internal {
2433 message: format!(
2434 "Expected exactly 1 namespace with id '{}', found {}",
2435 object_id, total_rows
2436 ),
2437 }
2438 .into());
2439 }
2440
2441 let object_id_array = Self::get_string_column(&batch, "object_id")?;
2442 let metadata_array = Self::get_string_column(&batch, "metadata")?;
2443
2444 let object_id_str = object_id_array.value(0);
2445 let metadata = if !metadata_array.is_null(0) {
2446 let metadata_str = metadata_array.value(0);
2447 match serde_json::from_str::<HashMap<String, String>>(metadata_str) {
2448 Ok(map) => Some(map),
2449 Err(e) => {
2450 return Err(NamespaceError::Internal {
2451 message: format!(
2452 "Failed to deserialize metadata for namespace '{}': {}",
2453 object_id, e
2454 ),
2455 }
2456 .into());
2457 }
2458 }
2459 } else {
2460 None
2461 };
2462
2463 let (namespace, name) = Self::parse_object_id(object_id_str);
2464 found_result = Some(NamespaceInfo {
2465 namespace,
2466 name,
2467 metadata,
2468 });
2469 }
2470
2471 Ok(found_result)
2472 }
2473
2474 async fn open_manifest_table(
2476 root: &str,
2477 storage_options: &Option<HashMap<String, String>>,
2478 session: Option<Arc<Session>>,
2479 ) -> Result<DatasetConsistencyWrapper> {
2480 let manifest_path = format!("{}/{}", root, MANIFEST_TABLE_NAME);
2481 log::debug!("Attempting to load manifest from {}", manifest_path);
2482 let store_options = ObjectStoreParams {
2483 storage_options_accessor: storage_options.as_ref().map(|opts| {
2484 Arc::new(
2485 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2486 opts.clone(),
2487 ),
2488 )
2489 }),
2490 ..Default::default()
2491 };
2492 let read_params = ReadParams {
2493 session,
2494 store_options: Some(store_options),
2495 ..Default::default()
2496 };
2497 let dataset = DatasetBuilder::from_uri(&manifest_path)
2498 .with_read_params(read_params)
2499 .load()
2500 .await?;
2501 ensure_readable(dataset.metadata())?;
2502 Ok(DatasetConsistencyWrapper::new(dataset))
2503 }
2504
2505 async fn ensure_manifest_table_up_to_date(
2512 root: &str,
2513 storage_options: &Option<HashMap<String, String>>,
2514 session: Option<Arc<Session>>,
2515 ) -> Result<DatasetConsistencyWrapper> {
2516 let manifest_path = format!("{}/{}", root, MANIFEST_TABLE_NAME);
2517 log::debug!("Attempting to load manifest from {}", manifest_path);
2518 let store_options = ObjectStoreParams {
2519 storage_options_accessor: storage_options.as_ref().map(|opts| {
2520 Arc::new(
2521 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2522 opts.clone(),
2523 ),
2524 )
2525 }),
2526 ..Default::default()
2527 };
2528 let read_params = ReadParams {
2529 session: session.clone(),
2530 store_options: Some(store_options.clone()),
2531 ..Default::default()
2532 };
2533 let dataset_result = DatasetBuilder::from_uri(&manifest_path)
2534 .with_read_params(read_params)
2535 .load()
2536 .await;
2537 match dataset_result {
2538 Ok(mut dataset) => {
2539 ensure_readable(dataset.metadata())?;
2542
2543 let needs_pk_migration = dataset
2545 .schema()
2546 .field("object_id")
2547 .map(|f| {
2548 !f.metadata
2549 .contains_key(LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
2550 })
2551 .unwrap_or(false);
2552
2553 if needs_pk_migration {
2554 ensure_writable(dataset.metadata())?;
2557 log::info!(
2558 "Migrating __manifest table to add primary key metadata on object_id"
2559 );
2560 dataset
2561 .update_field_metadata()
2562 .update("object_id", [(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "0")])
2563 .map_err(|e| {
2564 lance_core::Error::from(NamespaceError::Internal {
2565 message: format!(
2566 "Failed to find object_id field for migration: {:?}",
2567 e
2568 ),
2569 })
2570 })?
2571 .await
2572 .map_err(|e| {
2573 lance_core::Error::from(NamespaceError::Internal {
2574 message: format!("Failed to migrate primary key metadata: {:?}", e),
2575 })
2576 })?;
2577 }
2578
2579 Ok(DatasetConsistencyWrapper::new(dataset))
2580 }
2581 Err(err) if Self::is_not_found_load_error(&err) => {
2582 log::info!("Creating new manifest table at {}", manifest_path);
2583 let schema = Self::manifest_schema();
2584 let empty_batch = RecordBatch::new_empty(schema.clone());
2585 let reader = RecordBatchIterator::new(vec![Ok(empty_batch)], schema.clone());
2586
2587 let store_params = ObjectStoreParams {
2588 storage_options_accessor: storage_options.as_ref().map(|opts| {
2589 Arc::new(
2590 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2591 opts.clone(),
2592 ),
2593 )
2594 }),
2595 ..Default::default()
2596 };
2597 let write_params = WriteParams {
2598 session: session.clone(),
2599 store_params: Some(store_params),
2600 ..Default::default()
2601 };
2602
2603 let dataset =
2604 Dataset::write(Box::new(reader), &manifest_path, Some(write_params)).await;
2605
2606 match dataset {
2608 Ok(dataset) => {
2609 log::info!(
2610 "Successfully created manifest table at {}, version={}, uri={}",
2611 manifest_path,
2612 dataset.version().version,
2613 dataset.uri()
2614 );
2615 Ok(DatasetConsistencyWrapper::new(dataset))
2616 }
2617 Err(ref e)
2618 if matches!(
2619 e,
2620 LanceError::DatasetAlreadyExists { .. }
2621 | LanceError::CommitConflict { .. }
2622 | LanceError::IncompatibleTransaction { .. }
2623 | LanceError::RetryableCommitConflict { .. }
2624 ) =>
2625 {
2626 log::info!(
2628 "Manifest table was created by another process, loading it: {}",
2629 manifest_path
2630 );
2631 let recovery_store_options = ObjectStoreParams {
2632 storage_options_accessor: storage_options.as_ref().map(|opts| {
2633 Arc::new(
2634 lance_io::object_store::StorageOptionsAccessor::with_static_options(
2635 opts.clone(),
2636 ),
2637 )
2638 }),
2639 ..Default::default()
2640 };
2641 let recovery_read_params = ReadParams {
2642 session,
2643 store_options: Some(recovery_store_options),
2644 ..Default::default()
2645 };
2646 let dataset = DatasetBuilder::from_uri(&manifest_path)
2647 .with_read_params(recovery_read_params)
2648 .load()
2649 .await
2650 .map_err(|e| {
2651 lance_core::Error::from(NamespaceError::Internal {
2652 message: format!(
2653 "Failed to load manifest dataset after creation conflict: {}",
2654 e
2655 ),
2656 })
2657 })?;
2658 Ok(DatasetConsistencyWrapper::new(dataset))
2659 }
2660 Err(e) => Err(lance_core::Error::from(NamespaceError::Internal {
2661 message: format!("Failed to create manifest dataset: {:?}", e),
2662 })),
2663 }
2664 }
2665 Err(err) => Err(err),
2666 }
2667 }
2668
2669 fn apply_pagination(
2674 names: &mut Vec<String>,
2675 page_token: Option<String>,
2676 limit: Option<i32>,
2677 ) -> Option<String> {
2678 names.sort();
2679
2680 if let Some(start_after) = page_token {
2681 if let Some(index) = names
2682 .iter()
2683 .position(|name| name.as_str() > start_after.as_str())
2684 {
2685 names.drain(0..index);
2686 } else {
2687 names.clear();
2688 }
2689 }
2690
2691 if let Some(limit) = limit
2692 && limit >= 0
2693 {
2694 let limit = limit as usize;
2695 if names.len() > limit {
2696 let next_page_token = if limit > 0 {
2697 Some(names[limit - 1].clone())
2698 } else {
2699 None
2700 };
2701 names.truncate(limit);
2702 return next_page_token;
2703 }
2704 }
2705
2706 None
2707 }
2708}
2709
2710#[async_trait]
2711impl LanceNamespace for ManifestNamespace {
2712 fn namespace_id(&self) -> String {
2713 self.root.clone()
2714 }
2715
2716 async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
2717 let namespace_id = request.id.as_ref().ok_or_else(|| {
2718 lance_core::Error::from(NamespaceError::InvalidInput {
2719 message: "Namespace ID is required".to_string(),
2720 })
2721 })?;
2722
2723 let filter = if namespace_id.is_empty() {
2725 "object_type = 'table' AND NOT contains(object_id, '$')".to_string()
2727 } else {
2728 let prefix = namespace_id.join(DELIMITER);
2730 format!(
2731 "object_type = 'table' AND starts_with(object_id, '{}{}') AND NOT contains(substring(object_id, {}), '$')",
2732 prefix,
2733 DELIMITER,
2734 prefix.len() + 2
2735 )
2736 };
2737
2738 let mut scanner = self.manifest_scanner().await?;
2739 scanner.filter(&filter).map_err(|e| {
2740 lance_core::Error::from(NamespaceError::Internal {
2741 message: format!("Failed to filter: {:?}", e),
2742 })
2743 })?;
2744 scanner.project(&["object_id", "location"]).map_err(|e| {
2745 lance_core::Error::from(NamespaceError::Internal {
2746 message: format!("Failed to project: {:?}", e),
2747 })
2748 })?;
2749
2750 let batches = Self::execute_scanner(scanner).await?;
2751
2752 let mut table_entries = Vec::new();
2753 for batch in batches {
2754 if batch.num_rows() == 0 {
2755 continue;
2756 }
2757
2758 let object_id_array = Self::get_string_column(&batch, "object_id")?;
2759 let location_array = Self::get_string_column(&batch, "location")?;
2760 for i in 0..batch.num_rows() {
2761 let object_id = object_id_array.value(i);
2762 let location = location_array.value(i);
2763 let (_namespace, name) = Self::parse_object_id(object_id);
2764 table_entries.push((name, location.to_string()));
2765 }
2766 }
2767
2768 let mut tables: Vec<String> = if request.include_declared.unwrap_or(true) {
2769 table_entries.into_iter().map(|(name, _)| name).collect()
2770 } else {
2771 let mut stream = futures::stream::iter(table_entries.into_iter().map(
2772 |(name, location)| async move {
2773 if self.location_has_actual_manifests(&location).await? {
2778 Ok::<Option<String>, Error>(Some(name))
2779 } else {
2780 Ok::<Option<String>, Error>(None)
2781 }
2782 },
2783 ))
2784 .buffered(DECLARED_FILTER_CONCURRENCY);
2785
2786 let mut filtered = Vec::new();
2787 while let Some(result) = stream.next().await {
2788 if let Some(name) = result? {
2789 filtered.push(name);
2790 }
2791 }
2792 filtered
2793 };
2794
2795 let next_page_token =
2796 Self::apply_pagination(&mut tables, request.page_token, request.limit);
2797 let mut response = ListTablesResponse::new(tables);
2798 response.page_token = next_page_token;
2799 Ok(response)
2800 }
2801
2802 async fn describe_table(&self, request: DescribeTableRequest) -> Result<DescribeTableResponse> {
2803 let table_id = request.id.as_ref().ok_or_else(|| {
2804 lance_core::Error::from(NamespaceError::InvalidInput {
2805 message: "Table ID is required".to_string(),
2806 })
2807 })?;
2808
2809 if table_id.is_empty() {
2810 return Err(NamespaceError::InvalidInput {
2811 message: "Table ID cannot be empty".to_string(),
2812 }
2813 .into());
2814 }
2815
2816 let object_id = Self::str_object_id(table_id);
2817 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
2818
2819 let table_name = table_id.last().cloned().unwrap_or_default();
2821 let namespace_id: Vec<String> = if table_id.len() > 1 {
2822 table_id[..table_id.len() - 1].to_vec()
2823 } else {
2824 vec![]
2825 };
2826
2827 let load_detailed_metadata = request.load_detailed_metadata.unwrap_or(false);
2828 let should_check_declared =
2829 load_detailed_metadata || request.check_declared.unwrap_or(false);
2830 let vend_credentials = request.vend_credentials.unwrap_or(true);
2832
2833 match table_info {
2834 Some(info) => {
2835 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
2837
2838 let storage_options = if vend_credentials {
2839 self.storage_options.clone()
2840 } else {
2841 None
2842 };
2843 let is_only_declared = if should_check_declared {
2844 Some(!self.location_has_actual_manifests(&info.location).await?)
2845 } else {
2846 None
2847 };
2848
2849 if !load_detailed_metadata {
2850 return Ok(DescribeTableResponse {
2851 table: Some(table_name),
2852 namespace: Some(namespace_id),
2853 location: Some(table_uri.clone()),
2854 table_uri: Some(table_uri),
2855 storage_options,
2856 properties: info.metadata,
2857 is_only_declared,
2858 ..Default::default()
2859 });
2860 }
2861
2862 if is_only_declared == Some(true) {
2863 return Ok(DescribeTableResponse {
2864 table: Some(table_name),
2865 namespace: Some(namespace_id),
2866 location: Some(table_uri.clone()),
2867 table_uri: Some(table_uri),
2868 storage_options,
2869 properties: info.metadata,
2870 is_only_declared,
2871 ..Default::default()
2872 });
2873 }
2874
2875 let mut builder = DatasetBuilder::from_uri(&table_uri);
2876 if let Some(opts) = &self.storage_options {
2877 builder = builder.with_storage_options(opts.clone());
2878 }
2879 if let Some(session) = &self.session {
2880 builder = builder.with_session(session.clone());
2881 }
2882
2883 match builder.load().await {
2884 Ok(mut dataset) => {
2885 if let Some(requested_version) = request.version {
2887 dataset = dataset.checkout_version(requested_version as u64).await?;
2888 }
2889
2890 let version = dataset.version().version;
2891 let lance_schema = dataset.schema();
2892 let arrow_schema: arrow_schema::Schema = lance_schema.into();
2893 let json_schema = arrow_schema_to_json(&arrow_schema)?;
2894
2895 Ok(DescribeTableResponse {
2896 table: Some(table_name.clone()),
2897 namespace: Some(namespace_id.clone()),
2898 version: Some(version as i64),
2899 location: Some(table_uri.clone()),
2900 table_uri: Some(table_uri),
2901 schema: Some(Box::new(json_schema)),
2902 storage_options,
2903 properties: info.metadata.clone(),
2904 is_only_declared,
2905 ..Default::default()
2906 })
2907 }
2908 Err(err) => Err(NamespaceError::Internal {
2909 message: format!(
2910 "Table exists in manifest but failed to load dataset '{}': {}",
2911 object_id, err
2912 ),
2913 }
2914 .into()),
2915 }
2916 }
2917 None => Err(NamespaceError::TableNotFound {
2918 message: Self::format_table_id(table_id),
2919 }
2920 .into()),
2921 }
2922 }
2923
2924 async fn table_exists(&self, request: TableExistsRequest) -> Result<()> {
2925 let table_id = request.id.as_ref().ok_or_else(|| {
2926 lance_core::Error::from(NamespaceError::InvalidInput {
2927 message: "Table ID is required".to_string(),
2928 })
2929 })?;
2930
2931 if table_id.is_empty() {
2932 return Err(NamespaceError::InvalidInput {
2933 message: "Table ID cannot be empty".to_string(),
2934 }
2935 .into());
2936 }
2937
2938 let object_id = Self::str_object_id(table_id);
2939 let exists = self.manifest_contains_object(&object_id).await?;
2940 if exists {
2941 Ok(())
2942 } else {
2943 Err(NamespaceError::TableNotFound {
2944 message: Self::format_table_id(table_id),
2945 }
2946 .into())
2947 }
2948 }
2949
2950 async fn create_table(
2951 &self,
2952 request: CreateTableRequest,
2953 data: Bytes,
2954 ) -> Result<CreateTableResponse> {
2955 let table_id = request.id.as_ref().ok_or_else(|| {
2956 lance_core::Error::from(NamespaceError::InvalidInput {
2957 message: "Table ID is required".to_string(),
2958 })
2959 })?;
2960
2961 if table_id.is_empty() {
2962 return Err(NamespaceError::InvalidInput {
2963 message: "Table ID cannot be empty".to_string(),
2964 }
2965 .into());
2966 }
2967
2968 let (namespace, table_name) = Self::split_object_id(table_id);
2969 let object_id = Self::build_object_id(&namespace, &table_name);
2970
2971 self.ensure_manifest_writable().await?;
2974
2975 let existing_table = self.query_manifest_for_table(&object_id).await?;
2976 let existing_has_manifests = if let Some(existing_table) = &existing_table {
2977 Some(
2978 self.location_has_actual_manifests(&existing_table.location)
2979 .await?,
2980 )
2981 } else {
2982 None
2983 };
2984
2985 if existing_has_manifests == Some(false)
2986 && request
2987 .properties
2988 .as_ref()
2989 .is_some_and(|properties| !properties.is_empty())
2990 {
2991 return Err(NamespaceError::InvalidInput {
2992 message: format!(
2993 "create_table cannot set properties for already declared table '{}'",
2994 object_id
2995 ),
2996 }
2997 .into());
2998 }
2999
3000 let create_mode = if existing_has_manifests == Some(false) {
3001 CreateTableMode::Create
3002 } else {
3003 CreateTableMode::parse(request.mode.as_deref())?
3004 };
3005 let dir_name = if let Some(existing_table) = &existing_table {
3006 existing_table.location.clone()
3007 } else if namespace.is_empty() && self.dir_listing_enabled {
3008 format!("{}.lance", table_name)
3009 } else {
3010 Self::generate_dir_name(&object_id)
3011 };
3012 let table_uri = Self::construct_full_uri(&self.root, &dir_name)?;
3013 let overwriting_existing_table =
3014 existing_has_manifests == Some(true) && create_mode == CreateTableMode::Overwrite;
3015
3016 if existing_has_manifests == Some(true) {
3017 match create_mode {
3018 CreateTableMode::Create => {
3019 return Err(NamespaceError::TableAlreadyExists {
3020 message: table_name.clone(),
3021 }
3022 .into());
3023 }
3024 CreateTableMode::ExistOk => {
3025 let properties = existing_table
3026 .as_ref()
3027 .and_then(|table| table.metadata.clone());
3028 return Ok(CreateTableResponse {
3029 location: Some(table_uri),
3030 storage_options: self.storage_options.clone(),
3031 properties,
3032 ..Default::default()
3033 });
3034 }
3035 CreateTableMode::Overwrite => {}
3036 }
3037 }
3038
3039 if data.is_empty() {
3041 return Err(NamespaceError::InvalidInput {
3042 message: "Request data (Arrow IPC stream) is required for create_table".to_string(),
3043 }
3044 .into());
3045 }
3046
3047 let cursor = Cursor::new(data.to_vec());
3049 let stream_reader = StreamReader::try_new(cursor, None).map_err(|e| {
3050 lance_core::Error::from(NamespaceError::Internal {
3051 message: format!("Failed to read IPC stream: {:?}", e),
3052 })
3053 })?;
3054
3055 let batches: Vec<RecordBatch> = stream_reader
3056 .collect::<std::result::Result<Vec<_>, _>>()
3057 .map_err(|e| {
3058 lance_core::Error::from(NamespaceError::Internal {
3059 message: format!("Failed to collect batches: {:?}", e),
3060 })
3061 })?;
3062
3063 if batches.is_empty() {
3064 return Err(NamespaceError::Internal {
3065 message: "No data provided for table creation".to_string(),
3066 }
3067 .into());
3068 }
3069
3070 let schema = batches[0].schema();
3071 let batch_results: Vec<std::result::Result<RecordBatch, arrow_schema::ArrowError>> =
3072 batches.into_iter().map(Ok).collect();
3073 let reader = RecordBatchIterator::new(batch_results, schema);
3074
3075 let mut write_storage_options = self.storage_options.clone().unwrap_or_default();
3076 if let Some(request_storage_options) = request.storage_options.as_ref() {
3077 write_storage_options.extend(request_storage_options.clone());
3078 }
3079
3080 let store_params = ObjectStoreParams {
3081 storage_options_accessor: (!write_storage_options.is_empty()).then(|| {
3082 Arc::new(
3083 lance_io::object_store::StorageOptionsAccessor::with_static_options(
3084 write_storage_options,
3085 ),
3086 )
3087 }),
3088 ..Default::default()
3089 };
3090 let write_params = WriteParams {
3091 mode: create_mode.write_mode(),
3092 session: self.session.clone(),
3093 store_params: Some(store_params),
3094 ..Default::default()
3095 };
3096 let dataset = Dataset::write(Box::new(reader), &table_uri, Some(write_params))
3097 .await
3098 .map_err(|e| {
3099 lance_core::Error::from(NamespaceError::Internal {
3100 message: format!("Failed to write dataset: {:?}", e),
3101 })
3102 })?;
3103 let version = dataset.version().version as i64;
3104
3105 if overwriting_existing_table {
3106 let metadata =
3107 Self::serialize_metadata(request.properties.as_ref(), "table", &object_id)?;
3108 self.upsert_into_manifest_with_metadata(
3109 vec![ManifestEntry {
3110 object_id,
3111 object_type: ObjectType::Table,
3112 location: Some(dir_name),
3113 metadata,
3114 }],
3115 None,
3116 )
3117 .await?;
3118
3119 Ok(CreateTableResponse {
3120 version: Some(version),
3121 location: Some(table_uri),
3122 storage_options: self.storage_options.clone(),
3123 properties: request.properties,
3124 ..Default::default()
3125 })
3126 } else {
3127 match existing_table {
3128 Some(existing_table) => Ok(CreateTableResponse {
3129 version: Some(version),
3130 location: Some(table_uri),
3131 storage_options: self.storage_options.clone(),
3132 properties: existing_table.metadata,
3133 ..Default::default()
3134 }),
3135 None => {
3136 let metadata =
3137 Self::serialize_metadata(request.properties.as_ref(), "table", &object_id)?;
3138 self.insert_into_manifest_with_metadata(
3140 vec![ManifestEntry {
3141 object_id,
3142 object_type: ObjectType::Table,
3143 location: Some(dir_name.clone()),
3144 metadata,
3145 }],
3146 None,
3147 )
3148 .await?;
3149
3150 Ok(CreateTableResponse {
3151 version: Some(version),
3152 location: Some(table_uri),
3153 storage_options: self.storage_options.clone(),
3154 properties: request.properties,
3155 ..Default::default()
3156 })
3157 }
3158 }
3159 }
3160 }
3161
3162 async fn drop_table(&self, request: DropTableRequest) -> Result<DropTableResponse> {
3163 let table_id = request.id.as_ref().ok_or_else(|| {
3164 lance_core::Error::from(NamespaceError::InvalidInput {
3165 message: "Table ID is required".to_string(),
3166 })
3167 })?;
3168
3169 if table_id.is_empty() {
3170 return Err(NamespaceError::InvalidInput {
3171 message: "Table ID cannot be empty".to_string(),
3172 }
3173 .into());
3174 }
3175
3176 let (namespace, table_name) = Self::split_object_id(table_id);
3177 let object_id = Self::build_object_id(&namespace, &table_name);
3178
3179 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3181
3182 match table_info {
3183 Some(info) => {
3184 self.delete_from_manifest(&object_id).boxed().await?;
3186
3187 let table_path = self.base_path.clone().join(info.location.as_str());
3189 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3190
3191 self.object_store
3193 .remove_dir_all(table_path)
3194 .boxed()
3195 .await
3196 .map_err(|e| {
3197 lance_core::Error::from(NamespaceError::Internal {
3198 message: format!("Failed to delete table directory: {:?}", e),
3199 })
3200 })?;
3201
3202 Ok(DropTableResponse {
3203 id: request.id.clone(),
3204 location: Some(table_uri),
3205 ..Default::default()
3206 })
3207 }
3208 None => Err(NamespaceError::TableNotFound {
3209 message: table_name.to_string(),
3210 }
3211 .into()),
3212 }
3213 }
3214
3215 async fn list_namespaces(
3216 &self,
3217 request: ListNamespacesRequest,
3218 ) -> Result<ListNamespacesResponse> {
3219 let parent_namespace = request.id.as_ref().ok_or_else(|| {
3220 lance_core::Error::from(NamespaceError::InvalidInput {
3221 message: "Namespace ID is required".to_string(),
3222 })
3223 })?;
3224
3225 let filter = if parent_namespace.is_empty() {
3227 "object_type = 'namespace' AND NOT contains(object_id, '$')".to_string()
3229 } else {
3230 let prefix = parent_namespace.join(DELIMITER);
3232 format!(
3233 "object_type = 'namespace' AND starts_with(object_id, '{}{}') AND NOT contains(substring(object_id, {}), '$')",
3234 prefix,
3235 DELIMITER,
3236 prefix.len() + 2
3237 )
3238 };
3239
3240 let mut scanner = self.manifest_scanner().await?;
3241 scanner.filter(&filter).map_err(|e| {
3242 lance_core::Error::from(NamespaceError::Internal {
3243 message: format!("Failed to filter: {:?}", e),
3244 })
3245 })?;
3246 scanner.project(&["object_id"]).map_err(|e| {
3247 lance_core::Error::from(NamespaceError::Internal {
3248 message: format!("Failed to project: {:?}", e),
3249 })
3250 })?;
3251
3252 let batches = Self::execute_scanner(scanner).await?;
3253 let mut namespaces = Vec::new();
3254
3255 for batch in batches {
3256 if batch.num_rows() == 0 {
3257 continue;
3258 }
3259
3260 let object_id_array = Self::get_string_column(&batch, "object_id")?;
3261 for i in 0..batch.num_rows() {
3262 let object_id = object_id_array.value(i);
3263 let (_namespace, name) = Self::parse_object_id(object_id);
3264 namespaces.push(name);
3265 }
3266 }
3267
3268 let next_page_token =
3269 Self::apply_pagination(&mut namespaces, request.page_token, request.limit);
3270 let mut response = ListNamespacesResponse::new(namespaces);
3271 response.page_token = next_page_token;
3272 Ok(response)
3273 }
3274
3275 async fn describe_namespace(
3276 &self,
3277 request: DescribeNamespaceRequest,
3278 ) -> Result<DescribeNamespaceResponse> {
3279 let namespace_id = request.id.as_ref().ok_or_else(|| {
3280 lance_core::Error::from(NamespaceError::InvalidInput {
3281 message: "Namespace ID is required".to_string(),
3282 })
3283 })?;
3284
3285 if namespace_id.is_empty() {
3287 #[allow(clippy::needless_update)]
3288 return Ok(DescribeNamespaceResponse {
3289 properties: Some(HashMap::new()),
3290 ..Default::default()
3291 });
3292 }
3293
3294 let object_id = namespace_id.join(DELIMITER);
3296 let namespace_info = self.query_manifest_for_namespace(&object_id).await?;
3297
3298 match namespace_info {
3299 #[allow(clippy::needless_update)]
3300 Some(info) => Ok(DescribeNamespaceResponse {
3301 properties: info.metadata,
3302 ..Default::default()
3303 }),
3304 None => Err(NamespaceError::NamespaceNotFound {
3305 message: object_id.to_string(),
3306 }
3307 .into()),
3308 }
3309 }
3310
3311 async fn create_namespace(
3312 &self,
3313 request: CreateNamespaceRequest,
3314 ) -> Result<CreateNamespaceResponse> {
3315 let namespace_id = request.id.as_ref().ok_or_else(|| {
3316 lance_core::Error::from(NamespaceError::InvalidInput {
3317 message: "Namespace ID is required".to_string(),
3318 })
3319 })?;
3320
3321 if namespace_id.is_empty() {
3323 return Err(NamespaceError::NamespaceAlreadyExists {
3324 message: "root namespace".to_string(),
3325 }
3326 .into());
3327 }
3328
3329 if namespace_id.len() > 1 {
3331 self.validate_namespace_levels_exist(&namespace_id[..namespace_id.len() - 1])
3332 .await?;
3333 }
3334
3335 let object_id = namespace_id.join(DELIMITER);
3336 if self.manifest_contains_object(&object_id).await? {
3337 return Err(NamespaceError::NamespaceAlreadyExists {
3338 message: object_id.to_string(),
3339 }
3340 .into());
3341 }
3342
3343 let metadata =
3344 Self::serialize_metadata(request.properties.as_ref(), "namespace", &object_id)?;
3345
3346 self.insert_into_manifest_with_metadata(
3347 vec![ManifestEntry {
3348 object_id,
3349 object_type: ObjectType::Namespace,
3350 location: None,
3351 metadata,
3352 }],
3353 None,
3354 )
3355 .await?;
3356
3357 Ok(CreateNamespaceResponse {
3358 properties: request.properties,
3359 ..Default::default()
3360 })
3361 }
3362
3363 async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
3364 let namespace_id = request.id.as_ref().ok_or_else(|| {
3365 lance_core::Error::from(NamespaceError::InvalidInput {
3366 message: "Namespace ID is required".to_string(),
3367 })
3368 })?;
3369
3370 if namespace_id.is_empty() {
3372 return Err(NamespaceError::InvalidInput {
3373 message: "Root namespace cannot be dropped".to_string(),
3374 }
3375 .into());
3376 }
3377
3378 let object_id = namespace_id.join(DELIMITER);
3379
3380 if !self.manifest_contains_object(&object_id).boxed().await? {
3382 return Err(NamespaceError::NamespaceNotFound {
3383 message: object_id.to_string(),
3384 }
3385 .into());
3386 }
3387
3388 let escaped_id = object_id.replace('\'', "''");
3390 let prefix = format!("{}{}", escaped_id, DELIMITER);
3391 let filter = format!("starts_with(object_id, '{}')", prefix);
3392 let mut scanner = self.manifest_scanner().boxed().await?;
3393 scanner.filter(&filter).map_err(|e| {
3394 lance_core::Error::from(NamespaceError::Internal {
3395 message: format!("Failed to filter: {:?}", e),
3396 })
3397 })?;
3398 scanner.project::<&str>(&[]).map_err(|e| {
3399 lance_core::Error::from(NamespaceError::Internal {
3400 message: format!("Failed to project: {:?}", e),
3401 })
3402 })?;
3403 scanner.with_row_id();
3404 let count = scanner.count_rows().boxed().await.map_err(|e| {
3405 lance_core::Error::from(NamespaceError::Internal {
3406 message: format!("Failed to count rows: {:?}", e),
3407 })
3408 })?;
3409
3410 if count > 0 {
3411 return Err(NamespaceError::NamespaceNotEmpty {
3412 message: format!("'{}' (contains {} child objects)", object_id, count),
3413 }
3414 .into());
3415 }
3416
3417 self.delete_from_manifest(&object_id).boxed().await?;
3418
3419 Ok(DropNamespaceResponse::default())
3420 }
3421
3422 async fn namespace_exists(&self, request: NamespaceExistsRequest) -> Result<()> {
3423 let namespace_id = request.id.as_ref().ok_or_else(|| {
3424 lance_core::Error::from(NamespaceError::InvalidInput {
3425 message: "Namespace ID is required".to_string(),
3426 })
3427 })?;
3428
3429 if namespace_id.is_empty() {
3431 return Ok(());
3432 }
3433
3434 let object_id = namespace_id.join(DELIMITER);
3435 if self.manifest_contains_object(&object_id).await? {
3436 Ok(())
3437 } else {
3438 Err(NamespaceError::NamespaceNotFound {
3439 message: object_id.to_string(),
3440 }
3441 .into())
3442 }
3443 }
3444
3445 async fn declare_table(&self, request: DeclareTableRequest) -> Result<DeclareTableResponse> {
3446 let table_id = request.id.as_ref().ok_or_else(|| {
3447 lance_core::Error::from(NamespaceError::InvalidInput {
3448 message: "Table ID is required".to_string(),
3449 })
3450 })?;
3451
3452 if table_id.is_empty() {
3453 return Err(NamespaceError::InvalidInput {
3454 message: "Table ID cannot be empty".to_string(),
3455 }
3456 .into());
3457 }
3458
3459 let (namespace, table_name) = Self::split_object_id(table_id);
3460 let object_id = Self::build_object_id(&namespace, &table_name);
3461
3462 let existing = self.query_manifest_for_table(&object_id).await?;
3464 if existing.is_some() {
3465 return Err(NamespaceError::TableAlreadyExists {
3466 message: table_name.to_string(),
3467 }
3468 .into());
3469 }
3470
3471 let dir_name = if namespace.is_empty() && self.dir_listing_enabled {
3475 format!("{}.lance", table_name)
3477 } else {
3478 Self::generate_dir_name(&object_id)
3480 };
3481 let table_path = self.base_path.clone().join(dir_name.as_str());
3482 let table_uri = Self::construct_full_uri(&self.root, &dir_name)?;
3483
3484 if let Some(req_location) = &request.location {
3486 let req_location = req_location.trim_end_matches('/');
3487 if req_location != table_uri {
3488 return Err(NamespaceError::InvalidInput {
3489 message: format!(
3490 "Cannot declare table {} at location {}, must be at location {}",
3491 table_name, req_location, table_uri
3492 ),
3493 }
3494 .into());
3495 }
3496 }
3497
3498 self.ensure_manifest_writable().await?;
3499
3500 let reserved_file_path = table_path.clone().join(".lance-reserved");
3504 super::put_marker_file_atomic(
3505 &self.object_store,
3506 &reserved_file_path,
3507 &format!("table {}", table_name),
3508 )
3509 .await
3510 .map_err(|e| match e {
3511 super::MarkerFileError::AlreadyExists { .. } => {
3512 lance_core::Error::from(NamespaceError::TableAlreadyExists {
3513 message: table_name.to_string(),
3514 })
3515 }
3516 super::MarkerFileError::Other { message } => {
3517 lance_core::Error::from(NamespaceError::Internal { message })
3518 }
3519 })?;
3520
3521 let metadata = Self::serialize_metadata(request.properties.as_ref(), "table", &object_id)?;
3522
3523 self.insert_into_manifest_with_metadata(
3525 vec![ManifestEntry {
3526 object_id,
3527 object_type: ObjectType::Table,
3528 location: Some(dir_name),
3529 metadata,
3530 }],
3531 None,
3532 )
3533 .await?;
3534
3535 log::info!(
3536 "Declared table '{}' in manifest at {}",
3537 table_name,
3538 table_uri
3539 );
3540
3541 let vend_credentials = request.vend_credentials.unwrap_or(true);
3543 let storage_options = if vend_credentials {
3544 self.storage_options.clone()
3545 } else {
3546 None
3547 };
3548
3549 Ok(DeclareTableResponse {
3550 location: Some(table_uri),
3551 storage_options,
3552 properties: request.properties,
3553 ..Default::default()
3554 })
3555 }
3556
3557 async fn register_table(&self, request: RegisterTableRequest) -> Result<RegisterTableResponse> {
3558 let table_id = request.id.as_ref().ok_or_else(|| {
3559 lance_core::Error::from(NamespaceError::InvalidInput {
3560 message: "Table ID is required".to_string(),
3561 })
3562 })?;
3563
3564 if table_id.is_empty() {
3565 return Err(NamespaceError::InvalidInput {
3566 message: "Table ID cannot be empty".to_string(),
3567 }
3568 .into());
3569 }
3570
3571 let location = request.location.clone();
3572
3573 if location.contains("://") {
3576 return Err(NamespaceError::InvalidInput {
3577 message: format!(
3578 "Absolute URIs are not allowed for register_table. Location must be a relative path within the root directory: {}",
3579 location
3580 ),
3581 }
3582 .into());
3583 }
3584
3585 if location.starts_with('/') {
3586 return Err(NamespaceError::InvalidInput {
3587 message: format!(
3588 "Absolute paths are not allowed for register_table. Location must be a relative path within the root directory: {}",
3589 location
3590 ),
3591 }
3592 .into());
3593 }
3594
3595 if location.contains("..") {
3597 return Err(NamespaceError::InvalidInput {
3598 message: format!(
3599 "Path traversal is not allowed. Location must be a relative path within the root directory: {}",
3600 location
3601 ),
3602 }
3603 .into());
3604 }
3605
3606 let (namespace, table_name) = Self::split_object_id(table_id);
3607 let object_id = Self::build_object_id(&namespace, &table_name);
3608
3609 if !namespace.is_empty() {
3611 self.validate_namespace_levels_exist(&namespace).await?;
3612 }
3613
3614 if self.manifest_contains_object(&object_id).await? {
3616 return Err(NamespaceError::TableAlreadyExists {
3617 message: object_id.to_string(),
3618 }
3619 .into());
3620 }
3621
3622 self.insert_into_manifest(object_id, ObjectType::Table, Some(location.clone()))
3624 .await?;
3625
3626 Ok(RegisterTableResponse {
3627 location: Some(location),
3628 ..Default::default()
3629 })
3630 }
3631
3632 async fn deregister_table(
3633 &self,
3634 request: DeregisterTableRequest,
3635 ) -> Result<DeregisterTableResponse> {
3636 let table_id = request.id.as_ref().ok_or_else(|| {
3637 lance_core::Error::from(NamespaceError::InvalidInput {
3638 message: "Table ID is required".to_string(),
3639 })
3640 })?;
3641
3642 if table_id.is_empty() {
3643 return Err(NamespaceError::InvalidInput {
3644 message: "Table ID cannot be empty".to_string(),
3645 }
3646 .into());
3647 }
3648
3649 let (namespace, table_name) = Self::split_object_id(table_id);
3650 let object_id = Self::build_object_id(&namespace, &table_name);
3651
3652 let table_info = self.query_manifest_for_table(&object_id).await?;
3654
3655 let table_uri = match table_info {
3656 Some(info) => {
3657 self.delete_from_manifest(&object_id).boxed().await?;
3659 Self::construct_full_uri(&self.root, &info.location)?
3660 }
3661 None => {
3662 return Err(NamespaceError::TableNotFound {
3663 message: object_id.to_string(),
3664 }
3665 .into());
3666 }
3667 };
3668
3669 Ok(DeregisterTableResponse {
3670 id: request.id.clone(),
3671 location: Some(table_uri),
3672 ..Default::default()
3673 })
3674 }
3675
3676 async fn alter_table_add_columns(
3681 &self,
3682 request: AlterTableAddColumnsRequest,
3683 ) -> Result<AlterTableAddColumnsResponse> {
3684 let table_id = request
3685 .id
3686 .as_ref()
3687 .ok_or_else(|| Error::invalid_input_source("Table ID is required".into()))?;
3688
3689 if table_id.is_empty() {
3690 return Err(Error::invalid_input_source(
3691 "Table ID cannot be empty".into(),
3692 ));
3693 }
3694
3695 let object_id = Self::str_object_id(table_id);
3696 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3697
3698 match table_info {
3699 Some(info) => {
3700 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3701 let mut builder = DatasetBuilder::from_uri(&table_uri);
3704 if let Some(opts) = &self.storage_options {
3705 builder = builder.with_storage_options(opts.clone());
3706 }
3707 if let Some(session) = &self.session {
3708 builder = builder.with_session(session.clone());
3709 }
3710 let mut dataset = builder.load().await.map_err(|e| {
3711 Error::io_source(box_error(std::io::Error::other(format!(
3712 "Failed to open dataset: {}",
3713 e
3714 ))))
3715 })?;
3716
3717 let sql_expressions = super::build_sql_expressions(&request.new_columns)?;
3719
3720 dataset
3721 .add_columns(
3722 lance::dataset::NewColumnTransform::SqlExpressions(sql_expressions),
3723 None,
3724 None,
3725 )
3726 .await
3727 .map_err(|e| {
3728 convert_lance_commit_error(&e, "add_columns", Some(&object_id))
3732 })?;
3733
3734 let version = dataset.version().version as i64;
3735 Ok(AlterTableAddColumnsResponse::new(version))
3736 }
3737 None => Err(NamespaceError::TableNotFound { message: object_id }.into()),
3738 }
3739 }
3740
3741 async fn alter_table_alter_columns(
3746 &self,
3747 request: AlterTableAlterColumnsRequest,
3748 ) -> Result<AlterTableAlterColumnsResponse> {
3749 let table_id = request
3750 .id
3751 .as_ref()
3752 .ok_or_else(|| Error::invalid_input_source("Table ID is required".into()))?;
3753
3754 if table_id.is_empty() {
3755 return Err(Error::invalid_input_source(
3756 "Table ID cannot be empty".into(),
3757 ));
3758 }
3759
3760 let object_id = Self::str_object_id(table_id);
3761 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3762
3763 match table_info {
3764 Some(info) => {
3765 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3766 let mut builder = DatasetBuilder::from_uri(&table_uri);
3767 if let Some(opts) = &self.storage_options {
3768 builder = builder.with_storage_options(opts.clone());
3769 }
3770 if let Some(session) = &self.session {
3771 builder = builder.with_session(session.clone());
3772 }
3773 let mut dataset = builder.load().await.map_err(|e| {
3774 Error::io_source(box_error(std::io::Error::other(format!(
3775 "Failed to open dataset: {}",
3776 e
3777 ))))
3778 })?;
3779
3780 let alterations = super::build_column_alterations(&request.alterations)?;
3782
3783 dataset.alter_columns(&alterations).await.map_err(|e| {
3784 convert_lance_commit_error(&e, "alter_columns", Some(&object_id))
3785 })?;
3786
3787 let version = dataset.version().version as i64;
3788 Ok(AlterTableAlterColumnsResponse::new(version))
3789 }
3790 None => Err(NamespaceError::TableNotFound { message: object_id }.into()),
3791 }
3792 }
3793
3794 async fn alter_table_drop_columns(
3798 &self,
3799 request: AlterTableDropColumnsRequest,
3800 ) -> Result<AlterTableDropColumnsResponse> {
3801 let table_id = request
3802 .id
3803 .as_ref()
3804 .ok_or_else(|| Error::invalid_input_source("Table ID is required".into()))?;
3805
3806 if table_id.is_empty() {
3807 return Err(Error::invalid_input_source(
3808 "Table ID cannot be empty".into(),
3809 ));
3810 }
3811
3812 let object_id = Self::str_object_id(table_id);
3813 let table_info = self.query_manifest_for_table(&object_id).boxed().await?;
3814
3815 match table_info {
3816 Some(info) => {
3817 let table_uri = Self::construct_full_uri(&self.root, &info.location)?;
3818 let mut builder = DatasetBuilder::from_uri(&table_uri);
3819 if let Some(opts) = &self.storage_options {
3820 builder = builder.with_storage_options(opts.clone());
3821 }
3822 if let Some(session) = &self.session {
3823 builder = builder.with_session(session.clone());
3824 }
3825 let mut dataset = builder.load().await.map_err(|e| {
3826 Error::io_source(box_error(std::io::Error::other(format!(
3827 "Failed to open dataset: {}",
3828 e
3829 ))))
3830 })?;
3831
3832 let columns: Vec<&str> = request.columns.iter().map(|s| s.as_str()).collect();
3833 dataset.drop_columns(&columns).await.map_err(|e| {
3834 convert_lance_commit_error(&e, "drop_columns", Some(&object_id))
3835 })?;
3836
3837 let version = dataset.version().version as i64;
3838 Ok(AlterTableDropColumnsResponse::new(version))
3839 }
3840 None => Err(NamespaceError::TableNotFound { message: object_id }.into()),
3841 }
3842 }
3843}
3844
3845#[cfg(test)]
3846mod tests {
3847 use super::{
3848 BASE_OBJECTS_INDEX_NAME, ConflictResolution, CopyOnWriteMutation, DeleteObjectMutation,
3849 LANCE_DATA_DIR, LANCE_INDICES_DIR, MANIFEST_TABLE_NAME, ManifestBatchBuilder,
3850 ManifestEntry, ManifestIndexAccumulator, ManifestNamespace, ManifestOutputRow,
3851 ManifestRowValue, ManifestStreamMutation, OBJECT_ID_INDEX_NAME, OBJECT_TYPE_INDEX_NAME,
3852 ObjectType,
3853 };
3854 use crate::DirectoryNamespaceBuilder;
3855 use arrow::datatypes::DataType;
3856 use bytes::Bytes;
3857 use futures::StreamExt;
3858 use lance::index::DatasetIndexExt;
3859 use lance_core::utils::tempfile::TempStdDir;
3860 use lance_io::object_store::{ObjectStore, ObjectStoreParams, ObjectStoreRegistry};
3861 use lance_namespace::LanceNamespace;
3862 use lance_namespace::models::{
3863 CreateNamespaceRequest, CreateTableRequest, DeclareTableRequest, DescribeTableRequest,
3864 DropTableRequest, ListTablesRequest, TableExistsRequest,
3865 };
3866 use lance_table::feature_flags::FLAG_UNKNOWN;
3867 use lance_table::format::Fragment;
3868 use rstest::rstest;
3869 use std::collections::{HashMap, HashSet};
3870 use std::sync::Arc;
3871
3872 async fn create_manifest_namespace(
3873 root: &str,
3874 inline_optimization_enabled: bool,
3875 ) -> ManifestNamespace {
3876 create_manifest_namespace_with_retries(root, inline_optimization_enabled, None).await
3877 }
3878
3879 async fn create_manifest_namespace_with_retries(
3880 root: &str,
3881 inline_optimization_enabled: bool,
3882 commit_retries: Option<u32>,
3883 ) -> ManifestNamespace {
3884 let (object_store, base_path) = ObjectStore::from_uri_and_params(
3885 Arc::new(ObjectStoreRegistry::default()),
3886 root,
3887 &ObjectStoreParams::default(),
3888 )
3889 .await
3890 .unwrap();
3891 ManifestNamespace::from_directory(
3892 root.to_string(),
3893 None,
3894 None,
3895 object_store,
3896 base_path,
3897 true,
3898 inline_optimization_enabled,
3899 commit_retries,
3900 )
3901 .await
3902 .unwrap()
3903 }
3904
3905 struct CommitConflictAfterRewriteMutation {
3906 root: String,
3907 conflict_object_id: String,
3908 }
3909
3910 impl ManifestStreamMutation for CommitConflictAfterRewriteMutation {
3911 type Output = ();
3912
3913 fn process_existing_row(
3914 &mut self,
3915 row: ManifestRowValue,
3916 output: &mut ManifestBatchBuilder,
3917 index_data: &mut ManifestIndexAccumulator,
3918 ) -> lance_core::Result<()> {
3919 output.append(
3920 index_data,
3921 ManifestOutputRow {
3922 object_id: &row.object_id,
3923 object_type: row.object_type,
3924 location: row.location.as_deref(),
3925 metadata: row.metadata.as_deref(),
3926 base_objects: row.base_objects.as_deref(),
3927 },
3928 )
3929 }
3930
3931 fn append_rows(
3932 &mut self,
3933 output: &mut ManifestBatchBuilder,
3934 index_data: &mut ManifestIndexAccumulator,
3935 ) -> lance_core::Result<()> {
3936 output.append(
3937 index_data,
3938 ManifestOutputRow {
3939 object_id: "attempted_table",
3940 object_type: ObjectType::Table,
3941 location: Some("attempted_table.lance"),
3942 metadata: None,
3943 base_objects: None,
3944 },
3945 )
3946 }
3947
3948 fn finish(&self) -> CopyOnWriteMutation<Self::Output> {
3949 let root = self.root.clone();
3950 let object_id = self.conflict_object_id.clone();
3951 std::thread::spawn(move || {
3952 let runtime = tokio::runtime::Runtime::new().unwrap();
3953 runtime.block_on(async move {
3954 let writer = create_manifest_namespace(&root, false).await;
3955 writer
3956 .insert_into_manifest_with_metadata(
3957 vec![ManifestEntry {
3958 object_id,
3959 object_type: ObjectType::Table,
3960 location: Some("conflicting_table.lance".to_string()),
3961 metadata: None,
3962 }],
3963 None,
3964 )
3965 .await
3966 .unwrap();
3967 });
3968 })
3969 .join()
3970 .unwrap();
3971 CopyOnWriteMutation::updated(())
3972 }
3973 }
3974
3975 struct ConcurrentDeleteBeforeCommitMutation {
3979 inner: DeleteObjectMutation,
3980 root: String,
3981 target: String,
3982 }
3983
3984 impl ManifestStreamMutation for ConcurrentDeleteBeforeCommitMutation {
3985 type Output = ();
3986
3987 fn process_existing_row(
3988 &mut self,
3989 row: ManifestRowValue,
3990 output: &mut ManifestBatchBuilder,
3991 index_data: &mut ManifestIndexAccumulator,
3992 ) -> lance_core::Result<()> {
3993 self.inner.process_existing_row(row, output, index_data)
3994 }
3995
3996 fn append_rows(
3997 &mut self,
3998 output: &mut ManifestBatchBuilder,
3999 index_data: &mut ManifestIndexAccumulator,
4000 ) -> lance_core::Result<()> {
4001 self.inner.append_rows(output, index_data)
4002 }
4003
4004 fn finish(&self) -> CopyOnWriteMutation<Self::Output> {
4005 let root = self.root.clone();
4006 let target = self.target.clone();
4007 std::thread::spawn(move || {
4008 let runtime = tokio::runtime::Runtime::new().unwrap();
4009 runtime.block_on(async move {
4010 let writer = create_manifest_namespace(&root, false).await;
4011 writer.delete_from_manifest(&target).await.unwrap();
4012 });
4013 })
4014 .join()
4015 .unwrap();
4016 self.inner.finish()
4017 }
4018
4019 fn conflict_resolution(&self) -> ConflictResolution<Self::Output> {
4020 ConflictResolution::SucceedIfAbsent {
4021 object_id: self.target.clone(),
4022 output: (),
4023 }
4024 }
4025 }
4026
4027 async fn manifest_base_objects(
4028 manifest_ns: &ManifestNamespace,
4029 ) -> HashMap<String, Option<Vec<String>>> {
4030 let mut scanner = manifest_ns.manifest_scanner().await.unwrap();
4031 scanner.project(&["object_id", "base_objects"]).unwrap();
4032 let batches = ManifestNamespace::execute_scanner(scanner).await.unwrap();
4033 let mut rows = HashMap::new();
4034 for batch in batches {
4035 let object_ids = ManifestNamespace::get_string_column(&batch, "object_id").unwrap();
4036 let base_objects = ManifestNamespace::base_objects_column_values(&batch).unwrap();
4037 for (row, value) in base_objects.into_iter().enumerate() {
4038 rows.insert(object_ids.value(row).to_string(), value);
4039 }
4040 }
4041 rows
4042 }
4043
4044 async fn manifest_data_paths(manifest_ns: &ManifestNamespace) -> HashSet<String> {
4045 let data_dir = manifest_ns
4046 .base_path
4047 .clone()
4048 .join(MANIFEST_TABLE_NAME)
4049 .join(LANCE_DATA_DIR);
4050 let mut stream = manifest_ns.object_store.read_dir_all(&data_dir, None);
4051 let mut paths = HashSet::new();
4052 while let Some(meta) = stream.next().await.transpose().unwrap() {
4053 paths.insert(meta.location.to_string());
4054 }
4055 paths
4056 }
4057
4058 async fn manifest_index_paths(manifest_ns: &ManifestNamespace) -> HashSet<String> {
4059 let index_dir = manifest_ns
4060 .base_path
4061 .clone()
4062 .join(MANIFEST_TABLE_NAME)
4063 .join(LANCE_INDICES_DIR);
4064 let mut stream = manifest_ns.object_store.read_dir_all(&index_dir, None);
4065 let mut paths = HashSet::new();
4066 while let Some(meta) = stream.next().await.transpose().unwrap() {
4067 paths.insert(meta.location.to_string());
4068 }
4069 paths
4070 }
4071
4072 fn create_test_ipc_data() -> Vec<u8> {
4073 use arrow::array::{Int32Array, StringArray};
4074 use arrow::datatypes::{DataType, Field, Schema};
4075 use arrow::ipc::writer::StreamWriter;
4076 use arrow::record_batch::RecordBatch;
4077 use std::sync::Arc;
4078
4079 let schema = Arc::new(Schema::new(vec![
4080 Field::new("id", DataType::Int32, false),
4081 Field::new("name", DataType::Utf8, false),
4082 ]));
4083
4084 let batch = RecordBatch::try_new(
4085 schema.clone(),
4086 vec![
4087 Arc::new(Int32Array::from(vec![1, 2, 3])),
4088 Arc::new(StringArray::from(vec!["a", "b", "c"])),
4089 ],
4090 )
4091 .unwrap();
4092
4093 let mut buffer = Vec::new();
4094 {
4095 let mut writer = StreamWriter::try_new(&mut buffer, &schema).unwrap();
4096 writer.write(&batch).unwrap();
4097 writer.finish().unwrap();
4098 }
4099 buffer
4100 }
4101
4102 async fn set_manifest_table_metadata(temp_path: &str, key: &str, value: &str) {
4105 use lance::dataset::builder::DatasetBuilder;
4106 let mut ds = DatasetBuilder::from_uri(format!("{}/{}", temp_path, MANIFEST_TABLE_NAME))
4107 .load()
4108 .await
4109 .unwrap();
4110 ds.update_metadata([(key, value)]).await.unwrap();
4111 }
4112
4113 async fn create_namespace_with_one_table(temp_path: &str) {
4114 let ns = DirectoryNamespaceBuilder::new(temp_path)
4115 .build()
4116 .await
4117 .unwrap();
4118 let mut create_request = CreateTableRequest::new();
4119 create_request.id = Some(vec!["t1".to_string()]);
4120 ns.create_table(create_request, Bytes::from(create_test_ipc_data()))
4121 .await
4122 .unwrap();
4123 }
4124
4125 #[tokio::test]
4128 async fn test_manifest_has_no_feature_flags_by_default() {
4129 use lance::dataset::builder::DatasetBuilder;
4130 let temp_dir = TempStdDir::default();
4131 let temp_path = temp_dir.to_str().unwrap();
4132 create_namespace_with_one_table(temp_path).await;
4133
4134 let ds = DatasetBuilder::from_uri(format!("{}/{}", temp_path, MANIFEST_TABLE_NAME))
4135 .load()
4136 .await
4137 .unwrap();
4138 assert!(
4139 !ds.metadata()
4140 .contains_key(crate::dir::manifest_feature_flags::READER_FEATURE_FLAGS_KEY)
4141 );
4142 assert!(
4143 !ds.metadata()
4144 .contains_key(crate::dir::manifest_feature_flags::WRITER_FEATURE_FLAGS_KEY)
4145 );
4146 }
4147
4148 #[tokio::test]
4151 async fn test_unknown_reader_flag_blocks_access() {
4152 let temp_dir = TempStdDir::default();
4153 let temp_path = temp_dir.to_str().unwrap();
4154 create_namespace_with_one_table(temp_path).await;
4155 set_manifest_table_metadata(
4156 temp_path,
4157 crate::dir::manifest_feature_flags::READER_FEATURE_FLAGS_KEY,
4158 "1",
4159 )
4160 .await;
4161
4162 let err = DirectoryNamespaceBuilder::new(temp_path)
4163 .build()
4164 .await
4165 .expect_err("opening a manifest with an unknown reader flag should fail");
4166 assert!(
4167 err.to_string().to_lowercase().contains("upgrade"),
4168 "expected an upgrade error, got: {err}"
4169 );
4170 }
4171
4172 #[tokio::test]
4174 async fn test_unknown_writer_flag_blocks_writes_but_allows_reads() {
4175 let temp_dir = TempStdDir::default();
4176 let temp_path = temp_dir.to_str().unwrap();
4177 create_namespace_with_one_table(temp_path).await;
4178 set_manifest_table_metadata(
4179 temp_path,
4180 crate::dir::manifest_feature_flags::WRITER_FEATURE_FLAGS_KEY,
4181 "1",
4182 )
4183 .await;
4184
4185 let ns = DirectoryNamespaceBuilder::new(temp_path)
4186 .build()
4187 .await
4188 .expect("reads should still be allowed with only a writer flag set");
4189 let mut list_request = ListTablesRequest::new();
4190 list_request.id = Some(vec![]);
4191 assert_eq!(ns.list_tables(list_request).await.unwrap().tables.len(), 1);
4192
4193 let entries_before = dir_entry_names(temp_path);
4195 let mut create_request = CreateTableRequest::new();
4196 create_request.id = Some(vec!["t2".to_string()]);
4197 let err = ns
4198 .create_table(create_request, Bytes::from(create_test_ipc_data()))
4199 .await
4200 .expect_err("writing through an unknown writer flag should fail");
4201 assert!(
4202 err.to_string().to_lowercase().contains("upgrade"),
4203 "expected an upgrade error, got: {err}"
4204 );
4205 assert_eq!(
4206 entries_before,
4207 dir_entry_names(temp_path),
4208 "a refused create_table must not create an orphaned table directory"
4209 );
4210
4211 let mut create_ns = CreateNamespaceRequest::new();
4216 create_ns.id = Some(vec!["ns1".to_string()]);
4217 let err = ns
4218 .create_namespace(create_ns)
4219 .await
4220 .expect_err("create_namespace through an unknown writer flag should fail");
4221 assert!(
4222 err.to_string().to_lowercase().contains("upgrade"),
4223 "expected an upgrade error, got: {err}"
4224 );
4225
4226 let mut drop_request = DropTableRequest::new();
4227 drop_request.id = Some(vec!["t1".to_string()]);
4228 let err = ns
4229 .drop_table(drop_request)
4230 .await
4231 .expect_err("drop_table through an unknown writer flag should fail");
4232 assert!(
4233 err.to_string().to_lowercase().contains("upgrade"),
4234 "expected an upgrade error, got: {err}"
4235 );
4236 }
4237
4238 fn dir_entry_names(path: &str) -> std::collections::BTreeSet<String> {
4239 std::fs::read_dir(path)
4240 .unwrap()
4241 .map(|e| e.unwrap().file_name().to_string_lossy().into_owned())
4242 .collect()
4243 }
4244
4245 #[tokio::test]
4246 async fn test_manifest_rewrite_preserves_utf8_metadata_and_base_objects() {
4247 let temp_dir = TempStdDir::default();
4248 let temp_path = temp_dir.to_str().unwrap();
4249 let manifest_ns = create_manifest_namespace(temp_path, true).await;
4250
4251 manifest_ns
4252 .insert_into_manifest_with_metadata(
4253 vec![ManifestEntry {
4254 object_id: "view".to_string(),
4255 object_type: ObjectType::Table,
4256 location: Some("view.lance".to_string()),
4257 metadata: Some(r#"{"kind":"view"}"#.to_string()),
4258 }],
4259 Some(vec!["base_a".to_string(), "base_b".to_string()]),
4260 )
4261 .await
4262 .unwrap();
4263 manifest_ns
4264 .insert_into_manifest_with_metadata(
4265 vec![ManifestEntry {
4266 object_id: "other".to_string(),
4267 object_type: ObjectType::Namespace,
4268 location: None,
4269 metadata: Some(r#"{"kind":"namespace"}"#.to_string()),
4270 }],
4271 None,
4272 )
4273 .await
4274 .unwrap();
4275
4276 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4277 let metadata_field = dataset_guard.schema().field("metadata").unwrap();
4278 assert_eq!(metadata_field.data_type(), DataType::Utf8);
4279 drop(dataset_guard);
4280
4281 let base_objects = manifest_base_objects(&manifest_ns).await;
4282 assert_eq!(
4283 base_objects.get("view").cloned().unwrap(),
4284 Some(vec!["base_a".to_string(), "base_b".to_string()])
4285 );
4286 assert_eq!(base_objects.get("other").cloned().unwrap(), None);
4287 }
4288
4289 #[tokio::test]
4290 async fn test_manifest_rewrite_replacement_indices_are_versioned() {
4291 let temp_dir = TempStdDir::default();
4292 let temp_path = temp_dir.to_str().unwrap();
4293 let manifest_ns = create_manifest_namespace(temp_path, true).await;
4294
4295 manifest_ns
4296 .insert_into_manifest_with_metadata(
4297 vec![ManifestEntry {
4298 object_id: "table".to_string(),
4299 object_type: ObjectType::Table,
4300 location: Some("table.lance".to_string()),
4301 metadata: None,
4302 }],
4303 Some(vec!["base".to_string()]),
4304 )
4305 .await
4306 .unwrap();
4307
4308 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4309 let dataset_version = dataset_guard.version().version;
4310 let indices = dataset_guard.load_indices().await.unwrap();
4311 let names = indices
4312 .iter()
4313 .map(|index| index.name.as_str())
4314 .collect::<HashSet<_>>();
4315 assert!(names.contains(OBJECT_ID_INDEX_NAME));
4316 assert!(names.contains(OBJECT_TYPE_INDEX_NAME));
4317 assert!(names.contains(BASE_OBJECTS_INDEX_NAME));
4318 for index in indices.iter() {
4319 assert_eq!(index.dataset_version, dataset_version);
4320 assert!(!index.fragment_bitmap.as_ref().unwrap().is_empty());
4321 }
4322 }
4323
4324 #[tokio::test]
4325 async fn test_manifest_rewrite_empty_manifest_keeps_replacement_indices_valid() {
4326 let temp_dir = TempStdDir::default();
4327 let temp_path = temp_dir.to_str().unwrap();
4328 let manifest_ns = create_manifest_namespace(temp_path, true).await;
4329
4330 manifest_ns
4331 .insert_into_manifest_with_metadata(
4332 vec![ManifestEntry {
4333 object_id: "table".to_string(),
4334 object_type: ObjectType::Table,
4335 location: Some("table.lance".to_string()),
4336 metadata: None,
4337 }],
4338 None,
4339 )
4340 .await
4341 .unwrap();
4342 manifest_ns.delete_from_manifest("table").await.unwrap();
4343
4344 assert!(!manifest_ns.manifest_contains_object("table").await.unwrap());
4345 let mut scanner = manifest_ns.manifest_scanner().await.unwrap();
4346 scanner.project(&["object_id"]).unwrap();
4347 let rows = ManifestNamespace::execute_scanner(scanner)
4348 .await
4349 .unwrap()
4350 .into_iter()
4351 .map(|batch| batch.num_rows())
4352 .sum::<usize>();
4353 assert_eq!(rows, 0);
4354
4355 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4356 let dataset_version = dataset_guard.version().version;
4357 let indices = dataset_guard.load_indices().await.unwrap();
4358 let names = indices
4359 .iter()
4360 .map(|index| index.name.as_str())
4361 .collect::<HashSet<_>>();
4362 assert!(names.contains(OBJECT_ID_INDEX_NAME));
4363 assert!(names.contains(OBJECT_TYPE_INDEX_NAME));
4364 assert!(names.contains(BASE_OBJECTS_INDEX_NAME));
4365 for index in indices.iter() {
4366 assert_eq!(index.dataset_version, dataset_version);
4367 }
4368 }
4369
4370 #[tokio::test]
4371 async fn test_manifest_rewrite_fragment_bitmap_uses_overwrite_fragment_ids() {
4372 let temp_dir = TempStdDir::default();
4373 let temp_path = temp_dir.to_str().unwrap();
4374 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4375 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4376 let fragments = vec![Fragment::new(0), Fragment::new(0), Fragment::new(7)];
4377
4378 let manifest = ManifestNamespace::manifest_from_overwrite_transaction(
4379 dataset_guard.manifest(),
4380 dataset_guard.manifest().schema.clone(),
4381 &fragments,
4382 );
4383
4384 let fragment_ids = manifest
4385 .fragments
4386 .iter()
4387 .map(|fragment| fragment.id)
4388 .collect::<Vec<_>>();
4389 assert_eq!(fragment_ids, vec![0, 1, 7]);
4390 assert_eq!(
4391 ManifestNamespace::manifest_fragment_bitmap(&manifest)
4392 .unwrap()
4393 .into_iter()
4394 .collect::<Vec<_>>(),
4395 vec![0, 1, 7]
4396 );
4397 }
4398
4399 #[tokio::test]
4400 async fn test_manifest_writes_reject_unknown_writer_flag_before_staging() {
4401 let temp_dir = TempStdDir::default();
4402 let temp_path = temp_dir.to_str().unwrap();
4403 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4404 let data_paths_before = manifest_data_paths(&manifest_ns).await;
4405 let original_version = {
4406 let mut dataset = manifest_ns.manifest_dataset.get_mut().await.unwrap();
4407 let mut manifest = dataset.manifest().clone();
4408 manifest.writer_feature_flags |= FLAG_UNKNOWN << 1;
4409 let version = manifest.version;
4410 dataset.manifest = Arc::new(manifest);
4411 version
4412 };
4413
4414 let entries_before = dir_entry_names(temp_path);
4415 let mut declare_request = DeclareTableRequest::new();
4416 declare_request.id = Some(vec!["declared_table".to_string()]);
4417 let error = manifest_ns
4418 .declare_table(declare_request)
4419 .await
4420 .unwrap_err();
4421 assert!(
4422 error.to_string().to_lowercase().contains("upgrade"),
4423 "expected an upgrade error, got: {error}"
4424 );
4425 assert_eq!(dir_entry_names(temp_path), entries_before);
4426
4427 let mut create_request = CreateTableRequest::new();
4428 create_request.id = Some(vec!["new_table".to_string()]);
4429 let error = manifest_ns
4430 .create_table(create_request, Bytes::from(create_test_ipc_data()))
4431 .await
4432 .unwrap_err();
4433 assert!(
4434 error.to_string().to_lowercase().contains("upgrade"),
4435 "expected an upgrade error, got: {error}"
4436 );
4437 assert_eq!(dir_entry_names(temp_path), entries_before);
4438
4439 let error = manifest_ns
4440 .insert_into_manifest_with_metadata(
4441 vec![ManifestEntry {
4442 object_id: "table".to_string(),
4443 object_type: ObjectType::Table,
4444 location: Some("table.lance".to_string()),
4445 metadata: None,
4446 }],
4447 None,
4448 )
4449 .await
4450 .unwrap_err();
4451
4452 assert!(
4453 error.to_string().to_lowercase().contains("upgrade"),
4454 "expected an upgrade error, got: {error}"
4455 );
4456 assert_eq!(
4457 manifest_ns
4458 .manifest_dataset
4459 .get()
4460 .await
4461 .unwrap()
4462 .version()
4463 .version,
4464 original_version
4465 );
4466 assert_eq!(manifest_data_paths(&manifest_ns).await, data_paths_before);
4467 }
4468
4469 #[tokio::test]
4470 async fn test_manifest_noop_delete_uses_latest_snapshot() {
4471 let temp_dir = TempStdDir::default();
4472 let temp_path = temp_dir.to_str().unwrap();
4473 let stale_ns = create_manifest_namespace(temp_path, false).await;
4474 let writer_ns = create_manifest_namespace(temp_path, false).await;
4475
4476 writer_ns
4477 .insert_into_manifest_with_metadata(
4478 vec![ManifestEntry {
4479 object_id: "late_table".to_string(),
4480 object_type: ObjectType::Table,
4481 location: Some("late_table.lance".to_string()),
4482 metadata: None,
4483 }],
4484 None,
4485 )
4486 .await
4487 .unwrap();
4488
4489 stale_ns.delete_from_manifest("late_table").await.unwrap();
4490
4491 let check_ns = create_manifest_namespace(temp_path, false).await;
4492 assert!(
4493 !check_ns
4494 .manifest_contains_object("late_table")
4495 .await
4496 .unwrap()
4497 );
4498 }
4499
4500 #[tokio::test]
4501 async fn test_manifest_noop_delete_cleans_uncommitted_data_file() {
4502 let temp_dir = TempStdDir::default();
4503 let temp_path = temp_dir.to_str().unwrap();
4504 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4505
4506 manifest_ns
4507 .insert_into_manifest_with_metadata(
4508 vec![ManifestEntry {
4509 object_id: "table".to_string(),
4510 object_type: ObjectType::Table,
4511 location: Some("table.lance".to_string()),
4512 metadata: None,
4513 }],
4514 None,
4515 )
4516 .await
4517 .unwrap();
4518
4519 let before = manifest_data_paths(&manifest_ns).await;
4520 assert!(!before.is_empty());
4521
4522 manifest_ns
4523 .delete_from_manifest("missing_table")
4524 .await
4525 .unwrap();
4526
4527 let after = manifest_data_paths(&manifest_ns).await;
4528 assert_eq!(after, before);
4529 }
4530
4531 #[tokio::test]
4532 async fn test_manifest_final_commit_failure_cleans_uncommitted_rewrite_files() {
4533 let temp_dir = TempStdDir::default();
4534 let temp_path = temp_dir.to_str().unwrap();
4535 let manifest_ns = create_manifest_namespace_with_retries(temp_path, true, Some(0)).await;
4536
4537 manifest_ns
4538 .insert_into_manifest_with_metadata(
4539 vec![ManifestEntry {
4540 object_id: "table".to_string(),
4541 object_type: ObjectType::Table,
4542 location: Some("table.lance".to_string()),
4543 metadata: None,
4544 }],
4545 None,
4546 )
4547 .await
4548 .unwrap();
4549
4550 let before_data_paths = manifest_data_paths(&manifest_ns).await;
4551 let before_index_paths = manifest_index_paths(&manifest_ns).await;
4552
4553 let result = manifest_ns
4554 .rewrite_manifest("Failed to test manifest cleanup", || {
4555 CommitConflictAfterRewriteMutation {
4556 root: temp_path.to_string(),
4557 conflict_object_id: "conflicting_table".to_string(),
4558 }
4559 })
4560 .await;
4561 assert!(result.is_err());
4562
4563 let after_data_paths = manifest_data_paths(&manifest_ns).await;
4564 assert!(before_data_paths.is_subset(&after_data_paths));
4565 assert_eq!(after_data_paths.len(), before_data_paths.len() + 1);
4566 assert_eq!(manifest_index_paths(&manifest_ns).await, before_index_paths);
4567 assert!(
4568 manifest_ns
4569 .manifest_contains_object("conflicting_table")
4570 .await
4571 .unwrap()
4572 );
4573 assert!(
4574 !manifest_ns
4575 .manifest_contains_object("attempted_table")
4576 .await
4577 .unwrap()
4578 );
4579 }
4580
4581 #[tokio::test]
4582 async fn test_manifest_commit_visible_on_memory_store() {
4583 let manifest_ns = create_manifest_namespace("memory://test_commit_visible", false).await;
4588 manifest_ns
4589 .insert_into_manifest_with_metadata(
4590 vec![ManifestEntry {
4591 object_id: "table".to_string(),
4592 object_type: ObjectType::Table,
4593 location: Some("table.lance".to_string()),
4594 metadata: None,
4595 }],
4596 None,
4597 )
4598 .await
4599 .unwrap();
4600 assert!(manifest_ns.manifest_contains_object("table").await.unwrap());
4601 manifest_ns
4603 .insert_into_manifest_with_metadata(
4604 vec![ManifestEntry {
4605 object_id: "table2".to_string(),
4606 object_type: ObjectType::Table,
4607 location: Some("table2.lance".to_string()),
4608 metadata: None,
4609 }],
4610 None,
4611 )
4612 .await
4613 .unwrap();
4614 assert!(
4615 manifest_ns
4616 .manifest_contains_object("table2")
4617 .await
4618 .unwrap()
4619 );
4620 }
4621
4622 #[tokio::test]
4623 async fn test_manifest_commit_uses_inline_transaction() {
4624 let temp_dir = TempStdDir::default();
4625 let temp_path = temp_dir.to_str().unwrap();
4626 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4627
4628 manifest_ns
4629 .insert_into_manifest_with_metadata(
4630 vec![ManifestEntry {
4631 object_id: "table".to_string(),
4632 object_type: ObjectType::Table,
4633 location: Some("table.lance".to_string()),
4634 metadata: None,
4635 }],
4636 None,
4637 )
4638 .await
4639 .unwrap();
4640
4641 let dataset_guard = manifest_ns.manifest_dataset.get().await.unwrap();
4642 let manifest = dataset_guard.manifest();
4643 assert!(manifest.transaction_section.is_some());
4646 assert!(manifest.transaction_file.is_none());
4647 }
4648
4649 #[tokio::test]
4650 async fn test_manifest_commit_landed_attributes_data_file() {
4651 let temp_dir = TempStdDir::default();
4652 let temp_path = temp_dir.to_str().unwrap();
4653 let manifest_ns = create_manifest_namespace(temp_path, false).await;
4654
4655 manifest_ns
4656 .insert_into_manifest_with_metadata(
4657 vec![ManifestEntry {
4658 object_id: "table".to_string(),
4659 object_type: ObjectType::Table,
4660 location: Some("table.lance".to_string()),
4661 metadata: None,
4662 }],
4663 None,
4664 )
4665 .await
4666 .unwrap();
4667
4668 let dataset = Arc::new(manifest_ns.manifest_dataset.get().await.unwrap().clone());
4669 let version = dataset.manifest().version;
4670 let our_files = dataset
4671 .manifest()
4672 .fragments
4673 .iter()
4674 .flat_map(|fragment| fragment.files.iter())
4675 .map(|file| file.path.clone())
4676 .collect::<HashSet<_>>();
4677 assert!(!our_files.is_empty());
4678
4679 assert!(
4682 manifest_ns
4683 .manifest_commit_landed(&dataset, version, &our_files)
4684 .await
4685 );
4686 let other = HashSet::from(["missing.lance".to_string()]);
4688 assert!(
4689 !manifest_ns
4690 .manifest_commit_landed(&dataset, version, &other)
4691 .await
4692 );
4693 assert!(
4695 !manifest_ns
4696 .manifest_commit_landed(&dataset, version + 100, &our_files)
4697 .await
4698 );
4699 }
4700
4701 #[tokio::test]
4702 async fn test_manifest_delete_conflict_with_concurrent_delete_succeeds() {
4703 let temp_dir = TempStdDir::default();
4704 let temp_path = temp_dir.to_str().unwrap();
4705 let manifest_ns = create_manifest_namespace_with_retries(temp_path, false, Some(0)).await;
4706
4707 manifest_ns
4708 .insert_into_manifest_with_metadata(
4709 vec![ManifestEntry {
4710 object_id: "table".to_string(),
4711 object_type: ObjectType::Table,
4712 location: Some("table.lance".to_string()),
4713 metadata: None,
4714 }],
4715 None,
4716 )
4717 .await
4718 .unwrap();
4719 assert!(manifest_ns.manifest_contains_object("table").await.unwrap());
4720
4721 let result = manifest_ns
4725 .rewrite_manifest("Failed to delete from manifest", || {
4726 ConcurrentDeleteBeforeCommitMutation {
4727 inner: DeleteObjectMutation {
4728 object_id: "table".to_string(),
4729 deleted: false,
4730 },
4731 root: temp_path.to_string(),
4732 target: "table".to_string(),
4733 }
4734 })
4735 .await;
4736
4737 assert!(result.is_ok(), "delete should succeed: {result:?}");
4738 assert!(!manifest_ns.manifest_contains_object("table").await.unwrap());
4739 }
4740
4741 #[rstest]
4742 #[case::with_optimization(true)]
4743 #[case::without_optimization(false)]
4744 #[tokio::test]
4745 async fn test_manifest_namespace_basic_create_and_list(#[case] inline_optimization: bool) {
4746 let temp_dir = TempStdDir::default();
4747 let temp_path = temp_dir.to_str().unwrap();
4748
4749 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4751 .inline_optimization_enabled(inline_optimization)
4752 .build()
4753 .await
4754 .unwrap();
4755
4756 let mut request = ListTablesRequest::new();
4758 request.id = Some(vec![]);
4759 let response = dir_namespace.list_tables(request).await.unwrap();
4760 assert_eq!(response.tables.len(), 0);
4761
4762 let buffer = create_test_ipc_data();
4764 let mut create_request = CreateTableRequest::new();
4765 create_request.id = Some(vec!["test_table".to_string()]);
4766
4767 let _response = dir_namespace
4768 .create_table(create_request, Bytes::from(buffer))
4769 .await
4770 .unwrap();
4771
4772 let mut request = ListTablesRequest::new();
4774 request.id = Some(vec![]);
4775 let response = dir_namespace.list_tables(request).await.unwrap();
4776 assert_eq!(response.tables.len(), 1);
4777 assert_eq!(response.tables[0], "test_table");
4778 }
4779
4780 #[rstest]
4781 #[case::with_optimization(true)]
4782 #[case::without_optimization(false)]
4783 #[tokio::test]
4784 async fn test_manifest_namespace_table_exists(#[case] inline_optimization: bool) {
4785 let temp_dir = TempStdDir::default();
4786 let temp_path = temp_dir.to_str().unwrap();
4787
4788 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4789 .inline_optimization_enabled(inline_optimization)
4790 .build()
4791 .await
4792 .unwrap();
4793
4794 let mut request = TableExistsRequest::new();
4796 request.id = Some(vec!["nonexistent".to_string()]);
4797 let result = dir_namespace.table_exists(request).await;
4798 assert!(result.is_err());
4799
4800 let buffer = create_test_ipc_data();
4802 let mut create_request = CreateTableRequest::new();
4803 create_request.id = Some(vec!["test_table".to_string()]);
4804 dir_namespace
4805 .create_table(create_request, Bytes::from(buffer))
4806 .await
4807 .unwrap();
4808
4809 let mut request = TableExistsRequest::new();
4811 request.id = Some(vec!["test_table".to_string()]);
4812 let result = dir_namespace.table_exists(request).await;
4813 assert!(result.is_ok());
4814 }
4815
4816 #[rstest]
4817 #[case::with_optimization(true)]
4818 #[case::without_optimization(false)]
4819 #[tokio::test]
4820 async fn test_manifest_namespace_describe_table(#[case] inline_optimization: bool) {
4821 let temp_dir = TempStdDir::default();
4822 let temp_path = temp_dir.to_str().unwrap();
4823
4824 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4825 .inline_optimization_enabled(inline_optimization)
4826 .build()
4827 .await
4828 .unwrap();
4829
4830 let mut request = DescribeTableRequest::new();
4832 request.id = Some(vec!["nonexistent".to_string()]);
4833 let result = dir_namespace.describe_table(request).await;
4834 assert!(result.is_err());
4835
4836 let buffer = create_test_ipc_data();
4838 let mut create_request = CreateTableRequest::new();
4839 create_request.id = Some(vec!["test_table".to_string()]);
4840 dir_namespace
4841 .create_table(create_request, Bytes::from(buffer))
4842 .await
4843 .unwrap();
4844
4845 let mut request = DescribeTableRequest::new();
4847 request.id = Some(vec!["test_table".to_string()]);
4848 let response = dir_namespace.describe_table(request).await.unwrap();
4849 assert!(response.location.is_some());
4850 assert!(response.location.unwrap().contains("test_table"));
4851 }
4852
4853 #[rstest]
4854 #[case::with_optimization(true)]
4855 #[case::without_optimization(false)]
4856 #[tokio::test]
4857 async fn test_manifest_namespace_drop_table(#[case] inline_optimization: bool) {
4858 let temp_dir = TempStdDir::default();
4859 let temp_path = temp_dir.to_str().unwrap();
4860
4861 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4862 .inline_optimization_enabled(inline_optimization)
4863 .build()
4864 .await
4865 .unwrap();
4866
4867 let buffer = create_test_ipc_data();
4869 let mut create_request = CreateTableRequest::new();
4870 create_request.id = Some(vec!["test_table".to_string()]);
4871 dir_namespace
4872 .create_table(create_request, Bytes::from(buffer))
4873 .await
4874 .unwrap();
4875
4876 let mut request = ListTablesRequest::new();
4878 request.id = Some(vec![]);
4879 let response = dir_namespace.list_tables(request).await.unwrap();
4880 assert_eq!(response.tables.len(), 1);
4881
4882 let mut drop_request = DropTableRequest::new();
4884 drop_request.id = Some(vec!["test_table".to_string()]);
4885 let _response = dir_namespace.drop_table(drop_request).await.unwrap();
4886
4887 let mut request = ListTablesRequest::new();
4889 request.id = Some(vec![]);
4890 let response = dir_namespace.list_tables(request).await.unwrap();
4891 assert_eq!(response.tables.len(), 0);
4892 }
4893
4894 #[tokio::test]
4895 async fn test_list_tables_pagination_limit_zero() {
4896 let temp_dir = TempStdDir::default();
4897 let temp_path = temp_dir.to_str().unwrap();
4898
4899 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4900 .build()
4901 .await
4902 .unwrap();
4903
4904 let buffer = create_test_ipc_data();
4905 let mut create_request = CreateTableRequest::new();
4906 create_request.id = Some(vec!["alpha".to_string()]);
4907 dir_namespace
4908 .create_table(create_request, Bytes::from(buffer))
4909 .await
4910 .unwrap();
4911
4912 let response = dir_namespace
4913 .list_tables(ListTablesRequest {
4914 id: Some(vec![]),
4915 limit: Some(0),
4916 ..Default::default()
4917 })
4918 .await
4919 .unwrap();
4920
4921 assert!(response.tables.is_empty());
4922 assert!(response.page_token.is_none());
4923 }
4924
4925 #[rstest]
4926 #[case::with_optimization(true)]
4927 #[case::without_optimization(false)]
4928 #[tokio::test]
4929 async fn test_manifest_namespace_multiple_tables(#[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)
4934 .inline_optimization_enabled(inline_optimization)
4935 .build()
4936 .await
4937 .unwrap();
4938
4939 let buffer = create_test_ipc_data();
4941 for i in 1..=3 {
4942 let mut create_request = CreateTableRequest::new();
4943 create_request.id = Some(vec![format!("table{}", i)]);
4944 dir_namespace
4945 .create_table(create_request, Bytes::from(buffer.clone()))
4946 .await
4947 .unwrap();
4948 }
4949
4950 let mut request = ListTablesRequest::new();
4952 request.id = Some(vec![]);
4953 let response = dir_namespace.list_tables(request).await.unwrap();
4954 assert_eq!(response.tables.len(), 3);
4955 assert!(response.tables.contains(&"table1".to_string()));
4956 assert!(response.tables.contains(&"table2".to_string()));
4957 assert!(response.tables.contains(&"table3".to_string()));
4958 }
4959
4960 #[rstest]
4961 #[case::with_optimization(true)]
4962 #[case::without_optimization(false)]
4963 #[tokio::test]
4964 async fn test_directory_only_mode(#[case] inline_optimization: bool) {
4965 let temp_dir = TempStdDir::default();
4966 let temp_path = temp_dir.to_str().unwrap();
4967
4968 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
4970 .manifest_enabled(false)
4971 .inline_optimization_enabled(inline_optimization)
4972 .build()
4973 .await
4974 .unwrap();
4975
4976 let mut request = ListTablesRequest::new();
4978 request.id = Some(vec![]);
4979 let response = dir_namespace.list_tables(request).await.unwrap();
4980 assert_eq!(response.tables.len(), 0);
4981
4982 let buffer = create_test_ipc_data();
4984 let mut create_request = CreateTableRequest::new();
4985 create_request.id = Some(vec!["test_table".to_string()]);
4986
4987 let _response = dir_namespace
4989 .create_table(create_request, Bytes::from(buffer))
4990 .await
4991 .unwrap();
4992
4993 let mut request = ListTablesRequest::new();
4995 request.id = Some(vec![]);
4996 let response = dir_namespace.list_tables(request).await.unwrap();
4997 assert_eq!(response.tables.len(), 1);
4998 assert_eq!(response.tables[0], "test_table");
4999 }
5000
5001 #[rstest]
5002 #[case::with_optimization(true)]
5003 #[case::without_optimization(false)]
5004 #[tokio::test]
5005 async fn test_dual_mode_merge(#[case] inline_optimization: bool) {
5006 let temp_dir = TempStdDir::default();
5007 let temp_path = temp_dir.to_str().unwrap();
5008
5009 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5011 .manifest_enabled(true)
5012 .dir_listing_enabled(true)
5013 .inline_optimization_enabled(inline_optimization)
5014 .build()
5015 .await
5016 .unwrap();
5017
5018 let buffer = create_test_ipc_data();
5020 let mut create_request = CreateTableRequest::new();
5021 create_request.id = Some(vec!["table1".to_string()]);
5022 dir_namespace
5023 .create_table(create_request, Bytes::from(buffer))
5024 .await
5025 .unwrap();
5026
5027 let mut request = ListTablesRequest::new();
5029 request.id = Some(vec![]);
5030 let response = dir_namespace.list_tables(request).await.unwrap();
5031 assert_eq!(response.tables.len(), 1);
5032 assert_eq!(response.tables[0], "table1");
5033 }
5034
5035 #[rstest]
5036 #[case::with_optimization(true)]
5037 #[case::without_optimization(false)]
5038 #[tokio::test]
5039 async fn test_manifest_only_mode(#[case] inline_optimization: bool) {
5040 let temp_dir = TempStdDir::default();
5041 let temp_path = temp_dir.to_str().unwrap();
5042
5043 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5045 .manifest_enabled(true)
5046 .dir_listing_enabled(false)
5047 .inline_optimization_enabled(inline_optimization)
5048 .build()
5049 .await
5050 .unwrap();
5051
5052 let buffer = create_test_ipc_data();
5054 let mut create_request = CreateTableRequest::new();
5055 create_request.id = Some(vec!["test_table".to_string()]);
5056 dir_namespace
5057 .create_table(create_request, Bytes::from(buffer))
5058 .await
5059 .unwrap();
5060
5061 let mut request = ListTablesRequest::new();
5063 request.id = Some(vec![]);
5064 let response = dir_namespace.list_tables(request).await.unwrap();
5065 assert_eq!(response.tables.len(), 1);
5066 assert_eq!(response.tables[0], "test_table");
5067 }
5068
5069 #[rstest]
5070 #[case::with_optimization(true)]
5071 #[case::without_optimization(false)]
5072 #[tokio::test]
5073 async fn test_drop_nonexistent_table(#[case] inline_optimization: bool) {
5074 let temp_dir = TempStdDir::default();
5075 let temp_path = temp_dir.to_str().unwrap();
5076
5077 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5078 .inline_optimization_enabled(inline_optimization)
5079 .build()
5080 .await
5081 .unwrap();
5082
5083 let mut drop_request = DropTableRequest::new();
5085 drop_request.id = Some(vec!["nonexistent".to_string()]);
5086 let result = dir_namespace.drop_table(drop_request).await;
5087 assert!(result.is_err());
5088 }
5089
5090 #[rstest]
5091 #[case::with_optimization(true)]
5092 #[case::without_optimization(false)]
5093 #[tokio::test]
5094 async fn test_create_duplicate_table_fails(#[case] inline_optimization: bool) {
5095 let temp_dir = TempStdDir::default();
5096 let temp_path = temp_dir.to_str().unwrap();
5097
5098 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5099 .inline_optimization_enabled(inline_optimization)
5100 .build()
5101 .await
5102 .unwrap();
5103
5104 let buffer = create_test_ipc_data();
5106 let mut create_request = CreateTableRequest::new();
5107 create_request.id = Some(vec!["test_table".to_string()]);
5108 dir_namespace
5109 .create_table(create_request, Bytes::from(buffer.clone()))
5110 .await
5111 .unwrap();
5112
5113 let mut create_request = CreateTableRequest::new();
5115 create_request.id = Some(vec!["test_table".to_string()]);
5116 let result = dir_namespace
5117 .create_table(create_request, Bytes::from(buffer))
5118 .await;
5119 assert!(result.is_err());
5120 }
5121
5122 #[rstest]
5123 #[case::with_optimization(true)]
5124 #[case::without_optimization(false)]
5125 #[tokio::test]
5126 async fn test_create_child_namespace(#[case] inline_optimization: bool) {
5127 use lance_namespace::models::{
5128 CreateNamespaceRequest, ListNamespacesRequest, NamespaceExistsRequest,
5129 };
5130
5131 let temp_dir = TempStdDir::default();
5132 let temp_path = temp_dir.to_str().unwrap();
5133
5134 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5135 .inline_optimization_enabled(inline_optimization)
5136 .build()
5137 .await
5138 .unwrap();
5139
5140 let mut create_req = CreateNamespaceRequest::new();
5142 create_req.id = Some(vec!["ns1".to_string()]);
5143 let result = dir_namespace.create_namespace(create_req).await;
5144 assert!(
5145 result.is_ok(),
5146 "Failed to create child namespace: {:?}",
5147 result.err()
5148 );
5149
5150 let exists_req = NamespaceExistsRequest {
5152 id: Some(vec!["ns1".to_string()]),
5153 ..Default::default()
5154 };
5155 let result = dir_namespace.namespace_exists(exists_req).await;
5156 assert!(result.is_ok(), "Namespace should exist");
5157
5158 let list_req = ListNamespacesRequest {
5160 id: Some(vec![]),
5161 page_token: None,
5162 limit: None,
5163 ..Default::default()
5164 };
5165 let result = dir_namespace.list_namespaces(list_req).await;
5166 assert!(result.is_ok());
5167 let namespaces = result.unwrap();
5168 assert_eq!(namespaces.namespaces.len(), 1);
5169 assert_eq!(namespaces.namespaces[0], "ns1");
5170 }
5171
5172 #[rstest]
5173 #[case::with_optimization(true)]
5174 #[case::without_optimization(false)]
5175 #[tokio::test]
5176 async fn test_create_nested_namespace(#[case] inline_optimization: bool) {
5177 use lance_namespace::models::{
5178 CreateNamespaceRequest, ListNamespacesRequest, NamespaceExistsRequest,
5179 };
5180
5181 let temp_dir = TempStdDir::default();
5182 let temp_path = temp_dir.to_str().unwrap();
5183
5184 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5185 .inline_optimization_enabled(inline_optimization)
5186 .build()
5187 .await
5188 .unwrap();
5189
5190 let mut create_req = CreateNamespaceRequest::new();
5192 create_req.id = Some(vec!["parent".to_string()]);
5193 dir_namespace.create_namespace(create_req).await.unwrap();
5194
5195 let mut create_req = CreateNamespaceRequest::new();
5197 create_req.id = Some(vec!["parent".to_string(), "child".to_string()]);
5198 let result = dir_namespace.create_namespace(create_req).await;
5199 assert!(
5200 result.is_ok(),
5201 "Failed to create nested namespace: {:?}",
5202 result.err()
5203 );
5204
5205 let exists_req = NamespaceExistsRequest {
5207 id: Some(vec!["parent".to_string(), "child".to_string()]),
5208 ..Default::default()
5209 };
5210 let result = dir_namespace.namespace_exists(exists_req).await;
5211 assert!(result.is_ok(), "Nested namespace should exist");
5212
5213 let list_req = ListNamespacesRequest {
5215 id: Some(vec!["parent".to_string()]),
5216 page_token: None,
5217 limit: None,
5218 ..Default::default()
5219 };
5220 let result = dir_namespace.list_namespaces(list_req).await;
5221 assert!(result.is_ok());
5222 let namespaces = result.unwrap();
5223 assert_eq!(namespaces.namespaces.len(), 1);
5224 assert_eq!(namespaces.namespaces[0], "child");
5225 }
5226
5227 #[rstest]
5228 #[case::with_optimization(true)]
5229 #[case::without_optimization(false)]
5230 #[tokio::test]
5231 async fn test_create_namespace_without_parent_fails(#[case] inline_optimization: bool) {
5232 use lance_namespace::models::CreateNamespaceRequest;
5233
5234 let temp_dir = TempStdDir::default();
5235 let temp_path = temp_dir.to_str().unwrap();
5236
5237 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5238 .inline_optimization_enabled(inline_optimization)
5239 .build()
5240 .await
5241 .unwrap();
5242
5243 let mut create_req = CreateNamespaceRequest::new();
5245 create_req.id = Some(vec!["nonexistent_parent".to_string(), "child".to_string()]);
5246 let result = dir_namespace.create_namespace(create_req).await;
5247 assert!(result.is_err(), "Should fail when parent doesn't exist");
5248 }
5249
5250 #[rstest]
5251 #[case::with_optimization(true)]
5252 #[case::without_optimization(false)]
5253 #[tokio::test]
5254 async fn test_drop_child_namespace(#[case] inline_optimization: bool) {
5255 use lance_namespace::models::{
5256 CreateNamespaceRequest, DropNamespaceRequest, NamespaceExistsRequest,
5257 };
5258
5259 let temp_dir = TempStdDir::default();
5260 let temp_path = temp_dir.to_str().unwrap();
5261
5262 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5263 .inline_optimization_enabled(inline_optimization)
5264 .build()
5265 .await
5266 .unwrap();
5267
5268 let mut create_req = CreateNamespaceRequest::new();
5270 create_req.id = Some(vec!["ns1".to_string()]);
5271 dir_namespace.create_namespace(create_req).await.unwrap();
5272
5273 let mut drop_req = DropNamespaceRequest::new();
5275 drop_req.id = Some(vec!["ns1".to_string()]);
5276 let result = dir_namespace.drop_namespace(drop_req).await;
5277 assert!(
5278 result.is_ok(),
5279 "Failed to drop namespace: {:?}",
5280 result.err()
5281 );
5282
5283 let exists_req = NamespaceExistsRequest {
5285 id: Some(vec!["ns1".to_string()]),
5286 ..Default::default()
5287 };
5288 let result = dir_namespace.namespace_exists(exists_req).await;
5289 assert!(result.is_err(), "Namespace should not exist after drop");
5290 }
5291
5292 #[rstest]
5293 #[case::with_optimization(true)]
5294 #[case::without_optimization(false)]
5295 #[tokio::test]
5296 async fn test_drop_namespace_with_children_fails(#[case] inline_optimization: bool) {
5297 use lance_namespace::models::{CreateNamespaceRequest, DropNamespaceRequest};
5298
5299 let temp_dir = TempStdDir::default();
5300 let temp_path = temp_dir.to_str().unwrap();
5301
5302 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5303 .inline_optimization_enabled(inline_optimization)
5304 .build()
5305 .await
5306 .unwrap();
5307
5308 let mut create_req = CreateNamespaceRequest::new();
5310 create_req.id = Some(vec!["parent".to_string()]);
5311 dir_namespace.create_namespace(create_req).await.unwrap();
5312
5313 let mut create_req = CreateNamespaceRequest::new();
5314 create_req.id = Some(vec!["parent".to_string(), "child".to_string()]);
5315 dir_namespace.create_namespace(create_req).await.unwrap();
5316
5317 let mut drop_req = DropNamespaceRequest::new();
5319 drop_req.id = Some(vec!["parent".to_string()]);
5320 let result = dir_namespace.drop_namespace(drop_req).await;
5321 assert!(result.is_err(), "Should fail when namespace has children");
5322 }
5323
5324 #[rstest]
5325 #[case::with_optimization(true)]
5326 #[case::without_optimization(false)]
5327 #[tokio::test]
5328 async fn test_create_table_in_child_namespace(#[case] inline_optimization: bool) {
5329 use lance_namespace::models::{
5330 CreateNamespaceRequest, CreateTableRequest, ListTablesRequest,
5331 };
5332
5333 let temp_dir = TempStdDir::default();
5334 let temp_path = temp_dir.to_str().unwrap();
5335
5336 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5337 .inline_optimization_enabled(inline_optimization)
5338 .build()
5339 .await
5340 .unwrap();
5341
5342 let mut create_ns_req = CreateNamespaceRequest::new();
5344 create_ns_req.id = Some(vec!["ns1".to_string()]);
5345 dir_namespace.create_namespace(create_ns_req).await.unwrap();
5346
5347 let buffer = create_test_ipc_data();
5349 let mut create_table_req = CreateTableRequest::new();
5350 create_table_req.id = Some(vec!["ns1".to_string(), "table1".to_string()]);
5351 let result = dir_namespace
5352 .create_table(create_table_req, Bytes::from(buffer))
5353 .await;
5354 assert!(
5355 result.is_ok(),
5356 "Failed to create table in child namespace: {:?}",
5357 result.err()
5358 );
5359
5360 let list_req = ListTablesRequest {
5362 id: Some(vec!["ns1".to_string()]),
5363 page_token: None,
5364 limit: None,
5365 ..Default::default()
5366 };
5367 let result = dir_namespace.list_tables(list_req).await;
5368 assert!(result.is_ok());
5369 let tables = result.unwrap();
5370 assert_eq!(tables.tables.len(), 1);
5371 assert_eq!(tables.tables[0], "table1");
5372 }
5373
5374 #[rstest]
5375 #[case::with_optimization(true)]
5376 #[case::without_optimization(false)]
5377 #[tokio::test]
5378 async fn test_describe_child_namespace(#[case] inline_optimization: bool) {
5379 use lance_namespace::models::{CreateNamespaceRequest, DescribeNamespaceRequest};
5380
5381 let temp_dir = TempStdDir::default();
5382 let temp_path = temp_dir.to_str().unwrap();
5383
5384 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5385 .inline_optimization_enabled(inline_optimization)
5386 .build()
5387 .await
5388 .unwrap();
5389
5390 let mut properties = std::collections::HashMap::new();
5392 properties.insert("key1".to_string(), "value1".to_string());
5393
5394 let mut create_req = CreateNamespaceRequest::new();
5395 create_req.id = Some(vec!["ns1".to_string()]);
5396 create_req.properties = Some(properties.clone());
5397 dir_namespace.create_namespace(create_req).await.unwrap();
5398
5399 let describe_req = DescribeNamespaceRequest {
5401 id: Some(vec!["ns1".to_string()]),
5402 ..Default::default()
5403 };
5404 let result = dir_namespace.describe_namespace(describe_req).await;
5405 assert!(
5406 result.is_ok(),
5407 "Failed to describe namespace: {:?}",
5408 result.err()
5409 );
5410 let response = result.unwrap();
5411 assert!(response.properties.is_some());
5412 assert_eq!(
5413 response.properties.unwrap().get("key1"),
5414 Some(&"value1".to_string())
5415 );
5416 }
5417
5418 #[rstest]
5419 #[case::with_optimization(true)]
5420 #[case::without_optimization(false)]
5421 #[tokio::test]
5422 async fn test_concurrent_create_and_drop_single_instance(#[case] inline_optimization: bool) {
5423 use futures::future::join_all;
5424 use std::sync::Arc;
5425
5426 let temp_dir = TempStdDir::default();
5427 let temp_path = temp_dir.to_str().unwrap();
5428
5429 let dir_namespace = Arc::new(
5430 DirectoryNamespaceBuilder::new(temp_path)
5431 .inline_optimization_enabled(inline_optimization)
5432 .build()
5433 .await
5434 .unwrap(),
5435 );
5436
5437 let mut create_ns_request = CreateNamespaceRequest::new();
5440 create_ns_request.id = Some(vec!["test_ns".to_string()]);
5441 dir_namespace
5442 .create_namespace(create_ns_request)
5443 .await
5444 .unwrap();
5445
5446 let num_tables = 10;
5447 let mut handles = Vec::new();
5448
5449 for i in 0..num_tables {
5450 let ns = dir_namespace.clone();
5451 let handle = async move {
5452 let table_name = format!("concurrent_table_{}", i);
5453 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5454 let buffer = create_test_ipc_data();
5455
5456 let mut create_request = CreateTableRequest::new();
5458 create_request.id = Some(table_id.clone());
5459 ns.create_table(create_request, Bytes::from(buffer))
5460 .await
5461 .unwrap_or_else(|e| panic!("Failed to create table {}: {}", table_name, e));
5462
5463 let mut drop_request = DropTableRequest::new();
5465 drop_request.id = Some(table_id);
5466 ns.drop_table(drop_request)
5467 .await
5468 .unwrap_or_else(|e| panic!("Failed to drop table {}: {}", table_name, e));
5469
5470 Ok::<_, lance_core::Error>(())
5471 };
5472 handles.push(handle);
5473 }
5474
5475 let results = join_all(handles).await;
5476 for result in results {
5477 assert!(result.is_ok(), "All concurrent operations should succeed");
5478 }
5479
5480 let mut request = ListTablesRequest::new();
5482 request.id = Some(vec!["test_ns".to_string()]);
5483 let response = dir_namespace.list_tables(request).await.unwrap();
5484 assert_eq!(response.tables.len(), 0, "All tables should be dropped");
5485 }
5486
5487 #[rstest]
5488 #[case::with_optimization(true)]
5489 #[case::without_optimization(false)]
5490 #[tokio::test]
5491 async fn test_concurrent_create_and_drop_multiple_instances(#[case] inline_optimization: bool) {
5492 use futures::future::join_all;
5493
5494 let temp_dir = TempStdDir::default();
5495 let temp_path = temp_dir.to_str().unwrap().to_string();
5496
5497 let init_ns = DirectoryNamespaceBuilder::new(&temp_path)
5500 .inline_optimization_enabled(inline_optimization)
5501 .build()
5502 .await
5503 .unwrap();
5504 let mut create_ns_request = CreateNamespaceRequest::new();
5505 create_ns_request.id = Some(vec!["test_ns".to_string()]);
5506 init_ns.create_namespace(create_ns_request).await.unwrap();
5507
5508 let num_tables = 10;
5509 let mut handles = Vec::new();
5510
5511 for i in 0..num_tables {
5512 let path = temp_path.clone();
5513 let handle = async move {
5514 let ns = DirectoryNamespaceBuilder::new(&path)
5516 .inline_optimization_enabled(inline_optimization)
5517 .build()
5518 .await
5519 .unwrap();
5520
5521 let table_name = format!("multi_ns_table_{}", i);
5522 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5523 let buffer = create_test_ipc_data();
5524
5525 let mut create_request = CreateTableRequest::new();
5527 create_request.id = Some(table_id.clone());
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 let mut drop_request = DropTableRequest::new();
5534 drop_request.id = Some(table_id);
5535 ns.drop_table(drop_request)
5536 .await
5537 .unwrap_or_else(|e| panic!("Failed to drop table {}: {}", table_name, e));
5538
5539 Ok::<_, lance_core::Error>(())
5540 };
5541 handles.push(handle);
5542 }
5543
5544 let results = join_all(handles).await;
5545 for result in results {
5546 assert!(result.is_ok(), "All concurrent operations should succeed");
5547 }
5548
5549 let verify_ns = DirectoryNamespaceBuilder::new(&temp_path)
5551 .inline_optimization_enabled(inline_optimization)
5552 .build()
5553 .await
5554 .unwrap();
5555
5556 let mut request = ListTablesRequest::new();
5557 request.id = Some(vec!["test_ns".to_string()]);
5558 let response = verify_ns.list_tables(request).await.unwrap();
5559 assert_eq!(response.tables.len(), 0, "All tables should be dropped");
5560 }
5561
5562 #[rstest]
5563 #[case::with_optimization(true)]
5564 #[case::without_optimization(false)]
5565 #[tokio::test]
5566 async fn test_concurrent_create_then_drop_from_different_instance(
5567 #[case] inline_optimization: bool,
5568 ) {
5569 use futures::future::join_all;
5570
5571 let temp_dir = TempStdDir::default();
5572 let temp_path = temp_dir.to_str().unwrap().to_string();
5573
5574 let init_ns = DirectoryNamespaceBuilder::new(&temp_path)
5577 .inline_optimization_enabled(inline_optimization)
5578 .build()
5579 .await
5580 .unwrap();
5581 let mut create_ns_request = CreateNamespaceRequest::new();
5582 create_ns_request.id = Some(vec!["test_ns".to_string()]);
5583 init_ns.create_namespace(create_ns_request).await.unwrap();
5584
5585 let num_tables = 10;
5586
5587 let mut create_handles = Vec::new();
5589 for i in 0..num_tables {
5590 let path = temp_path.clone();
5591 let handle = async move {
5592 let ns = DirectoryNamespaceBuilder::new(&path)
5593 .inline_optimization_enabled(inline_optimization)
5594 .build()
5595 .await
5596 .unwrap();
5597
5598 let table_name = format!("cross_instance_table_{}", i);
5599 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5600 let buffer = create_test_ipc_data();
5601
5602 let mut create_request = CreateTableRequest::new();
5603 create_request.id = Some(table_id);
5604 ns.create_table(create_request, Bytes::from(buffer))
5605 .await
5606 .unwrap_or_else(|e| panic!("Failed to create table {}: {}", table_name, e));
5607
5608 Ok::<_, lance_core::Error>(())
5609 };
5610 create_handles.push(handle);
5611 }
5612
5613 let create_results = join_all(create_handles).await;
5614 for result in create_results {
5615 assert!(result.is_ok(), "All create operations should succeed");
5616 }
5617
5618 let mut drop_handles = Vec::new();
5620 for i in 0..num_tables {
5621 let path = temp_path.clone();
5622 let handle = async move {
5623 let ns = DirectoryNamespaceBuilder::new(&path)
5624 .inline_optimization_enabled(inline_optimization)
5625 .build()
5626 .await
5627 .unwrap();
5628
5629 let table_name = format!("cross_instance_table_{}", i);
5630 let table_id = vec!["test_ns".to_string(), table_name.clone()];
5631
5632 let mut drop_request = DropTableRequest::new();
5633 drop_request.id = Some(table_id);
5634 ns.drop_table(drop_request)
5635 .await
5636 .unwrap_or_else(|e| panic!("Failed to drop table {}: {}", table_name, e));
5637
5638 Ok::<_, lance_core::Error>(())
5639 };
5640 drop_handles.push(handle);
5641 }
5642
5643 let drop_results = join_all(drop_handles).await;
5644 for result in drop_results {
5645 assert!(result.is_ok(), "All drop operations should succeed");
5646 }
5647
5648 let verify_ns = DirectoryNamespaceBuilder::new(&temp_path)
5650 .inline_optimization_enabled(inline_optimization)
5651 .build()
5652 .await
5653 .unwrap();
5654
5655 let mut request = ListTablesRequest::new();
5656 request.id = Some(vec!["test_ns".to_string()]);
5657 let response = verify_ns.list_tables(request).await.unwrap();
5658 assert_eq!(response.tables.len(), 0, "All tables should be dropped");
5659 }
5660
5661 #[test]
5662 fn test_construct_full_uri_with_cloud_urls() {
5663 let s3_result =
5665 ManifestNamespace::construct_full_uri("s3://bucket/path/subdir", "table.lance")
5666 .unwrap();
5667 assert_eq!(
5668 s3_result, "s3://bucket/path/subdir/table.lance",
5669 "S3 URL should correctly append table name to nested path"
5670 );
5671
5672 let az_result =
5674 ManifestNamespace::construct_full_uri("az://container/path/subdir", "table.lance")
5675 .unwrap();
5676 assert_eq!(
5677 az_result, "az://container/path/subdir/table.lance",
5678 "Azure URL should correctly append table name to nested path"
5679 );
5680
5681 let gs_result =
5683 ManifestNamespace::construct_full_uri("gs://bucket/path/subdir", "table.lance")
5684 .unwrap();
5685 assert_eq!(
5686 gs_result, "gs://bucket/path/subdir/table.lance",
5687 "GCS URL should correctly append table name to nested path"
5688 );
5689
5690 let deep_result =
5692 ManifestNamespace::construct_full_uri("s3://bucket/a/b/c/d", "my_table.lance").unwrap();
5693 assert_eq!(
5694 deep_result, "s3://bucket/a/b/c/d/my_table.lance",
5695 "Deeply nested path should work correctly"
5696 );
5697
5698 let shallow_result =
5700 ManifestNamespace::construct_full_uri("s3://bucket", "table.lance").unwrap();
5701 assert_eq!(
5702 shallow_result, "s3://bucket/table.lance",
5703 "Single-level nested path should work correctly"
5704 );
5705
5706 let trailing_slash_result =
5708 ManifestNamespace::construct_full_uri("s3://bucket/path/subdir/", "table.lance")
5709 .unwrap();
5710 assert_eq!(
5711 trailing_slash_result, "s3://bucket/path/subdir/table.lance",
5712 "URL with existing trailing slash should still work"
5713 );
5714
5715 let empty_query_result =
5718 ManifestNamespace::construct_full_uri("s3://bucket/path?", "table.lance").unwrap();
5719 assert_eq!(
5720 empty_query_result, "s3://bucket/path/table.lance",
5721 "URL with empty query string should not include trailing '?'"
5722 );
5723
5724 let query_param_result =
5727 ManifestNamespace::construct_full_uri("s3://bucket/path?param=value", "table.lance")
5728 .unwrap();
5729 assert_eq!(
5730 query_param_result, "s3://bucket/path/table.lance",
5731 "URL with query parameters should have them stripped"
5732 );
5733 }
5734
5735 #[test]
5736 fn test_construct_full_uri_with_dollar_sign() {
5737 let result =
5738 ManifestNamespace::construct_full_uri("/tmp/root", "hash_workspace$test_table")
5739 .unwrap();
5740
5741 assert!(
5742 result.ends_with("/tmp/root/hash_workspace$test_table"),
5743 "local file URI should preserve dollar signs without adding empty path segments: {}",
5744 result
5745 );
5746 assert!(
5747 !result.contains("//hash_workspace$test_table"),
5748 "local file URI should not add a double slash before table directory: {}",
5749 result
5750 );
5751 }
5752
5753 #[test]
5754 fn test_construct_full_uri_with_nested_relative_location() {
5755 let result =
5756 ManifestNamespace::construct_full_uri("/tmp/root", "workspace/physical_table.lance")
5757 .unwrap();
5758
5759 assert!(
5760 result.ends_with("/tmp/root/workspace/physical_table.lance"),
5761 "nested relative location should preserve path separators: {}",
5762 result
5763 );
5764 assert!(
5765 !result.contains("%2Fphysical_table.lance"),
5766 "nested relative location should not encode path separators: {}",
5767 result
5768 );
5769 }
5770
5771 #[tokio::test]
5778 async fn test_concurrent_create_table_no_duplicates() {
5779 let temp_dir = TempStdDir::default();
5780 let temp_path = temp_dir.to_str().unwrap();
5781
5782 let ns1 = DirectoryNamespaceBuilder::new(temp_path)
5785 .inline_optimization_enabled(false)
5786 .build()
5787 .await
5788 .unwrap();
5789 let ns2 = DirectoryNamespaceBuilder::new(temp_path)
5790 .inline_optimization_enabled(false)
5791 .build()
5792 .await
5793 .unwrap();
5794
5795 let buffer = create_test_ipc_data();
5796
5797 let mut req1 = CreateTableRequest::new();
5798 req1.id = Some(vec!["race_table".to_string()]);
5799 let mut req2 = CreateTableRequest::new();
5800 req2.id = Some(vec!["race_table".to_string()]);
5801
5802 let (result1, result2) = tokio::join!(
5804 ns1.create_table(req1, Bytes::from(buffer.clone())),
5805 ns2.create_table(req2, Bytes::from(buffer.clone())),
5806 );
5807
5808 let success_count = [&result1, &result2].iter().filter(|r| r.is_ok()).count();
5810 let failure_count = [&result1, &result2].iter().filter(|r| r.is_err()).count();
5811 assert_eq!(
5812 success_count, 1,
5813 "Exactly one create should succeed, got: result1={:?}, result2={:?}",
5814 result1, result2
5815 );
5816 assert_eq!(
5817 failure_count, 1,
5818 "Exactly one create should fail, got: result1={:?}, result2={:?}",
5819 result1, result2
5820 );
5821
5822 let ns_check = DirectoryNamespaceBuilder::new(temp_path)
5824 .inline_optimization_enabled(false)
5825 .build()
5826 .await
5827 .unwrap();
5828 let mut list_request = ListTablesRequest::new();
5829 list_request.id = Some(vec![]);
5830 let response = ns_check.list_tables(list_request).await.unwrap();
5831 assert_eq!(
5832 response.tables.len(),
5833 1,
5834 "Should have exactly 1 table, found: {:?}",
5835 response.tables
5836 );
5837 assert_eq!(response.tables[0], "race_table");
5838
5839 let mut describe_request = DescribeTableRequest::new();
5841 describe_request.id = Some(vec!["race_table".to_string()]);
5842 let describe_result = ns_check.describe_table(describe_request).await;
5843 assert!(
5844 describe_result.is_ok(),
5845 "describe_table should not fail with duplicate entries: {:?}",
5846 describe_result
5847 );
5848 }
5849
5850 fn names(v: &[&str]) -> Vec<String> {
5853 v.iter().map(|s| s.to_string()).collect()
5854 }
5855
5856 #[test]
5857 fn test_apply_pagination_no_token_no_limit() {
5858 let mut n = names(&["b", "a", "c"]);
5859 let next = ManifestNamespace::apply_pagination(&mut n, None, None);
5860 assert_eq!(n, names(&["a", "b", "c"]));
5861 assert_eq!(next, None);
5862 }
5863
5864 #[test]
5865 fn test_apply_pagination_limit_truncates_and_returns_token() {
5866 let mut n = names(&["c", "a", "b"]);
5867 let next = ManifestNamespace::apply_pagination(&mut n, None, Some(2));
5868 assert_eq!(n, names(&["a", "b"]));
5869 assert_eq!(next, Some("b".to_string()));
5870 }
5871
5872 #[test]
5873 fn test_apply_pagination_limit_zero_returns_empty_no_token() {
5874 let mut n = names(&["a", "b", "c"]);
5875 let next = ManifestNamespace::apply_pagination(&mut n, None, Some(0));
5876 assert!(n.is_empty());
5877 assert_eq!(next, None);
5878 }
5879
5880 #[test]
5881 fn test_apply_pagination_page_token_in_list() {
5882 let mut n = names(&["a", "b", "c", "d"]);
5884 let next = ManifestNamespace::apply_pagination(&mut n, Some("b".to_string()), None);
5885 assert_eq!(n, names(&["c", "d"]));
5886 assert_eq!(next, None);
5887 }
5888
5889 #[test]
5890 fn test_apply_pagination_page_token_past_all_items() {
5891 let mut n = names(&["a", "b", "c"]);
5892 let next = ManifestNamespace::apply_pagination(&mut n, Some("z".to_string()), None);
5893 assert!(n.is_empty());
5894 assert_eq!(next, None);
5895 }
5896
5897 #[test]
5898 fn test_apply_pagination_token_and_limit_combined() {
5899 let mut n = names(&["a", "b", "c", "d", "e"]);
5900 let next = ManifestNamespace::apply_pagination(&mut n, Some("b".to_string()), Some(2));
5901 assert_eq!(n, names(&["c", "d"]));
5902 assert_eq!(next, Some("d".to_string()));
5903 }
5904
5905 #[rstest]
5906 #[case::with_optimization(true)]
5907 #[case::without_optimization(false)]
5908 #[tokio::test]
5909 async fn test_alter_table_add_columns(#[case] inline_optimization: bool) {
5910 use lance_namespace::models::{
5911 AddColumnsEntry, AlterTableAddColumnsRequest, DescribeTableRequest,
5912 };
5913
5914 let temp_dir = TempStdDir::default();
5915 let temp_path = temp_dir.to_str().unwrap();
5916
5917 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5918 .inline_optimization_enabled(inline_optimization)
5919 .build()
5920 .await
5921 .unwrap();
5922
5923 let buffer = create_test_ipc_data();
5925 let mut create_request = CreateTableRequest::new();
5926 create_request.id = Some(vec!["test_table".to_string()]);
5927 dir_namespace
5928 .create_table(create_request, Bytes::from(buffer))
5929 .await
5930 .unwrap();
5931
5932 let mut new_col = AddColumnsEntry::new("doubled_id".to_string());
5934 new_col.expression = Some(Some("id * 2".to_string()));
5935 let mut add_request = AlterTableAddColumnsRequest::new(vec![new_col]);
5936 add_request.id = Some(vec!["test_table".to_string()]);
5937
5938 let response = dir_namespace
5939 .alter_table_add_columns(add_request)
5940 .await
5941 .unwrap();
5942 assert!(response.version > 1);
5944
5945 let mut describe_request = DescribeTableRequest::new();
5947 describe_request.id = Some(vec!["test_table".to_string()]);
5948 describe_request.load_detailed_metadata = Some(true);
5949 let describe_response = dir_namespace
5950 .describe_table(describe_request)
5951 .await
5952 .unwrap();
5953 assert!(describe_response.schema.is_some());
5954
5955 let schema = describe_response.schema.unwrap();
5956 let field_names: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect();
5957 assert!(
5958 field_names.contains(&"doubled_id"),
5959 "Column 'doubled_id' should exist after add_columns, got: {:?}",
5960 field_names
5961 );
5962 }
5963
5964 #[rstest]
5965 #[case::with_optimization(true)]
5966 #[case::without_optimization(false)]
5967 #[tokio::test]
5968 async fn test_alter_table_add_columns_missing_id(#[case] inline_optimization: bool) {
5969 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
5970
5971 let temp_dir = TempStdDir::default();
5972 let temp_path = temp_dir.to_str().unwrap();
5973
5974 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5975 .inline_optimization_enabled(inline_optimization)
5976 .build()
5977 .await
5978 .unwrap();
5979
5980 let new_col = AddColumnsEntry::new("col".to_string());
5982 let request = AlterTableAddColumnsRequest::new(vec![new_col]);
5983 let result = dir_namespace.alter_table_add_columns(request).await;
5984 assert!(result.is_err(), "Should fail when table ID is missing");
5985 }
5986
5987 #[rstest]
5988 #[case::with_optimization(true)]
5989 #[case::without_optimization(false)]
5990 #[tokio::test]
5991 async fn test_alter_table_add_columns_nonexistent_table(#[case] inline_optimization: bool) {
5992 use lance_namespace::models::{AddColumnsEntry, AlterTableAddColumnsRequest};
5993
5994 let temp_dir = TempStdDir::default();
5995 let temp_path = temp_dir.to_str().unwrap();
5996
5997 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
5998 .inline_optimization_enabled(inline_optimization)
5999 .build()
6000 .await
6001 .unwrap();
6002
6003 let new_col = AddColumnsEntry::new("col".to_string());
6005 let mut request = AlterTableAddColumnsRequest::new(vec![new_col]);
6006 request.id = Some(vec!["nonexistent".to_string()]);
6007 let result = dir_namespace.alter_table_add_columns(request).await;
6008 assert!(result.is_err(), "Should fail when table does not exist");
6009 }
6010
6011 #[rstest]
6012 #[case::with_optimization(true)]
6013 #[case::without_optimization(false)]
6014 #[tokio::test]
6015 async fn test_alter_table_alter_columns_rename(#[case] inline_optimization: bool) {
6016 use lance_namespace::models::{
6017 AlterColumnsEntry, AlterTableAlterColumnsRequest, DescribeTableRequest,
6018 };
6019
6020 let temp_dir = TempStdDir::default();
6021 let temp_path = temp_dir.to_str().unwrap();
6022
6023 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6024 .inline_optimization_enabled(inline_optimization)
6025 .build()
6026 .await
6027 .unwrap();
6028
6029 let buffer = create_test_ipc_data();
6031 let mut create_request = CreateTableRequest::new();
6032 create_request.id = Some(vec!["test_table".to_string()]);
6033 dir_namespace
6034 .create_table(create_request, Bytes::from(buffer))
6035 .await
6036 .unwrap();
6037
6038 let mut entry = AlterColumnsEntry::new("name".to_string());
6040 entry.rename = Some(Some("full_name".to_string()));
6041 let mut alter_request = AlterTableAlterColumnsRequest::new(vec![entry]);
6042 alter_request.id = Some(vec!["test_table".to_string()]);
6043
6044 let response = dir_namespace
6045 .alter_table_alter_columns(alter_request)
6046 .await
6047 .unwrap();
6048 assert!(response.version > 1);
6049
6050 let mut describe_request = DescribeTableRequest::new();
6052 describe_request.id = Some(vec!["test_table".to_string()]);
6053 describe_request.load_detailed_metadata = Some(true);
6054 let describe_response = dir_namespace
6055 .describe_table(describe_request)
6056 .await
6057 .unwrap();
6058 assert!(describe_response.schema.is_some());
6059
6060 let schema = describe_response.schema.unwrap();
6061 let field_names: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect();
6062 assert!(
6063 field_names.contains(&"full_name"),
6064 "Column should be renamed to 'full_name', got: {:?}",
6065 field_names
6066 );
6067 assert!(
6068 !field_names.contains(&"name"),
6069 "Old column name 'name' should no longer exist, got: {:?}",
6070 field_names
6071 );
6072 }
6073
6074 #[rstest]
6075 #[case::with_optimization(true)]
6076 #[case::without_optimization(false)]
6077 #[tokio::test]
6078 async fn test_alter_table_alter_columns_missing_id(#[case] inline_optimization: bool) {
6079 use lance_namespace::models::{AlterColumnsEntry, AlterTableAlterColumnsRequest};
6080
6081 let temp_dir = TempStdDir::default();
6082 let temp_path = temp_dir.to_str().unwrap();
6083
6084 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6085 .inline_optimization_enabled(inline_optimization)
6086 .build()
6087 .await
6088 .unwrap();
6089
6090 let entry = AlterColumnsEntry::new("name".to_string());
6091 let request = AlterTableAlterColumnsRequest::new(vec![entry]);
6092 let result = dir_namespace.alter_table_alter_columns(request).await;
6093 assert!(result.is_err(), "Should fail when table ID is missing");
6094 }
6095
6096 #[rstest]
6097 #[case::with_optimization(true)]
6098 #[case::without_optimization(false)]
6099 #[tokio::test]
6100 async fn test_alter_table_drop_columns(#[case] inline_optimization: bool) {
6101 use lance_namespace::models::{AlterTableDropColumnsRequest, DescribeTableRequest};
6102
6103 let temp_dir = TempStdDir::default();
6104 let temp_path = temp_dir.to_str().unwrap();
6105
6106 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6107 .inline_optimization_enabled(inline_optimization)
6108 .build()
6109 .await
6110 .unwrap();
6111
6112 let buffer = create_test_ipc_data();
6114 let mut create_request = CreateTableRequest::new();
6115 create_request.id = Some(vec!["test_table".to_string()]);
6116 dir_namespace
6117 .create_table(create_request, Bytes::from(buffer))
6118 .await
6119 .unwrap();
6120
6121 let mut drop_request = AlterTableDropColumnsRequest::new(vec!["name".to_string()]);
6123 drop_request.id = Some(vec!["test_table".to_string()]);
6124
6125 let response = dir_namespace
6126 .alter_table_drop_columns(drop_request)
6127 .await
6128 .unwrap();
6129 assert!(response.version > 1);
6130
6131 let mut describe_request = DescribeTableRequest::new();
6133 describe_request.id = Some(vec!["test_table".to_string()]);
6134 describe_request.load_detailed_metadata = Some(true);
6135 let describe_response = dir_namespace
6136 .describe_table(describe_request)
6137 .await
6138 .unwrap();
6139 assert!(describe_response.schema.is_some());
6140
6141 let schema = describe_response.schema.unwrap();
6142 let field_names: Vec<&str> = schema.fields.iter().map(|f| f.name.as_str()).collect();
6143 assert!(
6144 !field_names.contains(&"name"),
6145 "Column 'name' should have been dropped, got: {:?}",
6146 field_names
6147 );
6148 assert!(
6149 field_names.contains(&"id"),
6150 "Column 'id' should still exist, got: {:?}",
6151 field_names
6152 );
6153 }
6154
6155 #[rstest]
6156 #[case::with_optimization(true)]
6157 #[case::without_optimization(false)]
6158 #[tokio::test]
6159 async fn test_alter_table_drop_columns_missing_id(#[case] inline_optimization: bool) {
6160 use lance_namespace::models::AlterTableDropColumnsRequest;
6161
6162 let temp_dir = TempStdDir::default();
6163 let temp_path = temp_dir.to_str().unwrap();
6164
6165 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6166 .inline_optimization_enabled(inline_optimization)
6167 .build()
6168 .await
6169 .unwrap();
6170
6171 let request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
6172 let result = dir_namespace.alter_table_drop_columns(request).await;
6173 assert!(result.is_err(), "Should fail when table ID is missing");
6174 }
6175
6176 #[rstest]
6177 #[case::with_optimization(true)]
6178 #[case::without_optimization(false)]
6179 #[tokio::test]
6180 async fn test_alter_table_drop_columns_nonexistent_table(#[case] inline_optimization: bool) {
6181 use lance_namespace::models::AlterTableDropColumnsRequest;
6182
6183 let temp_dir = TempStdDir::default();
6184 let temp_path = temp_dir.to_str().unwrap();
6185
6186 let dir_namespace = DirectoryNamespaceBuilder::new(temp_path)
6187 .inline_optimization_enabled(inline_optimization)
6188 .build()
6189 .await
6190 .unwrap();
6191
6192 let mut request = AlterTableDropColumnsRequest::new(vec!["col".to_string()]);
6193 request.id = Some(vec!["nonexistent".to_string()]);
6194 let result = dir_namespace.alter_table_drop_columns(request).await;
6195 assert!(result.is_err(), "Should fail when table does not exist");
6196 }
6197}