fix: torchdataset for lance does not support S3 (#4045)

Signed-off-by: jukejian <jukejian@bytedance.com>

Author: Jay-ju

python/python/lance/torch/data.py

@@ -11,7 +11,7 @@
 import math
 import warnings
 from pathlib import Path
-from typing import Callable, Dict, Iterable, List, Literal, Optional, Union
+from typing import Any, Callable, Dict, Iterable, List, Literal, Optional, Union
 
 import pyarrow as pa
 
@@ -182,6 +182,7 @@ def __init__(
         dataset: Union[torch.utils.data.Dataset, str, Path],
         batch_size: int,
         *args,
+        dataset_options: Optional[Dict[str, Any]] = None,
         columns: Optional[Union[List[str], Dict[str, str]]] = None,
         filter: Optional[str] = None,
         samples: Optional[int] = 0,
@@ -237,7 +238,8 @@ def __init__(
         """
         super().__init__()
         if isinstance(dataset, (str, Path)):
-            dataset = lance.dataset(dataset)
+            dataset_options = dataset_options or {}
+            dataset = lance.dataset(dataset, **dataset_options)
         self.dataset = dataset
         self.columns = columns
         self.batch_size = batch_size
@@ -378,16 +380,18 @@ def _blob_columns(self) -> List[str]:
 
 
 class SafeLanceDataset(torch.utils.data.Dataset):
-    def __init__(self, uri):
+    def __init__(self, uri, *, dataset_options=None, **kwargs):
+        super().__init__(**kwargs)
         self.uri = uri
+        self.dataset_options = dataset_options or {}
         self._len = self._safe_preload()
-        self._ds = None  # Deferred initialization
+        self._ds = None
 
     def _safe_preload(self):
         """Main-process safe metadata loading"""
-        ds = lance.dataset(self.uri)
+        ds = lance.dataset(self.uri, **self.dataset_options)
         length = ds.count_rows()
-        del ds  # Critical: release before spawning
+        del ds
         return length
 
     def __len__(self):

python/python/tests/torch_tests/test_data.py

@@ -292,10 +292,11 @@ def test_blob_api(tmp_path: Path):
     )
     tbl = pa.Table.from_arrays([ints, vals], schema=schema)
 
-    ds = lance.write_dataset(tbl, tmp_path / "data.lance")
+    uri = tmp_path / "data.lance"
+    dataset = lance.write_dataset(tbl, uri)
+
     torch_ds = LanceDataset(
-        ds,
-        batch_size=4,
+        uri, batch_size=4, dataset_options={"version": dataset.version}
     )
     with pytest.raises(NotImplementedError):
         next(iter(torch_ds))
@@ -314,7 +315,7 @@ def to_tensor_fn(batch, *args, **kwargs):
         return {"int": ints, "val": vals}
 
     torch_ds = LanceDataset(
-        ds,
+        dataset,
         batch_size=4,
         to_tensor_fn=to_tensor_fn,
     )