feat: FK Models (#584)

Co-authored-by: Antoine Bon <[email protected]>

Author: lukaszkolodziejczyk

mostlyai/sdk/_data/file/base.py

@@ -18,7 +18,7 @@
 import re
 import time
 from abc import abstractmethod
-from collections.abc import Generator, Iterable
+from collections.abc import Generator, Iterable, Iterator
 from enum import Enum
 from pathlib import Path
 from typing import Any
@@ -375,6 +375,17 @@ def handle_if_exists(self, if_exists: str = "fail") -> str:
                 return "a"
         return "w"
 
+    def iter_partitions(self) -> Iterator[tuple[int, Path, pd.DataFrame]]:
+        """Iterate over dataset partitions yielding (index, file_path, dataframe)."""
+        for idx, file_path in enumerate(self.dataset.files):
+            data = pd.read_parquet(file_path)
+            yield idx, Path(file_path), data
+
+    @property
+    def files(self) -> list[Path]:
+        """Get the list of partition files."""
+        return [Path(f) for f in self.dataset.files]
+
 
 class FileContainer(DataContainer):
     SCHEMES = ["http", "https"]