Merge remote-tracking branch 'refs/remotes/origin/main'

Author: Volker L.

pydala/catalog2.py

@@ -41,8 +41,7 @@ class Table:
     write_args: dict = field(default_factory=dict)
 
     def _load_parquet(self):
-
-
+        pass
 
 
 @dataclass

pydala/dataset.py

@@ -66,7 +66,7 @@ def __init__(
         )
         self._timestamp_column = timestamp_column
 
-        self.load_files()
+        #self.load_files()
 
         if self.has_files:
             if partitioning == "ignore":
@@ -732,7 +732,7 @@ def write_to_dataset(
             self.delete_files(del_files)
 
         self.clear_cache()
-        self.load_files()
+        #self.load_files()
 
 
 class ParquetDataset(PydalaDatasetMetadata, BaseDataset):

pydala/filesystem.py

@@ -1,6 +1,7 @@
 import datetime as dt
 import inspect
 import os
+import asyncio
 from datetime import datetime, timedelta
 from functools import wraps
 from pathlib import Path
@@ -13,6 +14,7 @@
 import pyarrow as pa
 import pyarrow.dataset as pds
 import pyarrow.parquet as pq
+import pyarrow.fs as pfs
 import s3fs
 from fsspec import AbstractFileSystem, filesystem
 from fsspec.implementations.cache_mapper import AbstractCacheMapper
@@ -26,6 +28,17 @@
 from .schema import shrink_large_string
 
 
+def get_credentials_from_fssspec(fs: AbstractFileSystem) -> dict[str, str]:
+    if "s3" in fs.protocol:
+        credendials = fs.s3._get_credentials()
+        return {
+            "access_key": credendials.access_key,
+            "secret_key": credendials.secret_key,
+            "session_token": credendials.token,
+            "endpoint_override": fs.s3._endpoint.host,
+        }
+
+
 def get_total_directory_size(directory: str):
     return sum(f.stat().st_size for f in Path(directory).glob("**/*") if f.is_file())
 
@@ -733,6 +746,12 @@ def sync_folder(
         self.cp(new_src, dst)
 
 
+def list_files_recursive(self, path:str, format:str=""):
+    bucket, prefix = path.split("/", maxsplit=1)
+    return [f["Key"] for f in asyncio.run(self.s3.list_objects_v2(Bucket=bucket, Prefix=prefix))["Contents"] if f["Key"].endswith(format)]
+
+
+
 AbstractFileSystem.read_parquet = read_parquet
 AbstractFileSystem.read_parquet_dataset = read_parquet_dataset
 AbstractFileSystem.write_parquet = write_parquet
@@ -763,6 +782,7 @@ def sync_folder(
 # AbstractFileSystem.parallel_mv = parallel_mv
 # AbstractFileSystem.parallel_rm = parallel_rm
 AbstractFileSystem.sync_folder = sync_folder
+AbstractFileSystem.list_files_recursive = list_files_recursive
 
 
 def FileSystem(
@@ -828,9 +848,104 @@ def FileSystem(
             same_names=same_names,
             **kwargs,
         )
+
+    return fs
+
+
+def PyArrowFileSystem(
+    bucket: str | None = None,
+    fs: AbstractFileSystem | None = None,
+    access_key: str | None = None,
+    secret_key: str | None = None,
+    session_token: str | None = None,
+    endpoint_override: str | None = None,
+    protocol: str | None = None,
+) -> pfs.FileSystem:
+    credentials = None
+    if fs is not None:
+        protocol = fs.protocol[0] if isinstance(fs.protocol, tuple) else fs.protocol
+
+        if protocol == "dir":
+            bucket = fs.path
+            fs = fs.fs
+            protocol = fs.protocol[0] if isinstance(fs.protocol, tuple) else fs.protocol
+
+        if protocol == "s3":
+            credentials = get_credentials_from_fssspec(fs)
+
+    if credentials is None:
+        credentials = {
+            "access_key": access_key,
+            "secret_key": secret_key,
+            "session_token": session_token,
+            "endpoint_override": endpoint_override,
+        }
+    if protocol == "s3":
+        fs = pfs.S3FileSystem(
+            **credentials,
+        )
+    elif protocol in ("file", "local", None):
+        fs = pfs.LocalFileSystem()
+
+    else:
+        fs = pfs.LocalFileSystem()
+
+    if bucket is not None:
+        if protocol in ["file", "local", "None"]:
+            bucket = os.path.abspath(bucket)
+
+        fs = pfs.SubTreeFileSystem(base_fs=fs, base_path=bucket)
+
     return fs
 
 
+# class FileSystem:
+#     def __init__(
+#         self,
+#         bucket: str | None = None,
+#         fs: AbstractFileSystem | None = None,
+#         profile: str | None = None,
+#         key: str | None = None,
+#         endpoint_url: str | None = None,
+#         secret: str | None = None,
+#         token: str | None = None,
+#         protocol: str | None = None,
+#         cached: bool = False,
+#         cache_storage="~/.tmp",
+#         check_files: bool = False,
+#         cache_check: int = 120,
+#         expire_time: int = 24 * 60 * 60,
+#         same_names: bool = False,
+#         **kwargs,
+#     ):
+#         self._fsspec_fs = FsSpecFileSystem(
+#             bucket=bucket,
+#             fs=fs,
+#             profile=profile,
+#             key=key,
+#             endpoint_url=endpoint_url,
+#             secret=secret,
+#             token=token,
+#             protocol=protocol,
+#             cached=cached,
+#             cache_storage=cache_storage,
+#             check_files=check_files,
+#             cache_check=cache_check,
+#             expire_time=expire_time,
+#             same_names=same_names,
+#             **kwargs,
+#         )
+#         self._pfs_fs = PyArrowFileSystem(
+#             bucket=bucket,
+#             fs=fs,
+#             access_key=key,
+#             secret_key=secret,
+#             session_token=token,
+#             endpoint_override=endpoint_url,
+#             protocol=protocol,
+#         )
+
+
 def clear_cache(fs: AbstractFileSystem | None):
     if hasattr(fs, "dir_cache"):
         if fs is not None:

pydala/helpers/_fsspec_util.py

@@ -106,7 +106,7 @@ def _check_file(self, path):
                 return fn
             logger.info(f"Downloading {self.protocol[0]}://{path}")
 
-    #def glob(self, path):
+    # def glob(self, path):
     #    return [self._strip_protocol(path)]
 
     def size(self, path):

pydala/helpers/datetime.py

@@ -86,8 +86,8 @@ def timestamp_from_string(
         datetime.datetime: The datetime object.
     """
     # Extract the timezone from the string if not provided
-    #tz = extract_timezone(timestamp) if tz is None else tz
-    #timestamp = timestamp.replace(tz, "").strip() if tz else timestamp
+    # tz = extract_timezone(timestamp) if tz is None else tz
+    # timestamp = timestamp.replace(tz, "").strip() if tz else timestamp
 
     pdl_timestamp = pdl.parse(timestamp, exact=exact, strict=strict)
 

pydala/metadata.py

@@ -107,6 +107,40 @@ def remove_from_metadata(
 
     return metadata
 
+def get_file_paths(metadata: pq.FileMetaData,) -> list[str]:
+    return [metadata.row_group(i).column(0).file_path for i in range(metadata.num_row_groups)]
+
+class FileMetadata:
+    def __init__(
+        self,
+        path:str,
+        filesystem: AbstractFileSystem | pfs.FileSystem | None = None,
+        bucket: str | None = None,
+        cached: bool = False,
+        **caching_options,
+        **kwargs,
+    ):
+        self._path = path
+        self._bucket = bucket
+        self._cached = cached
+        self._base_filesystem = filesystem
+        self._filesystem = FileSystem(
+            bucket=bucket, fs=filesystem, cached=cached, **caching_options
+        )
+
+        self._caching_options = caching_options
+
+        self.load_files()
+
+    def load_files(self):
+        self._files = self._filesystem.list_files_recursive(self._path)
+
+
+
+
+    @property
+    def fs(self):
+        return self._filesystem
 
 class ParquetDatasetMetadata:
     def __init__(
@@ -140,7 +174,7 @@ def __init__(
         )
 
         self._makedirs()
-        self.load_files()
+        #self.load_files()
 
         self._caching_options = caching_options
 
@@ -170,7 +204,10 @@ def _makedirs(self):
             self._filesystem.touch(os.path.join(self._path, "tmp.delete"))
             self._filesystem.rm(os.path.join(self._path, "tmp.delete"))
 
-    def load_files(self) -> None:
+    def load_files(self)->None:
+        self._files = get_file_paths(self._metadata)
+
+    def _ls_files(self) -> None:
         """
         Reloads the list of files in the dataset directory. This method should be called
         after adding or removing files from the directory to ensure that the dataset object
@@ -180,7 +217,7 @@ def load_files(self) -> None:
             None
         """
         self.clear_cache()
-        self._files = [
+        return [
             fn.replace(self._path, "").lstrip("/")
             for fn in sorted(
                 self._filesystem.glob(os.path.join(self._path, "**/*.parquet"))
@@ -200,7 +237,7 @@ def _collect_file_metadata(self, files: list[str] | None = None, **kwargs) -> No
             None
         """
         if files is None:
-            files = self._files
+            files = self._ls_files()
 
         file_metadata = collect_parquet_metadata(
             files=files,
@@ -211,8 +248,8 @@ def _collect_file_metadata(self, files: list[str] | None = None, **kwargs) -> No
 
         # if file_metadata:
         for f in file_metadata:
-            file_metadata[f.replace(self._path, "")].set_file_path(
-                f.split(self._path)[-1].lstrip("/")
+            file_metadata[f].set_file_path(
+                f
             )
 
         if self.has_file_metadata:
@@ -250,14 +287,14 @@ def update_file_metadata(self, files: list[str] | None = None, **kwargs) -> None
         """
 
         # Add new files to file_metadata
-        self.load_files()
+        all_files = self._ls_files()
 
         new_files = []
         rm_files = []
 
         if self.has_file_metadata:
-            new_files += sorted(set(self.files) - set(self._file_metadata.keys()))
-            rm_files += sorted(set(self._file_metadata.keys()) - set(self.files))
+            new_files += sorted(set(all_files) - set(self._file_metadata.keys()))
+            rm_files += sorted(set(self._file_metadata.keys()) - set(all_files))
 
         else:
             new_files += sorted(set(new_files + self._files))

pyproject.toml

@@ -56,8 +56,8 @@ dev-dependencies = [
     "datafusion>=42.0.0",
     "ibis>=3.3.0",
     "ibis-framework[duckdb,polars]>=9.5.0",
-    "obstore>=0.2.0",
-    "boto3>=1.35.49",
+    "ruff>=0.7.1",
+    "adlfs>=2024.7.0",
 ]
 managed = true