push more things down to rust layer

Author: jackye1995

protos/transaction.proto

@@ -68,8 +68,9 @@ message Transaction {
     map<string, bytes> schema_metadata = 3;
     // Key-value pairs to merge with existing config.
     map<string, string> config_upsert_values = 4;
-    // Additional base paths for data file distribution in multi-path layouts
-    repeated BasePath initial_data_paths = 5;
+    // Additional base paths for data file distribution in multi-path layouts.
+    // Each BasePath must have a properly assigned ID (non-zero).
+    repeated BasePath initial_bases = 5;
   }
 
   // Add or replace a new secondary index.

python/python/lance/__init__.py

@@ -11,7 +11,6 @@
 from . import log
 from .blob import BlobColumn, BlobFile
 from .dataset import (
-    DatasetBasePathConfig,
     DataStatistics,
     FieldStatistics,
     Index,
@@ -27,6 +26,7 @@
 )
 from .fragment import FragmentMetadata, LanceFragment
 from .lance import (
+    DatasetBasePath,
     FFILanceTableProvider,
     ScanStatistics,
     bytes_read_counter,
@@ -48,7 +48,7 @@
 __all__ = [
     "BlobColumn",
     "BlobFile",
-    "DatasetBasePathConfig",
+    "DatasetBasePath",
     "DataStatistics",
     "FieldStatistics",
     "FragmentMetadata",

python/python/lance/dataset.py

@@ -52,6 +52,7 @@
     CleanupStats,
     Compaction,
     CompactionMetrics,
+    DatasetBasePath,
     LanceSchema,
     ScanStatistics,
     _Dataset,
@@ -3550,39 +3551,6 @@ class AutoCleanupConfig(TypedDict):
     older_than_seconds: int
 
 
-@dataclass
-class DatasetBasePathConfig:
-    """Configuration for a dataset base path.
-
-    Used to specify additional storage locations for dataset data files.
-    This enables multi-bucket/multi-path dataset configurations.
-
-    Parameters
-    ----------
-    name : str
-        A unique name identifier for this base path. Used to reference
-        this base in subsequent operations.
-    path : str
-        The full URI of the base path (e.g., "s3://bucket/path", "/mnt/storage/data").
-        This is the root location where data will be stored.
-    is_dataset_root : bool, default False
-        Whether this path is a dataset root directory (with subdirectories like
-        'data/', '_deletions/') or a direct file directory. Default is False,
-        meaning data will be written directly to the path/data directory without
-        additional subdirectories for deletions, etc.
-
-    Examples
-    --------
-    >>> base1 = DatasetBasePathConfig(name="primary", path="s3://bucket1/dataset")
-    >>> base2 = DatasetBasePathConfig(name="archive", path="s3://bucket2/dataset",
-    ...                           is_dataset_root=True)
-    """
-
-    name: str
-    path: str
-    is_dataset_root: bool = False
-
-
 # LanceOperation is a namespace for operations that can be applied to a dataset.
 class LanceOperation:
     @staticmethod
@@ -5098,8 +5066,8 @@ def write_dataset(
             # Dataset doesn't exist yet, which is fine for append mode
             pass
 
-    # Resolve target_bases into new_bases and target_paths
-    new_bases, target_paths = _resolve_target_bases(
+    # Resolve target_bases into new_bases and target_base_ids
+    new_bases, target_base_ids = _resolve_target_bases(
         target_bases, existing_dataset, mode
     )
 
@@ -5116,7 +5084,7 @@ def write_dataset(
         "auto_cleanup_options": auto_cleanup_options,
         "transaction_properties": merged_properties,
         "new_bases": new_bases,
-        "target_paths": target_paths,
+        "target_bases": target_base_ids,
     }
 
     if commit_lock:
@@ -5238,24 +5206,19 @@ def _merge_message_to_properties(
     return merged_properties
 
 
-def _normalize_path(path: str) -> str:
-    """Normalize a path URI by removing trailing slashes."""
-    return path.rstrip("/")
-
-
 def _resolve_target_bases(
-    target_bases: Optional[List[Union[str, DatasetBasePathConfig]]],
+    target_bases: Optional[List[Union[str, DatasetBasePath]]],
     existing_dataset: Optional["LanceDataset"],
     mode: str,
-) -> Tuple[Optional[List[Dict[str, Any]]], Optional[List[str]]]:
+) -> Tuple[Optional[List[Dict[str, Any]]], Optional[List[int]]]:
     """
-    Parse target_bases into (new_bases, target_paths).
+    Parse target_bases into (new_bases, target_base_ids).
 
     Parameters
     ----------
-    target_bases : List[Union[str, DatasetBasePathConfig]], optional
+    target_bases : List[Union[str, DatasetBasePath]], optional
         List of base specifications. Can be:
-        - DatasetBasePathConfig objects (new bases to register)
+        - DatasetBasePath objects (new bases to register)
         - Strings (references to existing bases by name or path)
     existing_dataset : LanceDataset, optional
         The existing dataset (for append/overwrite modes)
@@ -5266,12 +5229,12 @@ def _resolve_target_bases(
     -------
     new_bases : List[Dict[str, Any]], optional
         List of new base configurations to register. Each dict contains:
+        - id: int (assigned sequentially starting from 1)
         - name: str
         - path: str
         - is_dataset_root: bool
-        IDs will be assigned later in Rust.
-    target_paths : List[str], optional
-        List of path URIs where data should be written.
+    target_base_ids : List[int], optional
+        List of base IDs where data should be written.
 
     Raises
     ------
@@ -5282,115 +5245,90 @@ def _resolve_target_bases(
         return None, None
 
     new_bases = []
-    target_paths = []
+    target_base_ids = []
+    next_base_id = 1  # Start assigning IDs from 1
 
     # Get existing bases from manifest if available
     existing_bases_by_name = {}
     existing_bases_by_path = {}
     if existing_dataset is not None:
-        try:
-            base_paths_dict = existing_dataset._ds.base_paths()
-            for base_path_dict in base_paths_dict.values():
-                # Convert dict to a simple object for easier access
-                class BasePath:
-                    def __init__(self, d):
-                        self.id = d["id"]
-                        self.name = d["name"]
-                        self.path = d["path"]
-                        self.is_dataset_root = d["is_dataset_root"]
-
-                base_path = BasePath(base_path_dict)
-                normalized_path = _normalize_path(base_path.path)
-                existing_bases_by_path[normalized_path] = base_path
-                if base_path.name:
-                    existing_bases_by_name[base_path.name] = base_path
-        except Exception:
-            # If we can't load base paths, assume no existing bases
-            pass
+        base_paths_dict = existing_dataset._ds.base_paths()
+        for base_path in base_paths_dict.values():
+            existing_bases_by_path[base_path.path] = base_path
+            if base_path.name:
+                existing_bases_by_name[base_path.name] = base_path
 
-    # Track new bases being added in this operation (by name and path)
+    # Track new bases being added in this operation (by name, path, and ID)
     new_bases_by_name = {}
     new_bases_by_path = {}
 
     for item in target_bases:
-        if isinstance(item, DatasetBasePathConfig):
-            # Validate that new bases can only be registered in CREATE or OVERWRITE mode
-            if mode == "append":
+        if isinstance(item, DatasetBasePath):
+            # Validate that new bases can only be registered in CREATE mode
+            if mode != "create":
                 raise ValueError(
-                    "Cannot register new bases in APPEND mode. "
-                    "New bases can only be registered in CREATE or OVERWRITE "
-                    "mode. For APPEND mode, reference existing bases by name "
-                    "or path as strings."
+                    f"Cannot register new bases in {mode.upper()} mode. "
+                    "New bases can only be registered in CREATE mode. "
+                    f"For {mode.upper()} mode, reference existing bases by name "
+                    "or path as strings. OVERWRITE mode always inherits existing base_paths."
                 )
 
-            # New base definition
+            # Assign ID to this base
+            base_id = next_base_id
+            next_base_id += 1
+
+            # New base definition with assigned ID
             base_dict = {
+                "id": base_id,
                 "name": item.name,
                 "path": item.path,
                 "is_dataset_root": item.is_dataset_root,
             }
             new_bases.append(base_dict)
-            # Don't add to target_paths - DatasetBasePathConfig is just registering
+            # Don't add to target_base_ids - DatasetBasePath is just registering
             # a base, not indicating where to write. Only string references
             # specify write targets.
 
             # Track for resolution within this same operation
-            new_bases_by_name[item.name] = item
-            new_bases_by_path[_normalize_path(item.path)] = item
+            # Store the base with its assigned ID
+            temp_base = type('obj', (object,), {
+                'id': base_id,
+                'name': item.name,
+                'path': item.path,
+                'is_dataset_root': item.is_dataset_root
+            })()
+            new_bases_by_name[item.name] = temp_base
+            new_bases_by_path[item.path] = temp_base
 
         elif isinstance(item, str):
             # String reference - resolve by name first, then by path
             resolved = False
 
             # 1. Try to match by name in newly defined bases
             if item in new_bases_by_name:
-                target_paths.append(new_bases_by_name[item].path)
+                target_base_ids.append(new_bases_by_name[item].id)
                 resolved = True
             # 2. Try to match by name in existing bases
             elif item in existing_bases_by_name:
-                target_paths.append(existing_bases_by_name[item].path)
+                target_base_ids.append(existing_bases_by_name[item].id)
                 resolved = True
             # 3. Try to match by path in newly defined bases
-            elif _normalize_path(item) in new_bases_by_path:
-                target_paths.append(item)
+            elif item in new_bases_by_path:
+                target_base_ids.append(new_bases_by_path[item].id)
                 resolved = True
             # 4. Try to match by path in existing bases
-            elif _normalize_path(item) in existing_bases_by_path:
-                target_paths.append(item)
+            elif item in existing_bases_by_path:
+                target_base_ids.append(existing_bases_by_path[item].id)
                 resolved = True
 
             if not resolved:
-                # Build helpful error message
-                available_names = list(existing_bases_by_name.keys()) + list(
-                    new_bases_by_name.keys()
-                )
-                available_paths = list(existing_bases_by_path.keys()) + list(
-                    new_bases_by_path.keys()
-                )
-
-                error_msg = f"Base '{item}' not found. "
-                if available_names or available_paths:
-                    error_msg += "Available bases:\n"
-                    for name in available_names:
-                        base = existing_bases_by_name.get(
-                            name
-                        ) or new_bases_by_name.get(name)
-                        error_msg += f"  - {name}: {base.path}\n"
-                    for path in available_paths:
-                        if path not in [
-                            b.path for b in existing_bases_by_name.values()
-                        ] and path not in [b.path for b in new_bases_by_name.values()]:
-                            error_msg += f"  - {path}\n"
-                else:
-                    error_msg += "No bases are currently registered."
-
-                raise ValueError(error_msg)
+                raise ValueError(f"Base '{item}' not found.")
         else:
             raise TypeError(
-                f"target_bases must contain str or DatasetBasePathConfig, got {type(item)}"
+                f"target_bases must contain str or DatasetBasePath, got {type(item)}"
             )
 
-    return (new_bases if new_bases else None, target_paths if target_paths else None)
+    return (new_bases if new_bases else None, target_base_ids if target_base_ids else None)
 
 
 class VectorIndexReader: