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