feat: expose base scoped store bindings to python (#6547)

Tested: Python can write to different S3 buckets within the same
multi-base dataset, using different storage configs for each
base.configs.

```
from pathlib import Path
import shutil

import lance
import pyarrow as pa
from lance import DatasetBasePath


def main():
    path0 = "s3://ceshi/las/lance/multi-base-write2"
    path1 = "s3://zy-test-lance/zy-test-lance/multi-base-write2/"

    storage_options_0 = {
        "access_key_id": "",
        "secret_access_key": "==",
        "aws_endpoint": "",
        "aws_region": "",
        "virtual_hosted_style_request": "true",
    }

    storage_options_1 = {
        "access_key_id": "",
        "secret_access_key": "==",
        "aws_region": "",
        "aws_endpoint": "",
        "virtual_hosted_style_request": "true"
    }

    base_store_params = {
        path1: storage_options_1,
    }

    local_rows = pa.Table.from_pylist(
        [
            {"name": "Alice", "age": 20},
            {"name": "Jack", "age": 20},
        ]
    )

    s3_rows = pa.Table.from_pylist(
        [
            {"name": "Bob", "age": 30},
        ]
    )

    ds = lance.write_dataset(
        local_rows,
        path0,
        mode="create",
        initial_bases=[
            DatasetBasePath(path0, name="s3", is_dataset_root=True),
            DatasetBasePath(path1, name="s3"),
        ],
        target_bases=["s3"],
        base_store_params=base_store_params,
        storage_options=storage_options_0,
    )

    ds = lance.write_dataset(
        s3_rows,
        ds,
        mode="append",
        target_bases=["s3"],
        base_store_params=base_store_params,
        storage_options=storage_options_0,
    )


if __name__ == "__main__":
    main()


```

---------

Co-authored-by: zhangyue19921010 <zhangyue.1010@bytedance.com>

Author: zhangyue19921010

python/python/lance/__init__.py

@@ -101,6 +101,7 @@ def dataset(
     session: Optional[Session] = None,
     namespace_client: Optional[LanceNamespace] = None,
     table_id: Optional[List[str]] = None,
+    base_store_params: Optional[Dict[str, Dict[str, str]]] = None,
 ) -> LanceDataset:
     """
     Opens the Lance dataset from the address specified.
@@ -169,6 +170,12 @@ def dataset(
     table_id : optional, List[str]
         The table identifier when using a namespace (e.g., ["my_table"]).
         Must be provided together with `namespace_client`. Cannot be used with `uri`.
+    base_store_params : dict of str to dict, optional
+        Runtime-only object store parameters keyed by base path URI. Each key
+        is a base path URI (e.g., "s3://bucket/path") and each value is a dict
+        of storage options (credentials, endpoint, etc.) for that base.  When a base
+        has no explicit entry here, the top-level ``storage_options`` is
+        used as a fallback.
 
     Notes
     -----
@@ -244,6 +251,7 @@ def dataset(
         namespace_client=namespace_client,
         table_id=table_id,
         namespace_client_managed_versioning=namespace_client_managed_versioning,
+        base_store_params=base_store_params,
     )
     if version is None and asof is not None:
         ts_cutoff = sanitize_ts(asof)
@@ -270,6 +278,7 @@ def dataset(
                 namespace_client=namespace_client,
                 table_id=table_id,
                 namespace_client_managed_versioning=namespace_client_managed_versioning,
+                base_store_params=base_store_params,
             )
     else:
         return ds

python/python/lance/dataset.py

@@ -583,10 +583,12 @@ def __init__(
         namespace_client: Optional[Any] = None,
         table_id: Optional[List[str]] = None,
         namespace_client_managed_versioning: bool = False,
+        base_store_params: Optional[Dict[str, Dict[str, str]]] = None,
     ):
         uri = os.fspath(uri) if isinstance(uri, Path) else uri
         self._uri = uri
         self._storage_options = storage_options
+        self._base_store_params = base_store_params
 
         # Handle deprecation warning for index_cache_size
         if index_cache_size is not None:
@@ -621,6 +623,7 @@ def __init__(
             namespace_client=namespace_client,
             table_id=table_id,
             namespace_client_managed_versioning=namespace_client_managed_versioning,
+            base_store_params=base_store_params,
         )
         self._default_scan_options = default_scan_options
         self._read_params = read_params
@@ -634,6 +637,7 @@ def __deserialize__(
         manifest: bytes,
         default_scan_options: Optional[Dict[str, Any]],
         read_params: Optional[Dict[str, Any]] = None,
+        base_store_params: Optional[Dict[str, Dict[str, str]]] = None,
     ):
         return cls(
             uri,
@@ -642,6 +646,7 @@ def __deserialize__(
             serialized_manifest=manifest,
             default_scan_options=default_scan_options,
             read_params=read_params,
+            base_store_params=base_store_params,
         )
 
     def __reduce__(self):
@@ -652,6 +657,7 @@ def __reduce__(self):
             self._ds.serialized_manifest(),
             self._default_scan_options,
             self._read_params,
+            self._base_store_params,
         )
 
     def __getstate__(self):
@@ -662,29 +668,34 @@ def __getstate__(self):
             self._ds.serialized_manifest(),
             self._default_scan_options,
             self._read_params,
+            self._base_store_params,
         )
 
     def __setstate__(self, state):
-        # Handle backwards compatibility - state may not have read_params
+        # Handle backwards compatibility - state may not have read_params or
+        # base_store_params.
         (
             self._uri,
             self._storage_options,
             version,
             manifest,
             default_scan_options,
-            *rest,  # Capture optional read_params
+            *rest,  # Capture optional read_params and base_store_params.
         ) = state
         read_params = rest[0] if rest else None
+        base_store_params = rest[1] if len(rest) > 1 else None
         self._ds = _Dataset(
             self._uri,
             version,
             storage_options=self._storage_options,
             manifest=manifest,
             default_scan_options=default_scan_options,
             read_params=read_params,
+            base_store_params=base_store_params,
         )
         self._default_scan_options = default_scan_options
         self._read_params = read_params
+        self._base_store_params = base_store_params
         self._namespace_client = None
         self._table_id = None
         self._namespace_client_managed_versioning = False
@@ -693,6 +704,7 @@ def __copy__(self):
         ds = LanceDataset.__new__(LanceDataset)
         ds._uri = self._uri
         ds._storage_options = self._storage_options
+        ds._base_store_params = self._base_store_params
         ds._namespace_client = self._namespace_client
         ds._table_id = self._table_id
         ds._namespace_client_managed_versioning = (
@@ -792,6 +804,7 @@ def create_branch(
         ds._ds = new_ds
         ds._uri = new_ds.uri
         ds._storage_options = self._storage_options
+        ds._base_store_params = self._base_store_params
         ds._namespace_client = self._namespace_client
         ds._table_id = self._table_id
         ds._default_scan_options = self._default_scan_options
@@ -3902,6 +3915,15 @@ def session(self) -> Session:
         """
         return self._ds.session()
 
+    @staticmethod
+    def _inherit_base_store_params(
+        dataset_or_uri: Union[str, Path, LanceDataset, None],
+        base_store_params: Optional[Dict[str, Dict[str, str]]],
+    ) -> Optional[Dict[str, Dict[str, str]]]:
+        if base_store_params is None and isinstance(dataset_or_uri, LanceDataset):
+            return dataset_or_uri._base_store_params
+        return base_store_params
+
     @staticmethod
     def _commit(
         base_uri: Union[str, Path],
@@ -3931,6 +3953,7 @@ def commit(
         namespace_client: Optional["LanceNamespace"] = None,
         table_id: Optional[List[str]] = None,
         namespace_client_managed_versioning: bool = False,
+        base_store_params: Optional[Dict[str, Dict[str, str]]] = None,
     ) -> LanceDataset:
         """Create a new version of dataset
 
@@ -4001,6 +4024,8 @@ def commit(
         table_id : List[str], optional
             The table identifier within the namespace (e.g., ["workspace", "table"]).
             Must be provided together with namespace_client.
+        base_store_params : dict of str to dict, optional
+            Runtime-only object store parameters keyed by base path URI.
 
         Returns
         -------
@@ -4028,6 +4053,10 @@ def commit(
         2  3  c
         3  4  d
         """
+        base_store_params = LanceDataset._inherit_base_store_params(
+            base_uri, base_store_params
+        )
+
         if isinstance(base_uri, Path):
             base_uri = str(base_uri)
         elif isinstance(base_uri, LanceDataset):
@@ -4101,6 +4130,7 @@ def commit(
 
         ds = LanceDataset.__new__(LanceDataset)
         ds._storage_options = storage_options
+        ds._base_store_params = base_store_params
         ds._namespace_client = namespace_client
         ds._table_id = table_id
         ds._namespace_client_managed_versioning = namespace_client_managed_versioning
@@ -4119,6 +4149,7 @@ def commit_batch(
         enable_v2_manifest_paths: Optional[bool] = None,
         detached: Optional[bool] = False,
         max_retries: int = 20,
+        base_store_params: Optional[Dict[str, Dict[str, str]]] = None,
     ) -> BulkCommitResult:
         """Create a new version of dataset with multiple transactions.
 
@@ -4161,6 +4192,8 @@ def commit_batch(
             the future.
         max_retries : int
             The maximum number of retries to perform when committing the dataset.
+        base_store_params : dict of str to dict, optional
+            Runtime-only object store parameters keyed by base path URI.
 
         Returns
         -------
@@ -4170,6 +4203,10 @@ def commit_batch(
             merged: Transaction
                 The merged transaction that was applied to the dataset.
         """
+        base_store_params = LanceDataset._inherit_base_store_params(
+            dest, base_store_params
+        )
+
         if isinstance(dest, Path):
             dest = str(dest)
         elif isinstance(dest, LanceDataset):
@@ -4198,6 +4235,7 @@ def commit_batch(
         ds._ds = new_ds
         ds._uri = new_ds.uri
         ds._storage_options = storage_options
+        ds._base_store_params = base_store_params
         ds._namespace_client = None
         ds._table_id = None
         ds._default_scan_options = None
@@ -6361,6 +6399,7 @@ def write_dataset(
     transaction_properties: Optional[Dict[str, str]] = None,
     initial_bases: Optional[List[DatasetBasePath]] = None,
     target_bases: Optional[List[str]] = None,
+    base_store_params: Optional[Dict[str, Dict[str, str]]] = None,
     external_blob_mode: Literal["reference", "ingest"] = "reference",
     allow_external_blob_outside_bases: bool = False,
     blob_pack_file_size_threshold: Optional[int] = None,
@@ -6458,6 +6497,12 @@ def write_dataset(
 
         **CREATE mode**: References must match bases in `initial_bases`
         **APPEND/OVERWRITE modes**: References must match bases in the existing manifest
+    base_store_params : dict of str to dict, optional
+        Runtime-only object store parameters keyed by base path URI. Each key
+        is a base path URI (e.g., "s3://bucket/path") and each value is a dict
+        of storage options (credentials, endpoint, etc.) for that base. These
+        are not persisted to the manifest. When a base has no explicit entry
+        here, the top-level ``storage_options`` is used as a fallback.
     external_blob_mode: {"reference", "ingest"}, default "reference"
         How external blob URIs are handled on write.
 
@@ -6578,6 +6623,7 @@ def write_dataset(
     reader = _coerce_reader(data_obj, schema)
     _validate_schema(reader.schema)
     # TODO add support for passing in LanceDataset and LanceScanner here
+    base_store_params = LanceDataset._inherit_base_store_params(uri, base_store_params)
 
     # Merge properties and commit_message with priority to commit_message
     merged_properties = _merge_message_to_properties(
@@ -6598,6 +6644,7 @@ def write_dataset(
         "transaction_properties": merged_properties,
         "initial_bases": initial_bases,
         "target_bases": target_bases,
+        "base_store_params": base_store_params,
         "external_blob_mode": external_blob_mode,
         "allow_external_blob_outside_bases": allow_external_blob_outside_bases,
         "blob_pack_file_size_threshold": blob_pack_file_size_threshold,
@@ -6628,6 +6675,7 @@ def write_dataset(
 
     ds = LanceDataset.__new__(LanceDataset)
     ds._storage_options = storage_options
+    ds._base_store_params = base_store_params
     ds._namespace_client = namespace_client
     ds._table_id = table_id
     ds._namespace_client_managed_versioning = namespace_client_managed_versioning

python/src/dataset.rs

@@ -513,7 +513,7 @@ impl Dataset {
     #[allow(clippy::too_many_arguments)]
     #[allow(deprecated)]
     #[new]
-    #[pyo3(signature=(uri, version=None, block_size=None, index_cache_size=None, metadata_cache_size=None, commit_handler=None, storage_options=None, manifest=None, metadata_cache_size_bytes=None, index_cache_size_bytes=None, read_params=None, session=None, namespace_client=None, table_id=None, namespace_client_managed_versioning=false))]
+    #[pyo3(signature=(uri, version=None, block_size=None, index_cache_size=None, metadata_cache_size=None, commit_handler=None, storage_options=None, manifest=None, metadata_cache_size_bytes=None, index_cache_size_bytes=None, read_params=None, session=None, namespace_client=None, table_id=None, namespace_client_managed_versioning=false, base_store_params=None))]
     fn new(
         py: Python,
         uri: String,
@@ -531,6 +531,7 @@ impl Dataset {
         namespace_client: Option<&Bound<'_, PyAny>>,
         table_id: Option<Vec<String>>,
         namespace_client_managed_versioning: bool,
+        base_store_params: Option<HashMap<String, HashMap<String, String>>>,
     ) -> PyResult<Self> {
         let mut params = ReadParams::default();
         if let Some(metadata_cache_size_bytes) = metadata_cache_size_bytes {
@@ -653,6 +654,17 @@ impl Dataset {
             }
         }
 
+        if let Some(base_store_params) = base_store_params {
+            for (base_path, opts) in base_store_params {
+                let accessor = Arc::new(StorageOptionsAccessor::with_static_options(opts));
+                let store_params = ObjectStoreParams {
+                    storage_options_accessor: Some(accessor),
+                    ..Default::default()
+                };
+                builder = builder.with_base_store_params(base_path, store_params);
+            }
+        }
+
         let dataset = rt().block_on(Some(py), builder.load())?;
 
         match dataset {
@@ -3620,6 +3632,20 @@ pub fn get_write_params(options: &Bound<'_, PyDict>) -> PyResult<Option<WritePar
             p = p.with_target_base_names_or_paths(target_bases_list);
         }
 
+        // Handle base_store_params: per-base storage options keyed by base path URI
+        if let Some(base_store_params) =
+            get_dict_opt::<HashMap<String, HashMap<String, String>>>(options, "base_store_params")?
+        {
+            for (base_path, opts) in base_store_params {
+                let accessor = Arc::new(StorageOptionsAccessor::with_static_options(opts));
+                let store_params = ObjectStoreParams {
+                    storage_options_accessor: Some(accessor),
+                    ..Default::default()
+                };
+                p = p.with_base_store_params(base_path, store_params);
+            }
+        }
+
         if let Some(allow_external) =
             get_dict_opt::<bool>(options, "allow_external_blob_outside_bases")?
         {