simplify

python/deltalake/_internal.pyi

@@ -205,7 +205,7 @@ class ObjectOutputStream:
 class DeltaFileSystemHandler:
     """Implementation of pyarrow.fs.FileSystemHandler for use with pyarrow.fs.PyFileSystem"""
 
-    def __init__(self, root: str, options: dict[str, str] | None = None) -> None: ...
+    def __init__(self, root: str, options: dict[str, str] | None = None, known_sizes: dict[str, int] | None = None) -> None: ...
     def get_type_name(self) -> str: ...
     def copy_file(self, src: str, dst: str) -> None:
         """Copy a file.

python/deltalake/fs.py

@@ -11,48 +11,24 @@ class DeltaStorageHandler(DeltaFileSystemHandler, FileSystemHandler):
     """
     DeltaStorageHandler is a concrete implementations of a PyArrow FileSystemHandler.
     """
-
-    known_sizes: Dict[str, int] = {}
-
-    def __new__(  # type:ignore
-        cls,
-        table_uri: str,
-        storage_options: Optional[Dict[str, str]] = None,
-        known_sizes: Optional[Dict[str, int]] = None,
-    ):
-        return super().__new__(
-            cls, table_uri=table_uri, options=storage_options  # type:ignore
-        )
-
-    def __init__(
-        self,
-        table_uri: str,
-        storage_options: Optional[Dict[str, str]] = None,
-        known_sizes: Optional[Dict[str, int]] = None,
-    ):
-        if known_sizes:
-            self.known_sizes = known_sizes
-        return
-
-    def open_input_file(self, path: str, size: Optional[int] = None) -> pa.PythonFile:
+
+    def open_input_file(self, path: str) -> pa.PythonFile:
         """
         Open an input file for random access reading.
 
         :param source: The source to open for reading.
         :return:  NativeFile
         """
-        size = self.known_sizes.get(path)
-        return pa.PythonFile(DeltaFileSystemHandler.open_input_file(self, path, size))
+        return pa.PythonFile(DeltaFileSystemHandler.open_input_file(self, path))
 
-    def open_input_stream(self, path: str, size: Optional[int] = None) -> pa.PythonFile:
+    def open_input_stream(self, path: str) -> pa.PythonFile:
         """
         Open an input stream for sequential reading.
 
         :param source: The source to open for reading.
         :return:  NativeFile
         """
-        size = self.known_sizes.get(path)
-        return pa.PythonFile(DeltaFileSystemHandler.open_input_file(self, path, size))
+        return pa.PythonFile(DeltaFileSystemHandler.open_input_file(self, path))
 
     def open_output_stream(
         self, path: str, metadata: Optional[Dict[str, str]] = None

python/src/filesystem.rs

@@ -25,6 +25,7 @@ pub struct DeltaFileSystemHandler {
     pub(crate) inner: Arc<DynObjectStore>,
     pub(crate) rt: Arc<Runtime>,
     pub(crate) config: FsConfig,
+    pub(crate) known_sizes: Option<HashMap<String, i64>>,
 }
 
 impl DeltaFileSystemHandler {
@@ -41,8 +42,12 @@ impl DeltaFileSystemHandler {
 #[pymethods]
 impl DeltaFileSystemHandler {
     #[new]
-    #[pyo3(signature = (table_uri, options = None))]
-    fn new(table_uri: &str, options: Option<HashMap<String, String>>) -> PyResult<Self> {
+    #[pyo3(signature = (table_uri, options = None, known_sizes = None))]
+    fn new(
+        table_uri: &str,
+        options: Option<HashMap<String, String>>,
+        known_sizes: Option<HashMap<String, i64>>,
+    ) -> PyResult<Self> {
         let storage = DeltaTableBuilder::from_uri(table_uri)
             .with_storage_options(options.clone().unwrap_or_default())
             .build_storage()
@@ -54,6 +59,7 @@ impl DeltaFileSystemHandler {
                 root_url: table_uri.into(),
                 options: options.unwrap_or_default(),
             },
+            known_sizes: known_sizes,
         })
     }
 
@@ -236,15 +242,20 @@ impl DeltaFileSystemHandler {
         Ok(())
     }
 
-    fn open_input_file(&self, path: String, size: Option<i64>) -> PyResult<ObjectInputFile> {
+    fn open_input_file(&self, path: String) -> PyResult<ObjectInputFile> {
+        let size = match &self.known_sizes {
+            Some(sz) => sz.get(&path),
+            None => None,
+        };
+
         let path = Self::parse_path(&path);
         let file = self
             .rt
             .block_on(ObjectInputFile::try_new(
                 Arc::clone(&self.rt),
                 self.inner.clone(),
                 path,
-                size,
+                size.copied(),
             ))
             .map_err(PythonError::from)?;
         Ok(file)

python/src/lib.rs

@@ -562,6 +562,7 @@ impl RawDeltaTable {
             inner: self._table.object_store(),
             rt: Arc::new(rt()?),
             config: self._config.clone(),
+            known_sizes: None
         })
     }