retrievalbase 2.2.0__tar.gz → 2.3.0__tar.gz

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/CHANGELOG.md

@@ -1,3 +1,10 @@
+# [2.3.0](https://gitlab.com/efysent/agentic-core/retrievalbase/compare/v2.2.0...v2.3.0) (2026-05-24)
+
+
+### Features
+
+* add recursive dataset connector loading ([3e7e46f](https://gitlab.com/efysent/agentic-core/retrievalbase/commit/3e7e46f6d8533612cb3e47375985aeb3ebf156ff))
+
 # [2.2.0](https://gitlab.com/efysent/agentic-core/retrievalbase/compare/v2.1.3...v2.2.0) (2026-05-22)
 
 

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: retrievalbase
-Version: 2.2.0
+Version: 2.3.0
 Author-email: jalal <jalalkhaldi3@gmail.com>
 Requires-Python: <3.13,>=3.11
 Requires-Dist: faiss-cpu<2.0.0,>=1.13.2

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "retrievalbase"
-version = "2.2.0"
+version = "2.3.0"
 description = ""
 authors = [
     { name = "jalal", email = "jalalkhaldi3@gmail.com" }

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/src/retrievalbase/connector/__init__.py

@@ -29,6 +29,10 @@ class DatasetConnector[TCDatasetConnector: DatasetConnectorSettings](
     def _load(self) -> pl.DataFrame | pl.LazyFrame:
         raise NotImplementedError
 
+    @abstractmethod
+    def _load_recursive(self, *paths: str) -> pl.DataFrame | pl.LazyFrame:
+        raise NotImplementedError
+
     @abstractmethod
     def to(self, ds: "Dataset[Any]") -> None:
         raise NotImplementedError
@@ -57,6 +61,26 @@ class DatasetConnector[TCDatasetConnector: DatasetConnectorSettings](
 
         return PolarsTextDataset.from_polars(df)
 
+    def load_recursive(self, *paths: str) -> "Dataset[pl.DataFrame | pl.LazyFrame]":
+        from retrievalbase.dataset.polars import PolarsDataset
+
+        _logger.info(f"Loading recursive dataset | connector={self.__class__.__name__}")
+
+        df = self._load_recursive(*paths)
+        self._log_polars_info(df)
+
+        return PolarsDataset.from_polars(df)
+
+    def load_recursive_text(self, *paths: str) -> "TextDataset[pl.DataFrame | pl.LazyFrame]":
+        from retrievalbase.dataset.polars import PolarsTextDataset
+
+        _logger.info(f"Loading recursive text dataset | connector={self.__class__.__name__}")
+
+        df = self._load_recursive(*paths)
+        self._log_polars_info(df)
+
+        return PolarsTextDataset.from_polars(df)
+
     # ------------------------------------------------------------------
     # Helpers
     # ------------------------------------------------------------------

retrievalbase-2.3.0/src/retrievalbase/connector/minio.py

@@ -0,0 +1,96 @@
+import io
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
+
+import polars as pl
+from minio import Minio
+from minio.error import S3Error
+
+from retrievalbase.connector import DatasetConnector
+from retrievalbase.connector.settings import MinioDatasetConnectorSettings
+from retrievalbase.exceptions import MinioParquetObjectsNotFoundError
+
+if TYPE_CHECKING:
+    from retrievalbase.dataset import Dataset
+
+
+class MinioDatasetConnector(DatasetConnector[MinioDatasetConnectorSettings]):
+    def __init__(self, config: MinioDatasetConnectorSettings):
+        super().__init__(config)
+        self.client = Minio(
+            self.config.endpoint.replace("http://", "").replace("https://", ""),
+            access_key=self.config.access_key.get_secret_value(),
+            secret_key=self.config.secret_key.get_secret_value(),
+            secure=self.config.endpoint.startswith("https://"),
+        )
+
+    def _load(self) -> pl.DataFrame | pl.LazyFrame:
+        return self._read_parquet_object(self.config.bucket, self.config.key)
+
+    def _load_recursive(self, *paths: str) -> pl.DataFrame:
+        dataframes = [
+            self._read_parquet_object(bucket, object_name)
+            for bucket, object_name in self._iter_parquet_objects(paths or (self.config.key,))
+        ]
+        if not dataframes:
+            raise MinioParquetObjectsNotFoundError(paths or (self.config.key,))
+        return pl.concat(dataframes)
+
+    def _read_parquet_object(self, bucket: str, object_name: str) -> pl.DataFrame:
+        response = self.client.get_object(bucket, object_name)
+        try:
+            buffer = io.BytesIO(response.read())
+        finally:
+            response.close()
+            response.release_conn()
+        return pl.read_parquet(buffer)
+
+    def _iter_parquet_objects(self, paths: tuple[str, ...]) -> list[tuple[str, str]]:
+        objects: list[tuple[str, str]] = []
+        for path in paths:
+            bucket, object_path = self._resolve_path(path)
+            if object_path.endswith(".parquet"):
+                objects.append((bucket, object_path))
+                continue
+            prefix = object_path.rstrip("/")
+            if prefix:
+                prefix = f"{prefix}/"
+            objects.extend(
+                (bucket, item.object_name)
+                for item in self.client.list_objects(bucket, prefix=prefix, recursive=True)
+                if item.object_name.endswith(".parquet")
+            )
+        return sorted(objects)
+
+    def _resolve_path(self, path: str) -> tuple[str, str]:
+        parsed = urlparse(path)
+        if parsed.scheme in {"s3", "minio"}:
+            return parsed.netloc, parsed.path.lstrip("/")
+
+        normalized = path.lstrip("/")
+        bucket, separator, key = normalized.partition("/")
+        if separator and bucket == self.config.bucket:
+            return bucket, key
+        return self.config.bucket, normalized
+
+    def to(self, ds: "Dataset[Any]") -> None:
+        df = ds.polars
+        buffer = io.BytesIO()
+        df.write_parquet(buffer)
+        buffer.seek(0)
+        self.client.put_object(
+            bucket_name=self.config.bucket,
+            object_name=self.config.key,
+            data=buffer,
+            length=buffer.getbuffer().nbytes,
+            content_type="application/octet-stream",
+        )
+
+    def target_exists(self, target: str) -> bool:
+        try:
+            self.client.stat_object(self.config.bucket, target)
+        except S3Error as error:
+            if error.code in {"NoSuchBucket", "NoSuchKey"}:
+                return False
+            raise
+        return True

retrievalbase-2.3.0/src/retrievalbase/connector/parquet.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+
+from retrievalbase.connector import DatasetConnector
+from retrievalbase.connector.settings import ParquetDatasetConnectorSettings
+from retrievalbase.exceptions import ParquetFilesNotFoundError
+
+if TYPE_CHECKING:
+    from retrievalbase.dataset import Dataset
+
+
+class ParquetDatasetConnector(DatasetConnector[ParquetDatasetConnectorSettings]):
+    def __init__(self, config: ParquetDatasetConnectorSettings):
+        super().__init__(config)
+
+    def _load(self) -> pl.DataFrame | pl.LazyFrame:
+        return pl.scan_parquet(self.config.path) if self.config.lazy else pl.read_parquet(self.config.path)
+
+    def _load_recursive(self, *paths: str) -> pl.DataFrame | pl.LazyFrame:
+        parquet_paths = self._iter_parquet_paths(paths or (self.config.path,))
+        if not parquet_paths:
+            raise ParquetFilesNotFoundError(paths or (self.config.path,))
+        if self.config.lazy:
+            return pl.scan_parquet([str(path) for path in parquet_paths])
+        return pl.concat([pl.read_parquet(path) for path in parquet_paths])
+
+    def _iter_parquet_paths(self, paths: tuple[str, ...]) -> list[Path]:
+        parquet_paths: list[Path] = []
+        for raw_path in paths:
+            path = Path(raw_path)
+            if path.suffix == ".parquet":
+                parquet_paths.append(path)
+                continue
+            parquet_paths.extend(path.rglob("*.parquet"))
+        return sorted(parquet_paths)
+
+    def to(self, ds: "Dataset[Any]") -> None:
+        ds.polars.write_parquet(self.config.path)
+
+    def target_exists(self, target: str) -> bool:
+        return Path(target).exists()

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/src/retrievalbase/exceptions.py

@@ -37,6 +37,28 @@ class DatasetError(RetrievalBaseError):
     pass
 
 
+class DatasetConnectorError(RetrievalBaseError):
+    """Base error for dataset connector failures."""
+
+
+class MinioParquetObjectsNotFoundError(DatasetConnectorError, FileNotFoundError):
+    """Raised when a MinIO path does not contain parquet objects."""
+
+    def __init__(self, paths: tuple[str, ...]):
+        self.paths = paths
+        roots = ", ".join(paths)
+        super().__init__(f"No parquet objects found in MinIO paths: {roots}")
+
+
+class ParquetFilesNotFoundError(DatasetConnectorError, FileNotFoundError):
+    """Raised when local paths do not contain parquet files."""
+
+    def __init__(self, paths: tuple[str, ...]):
+        self.paths = paths
+        roots = ", ".join(paths)
+        super().__init__(f"No parquet files found in paths: {roots}")
+
+
 class DatasetSplitError(DatasetError, ValueError):
     """Raised when dataset split parameters are invalid."""
 

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/tests/fixtures/components.py

@@ -76,6 +76,9 @@ class FakeDatasetConnector(DatasetConnector[FakeDatasetConnectorSettings]):
     def _load(self) -> pl.DataFrame:
         return pl.DataFrame(self.config.rows)
 
+    def _load_recursive(self, *paths: str) -> pl.DataFrame:
+        return self._load()
+
     def to(self, ds: Any) -> None:
         self.__class__.last_written = ds.polars
 

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/tests/unit/test_connector/test_connectors.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+from types import SimpleNamespace
 from typing import TypedDict, cast
 
 import polars as pl
@@ -17,6 +18,7 @@ from retrievalbase.connector.settings import (
 )
 from retrievalbase.dataset import Dataset, TextDataset
 from retrievalbase.dataset.polars import PolarsDataset, PolarsTextDataset
+from retrievalbase.exceptions import MinioParquetObjectsNotFoundError, ParquetFilesNotFoundError
 from tests.fixtures.data import make_text_dataframe
 
 
@@ -96,6 +98,36 @@ def test_parquet_connector_reports_target_existence(tmp_path) -> None:
     assert connector.target_exists(str(missing_path)) is False
 
 
+def test_parquet_connector_loads_recursive_paths(tmp_path) -> None:
+    root = tmp_path / "root"
+    nested = root / "nested"
+    nested.mkdir(parents=True)
+    first = make_text_dataframe([{"page_content": "one", "metadata": {"doc_id": "a"}}])
+    second = make_text_dataframe([{"page_content": "two", "metadata": {"doc_id": "b"}}])
+    ignored = root / "ignored.txt"
+    first.write_parquet(root / "first.parquet")
+    second.write_parquet(nested / "second.parquet")
+    ignored.write_text("ignore me")
+
+    connector = ParquetDatasetConnector(ParquetDatasetConnectorSettings(module_path="x", path=str(root), lazy=False))
+
+    loaded = connector.load_recursive_text()
+
+    assert loaded.polars.to_dict(as_series=False) == {
+        "page_content": ["one", "two"],
+        "metadata": [{"doc_id": "a"}, {"doc_id": "b"}],
+    }
+
+
+def test_parquet_connector_raises_when_recursive_path_has_no_parquet(tmp_path) -> None:
+    empty = tmp_path / "empty"
+    empty.mkdir()
+    connector = ParquetDatasetConnector(ParquetDatasetConnectorSettings(module_path="x", path=str(empty), lazy=False))
+
+    with pytest.raises(ParquetFilesNotFoundError, match="No parquet files found"):
+        connector.load_recursive()
+
+
 def test_minio_connector_reads_and_writes_parquet_payloads(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -207,3 +239,94 @@ def test_minio_connector_reports_target_existence(monkeypatch: pytest.MonkeyPatc
     assert connector.target_exists("sample.parquet") is True
     assert connector.target_exists("missing.parquet") is False
     assert calls == [("datasets", "sample.parquet"), ("datasets", "missing.parquet")]
+
+
+def test_minio_connector_loads_recursive_parquet_prefixes(monkeypatch: pytest.MonkeyPatch) -> None:
+    frames = {
+        "a/df.parquet": pl.DataFrame({"id": [1], "text": ["one"]}),
+        "a/B/df2.parquet": pl.DataFrame({"id": [2], "text": ["two"]}),
+        "another_root/df3.parquet": pl.DataFrame({"id": [3], "text": ["three"]}),
+        "another_root/ignored.txt": pl.DataFrame({"id": [99], "text": ["ignored"]}),
+    }
+    calls: dict[str, list[object]] = {"listed": [], "read": []}
+
+    def parquet_payload(df: pl.DataFrame) -> bytes:
+        buffer = io.BytesIO()
+        df.write_parquet(buffer)
+        return buffer.getvalue()
+
+    class FakeResponse:
+        def __init__(self, payload: bytes) -> None:
+            self.payload = payload
+
+        def read(self) -> bytes:
+            return self.payload
+
+        def close(self) -> None:
+            pass
+
+        def release_conn(self) -> None:
+            pass
+
+    class FakeMinio:
+        def __init__(self, endpoint: str, access_key: str, secret_key: str, secure: bool) -> None:
+            pass
+
+        def list_objects(self, bucket: str, prefix: str, recursive: bool) -> list[SimpleNamespace]:
+            calls["listed"].append((bucket, prefix, recursive))
+            return [SimpleNamespace(object_name=key) for key in frames if key.startswith(prefix)]
+
+        def get_object(self, bucket: str, key: str) -> FakeResponse:
+            calls["read"].append((bucket, key))
+            return FakeResponse(parquet_payload(frames[key]))
+
+    monkeypatch.setattr("retrievalbase.connector.minio.Minio", FakeMinio)
+
+    connector = MinioDatasetConnector(
+        MinioDatasetConnectorSettings(
+            module_path="x",
+            endpoint="https://minio.local",
+            bucket="datasets",
+            key="a",
+            access_key=SecretStr("access"),
+            secret_key=SecretStr("secret"),
+        )
+    )
+
+    loaded = connector.load_recursive("datasets/a", "s3://datasets/another_root")
+
+    assert calls["listed"] == [("datasets", "a/", True), ("datasets", "another_root/", True)]
+    assert calls["read"] == [
+        ("datasets", "a/B/df2.parquet"),
+        ("datasets", "a/df.parquet"),
+        ("datasets", "another_root/df3.parquet"),
+    ]
+    assert loaded.polars.to_dict(as_series=False) == {
+        "id": [2, 1, 3],
+        "text": ["two", "one", "three"],
+    }
+
+
+def test_minio_connector_raises_when_prefix_has_no_parquet(monkeypatch: pytest.MonkeyPatch) -> None:
+    class FakeMinio:
+        def __init__(self, endpoint: str, access_key: str, secret_key: str, secure: bool) -> None:
+            pass
+
+        def list_objects(self, bucket: str, prefix: str, recursive: bool) -> list[SimpleNamespace]:
+            return [SimpleNamespace(object_name="empty/readme.txt")]
+
+    monkeypatch.setattr("retrievalbase.connector.minio.Minio", FakeMinio)
+
+    connector = MinioDatasetConnector(
+        MinioDatasetConnectorSettings(
+            module_path="x",
+            endpoint="https://minio.local",
+            bucket="datasets",
+            key="empty",
+            access_key=SecretStr("access"),
+            secret_key=SecretStr("secret"),
+        )
+    )
+
+    with pytest.raises(MinioParquetObjectsNotFoundError, match="No parquet objects found"):
+        connector.load_recursive()

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/tests/unit/test_evaluation/test_base_contracts.py

@@ -30,6 +30,9 @@ class RaisingConnector(DatasetConnector[DatasetConnectorSettings]):
     def _load(self) -> pl.DataFrame | pl.LazyFrame:
         return super()._load()
 
+    def _load_recursive(self, *paths: str) -> pl.DataFrame | pl.LazyFrame:
+        return super()._load_recursive(*paths)
+
     def to(self, ds) -> None:
         super().to(ds)
 
@@ -95,6 +98,8 @@ def test_base_contracts_raise_not_implemented_and_noops() -> None:
 
     with pytest.raises(NotImplementedError):
         connector._load()
+    with pytest.raises(NotImplementedError):
+        connector._load_recursive("x")
     with pytest.raises(NotImplementedError):
         connector.to(PolarsTextDataset.from_records([("x", {"doc_id": "1"})]))
     with pytest.raises(NotImplementedError):

{retrievalbase-2.2.0 → retrievalbase-2.3.0}/uv.lock

@@ -2584,7 +2584,7 @@ wheels = [
 
 [[package]]
 name = "retrievalbase"
-version = "2.2.0"
+version = "2.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "faiss-cpu" },

retrievalbase-2.2.0/src/retrievalbase/connector/minio.py

@@ -1,55 +0,0 @@
-import io
-from typing import TYPE_CHECKING, Any
-
-import polars as pl
-from minio import Minio
-from minio.error import S3Error
-
-from retrievalbase.connector import DatasetConnector
-from retrievalbase.connector.settings import MinioDatasetConnectorSettings
-
-if TYPE_CHECKING:
-    from retrievalbase.dataset import Dataset
-
-
-class MinioDatasetConnector(DatasetConnector[MinioDatasetConnectorSettings]):
-    def __init__(self, config: MinioDatasetConnectorSettings):
-        super().__init__(config)
-        self.client = Minio(
-            self.config.endpoint.replace("http://", "").replace("https://", ""),
-            access_key=self.config.access_key.get_secret_value(),
-            secret_key=self.config.secret_key.get_secret_value(),
-            secure=self.config.endpoint.startswith("https://"),
-        )
-
-    def _load(self) -> pl.DataFrame | pl.LazyFrame:
-        response = self.client.get_object(self.config.bucket, self.config.key)
-        try:
-            buffer = io.BytesIO(response.read())
-        finally:
-            response.close()
-            response.release_conn()
-        df = pl.read_parquet(buffer)
-        return df
-
-    def to(self, ds: "Dataset[Any]") -> None:
-        df = ds.polars
-        buffer = io.BytesIO()
-        df.write_parquet(buffer)
-        buffer.seek(0)
-        self.client.put_object(
-            bucket_name=self.config.bucket,
-            object_name=self.config.key,
-            data=buffer,
-            length=buffer.getbuffer().nbytes,
-            content_type="application/octet-stream",
-        )
-
-    def target_exists(self, target: str) -> bool:
-        try:
-            self.client.stat_object(self.config.bucket, target)
-        except S3Error as error:
-            if error.code in {"NoSuchBucket", "NoSuchKey"}:
-                return False
-            raise
-        return True

retrievalbase-2.2.0/src/retrievalbase/connector/parquet.py

@@ -1,24 +0,0 @@
-from pathlib import Path
-from typing import TYPE_CHECKING, Any
-
-import polars as pl
-
-from retrievalbase.connector import DatasetConnector
-from retrievalbase.connector.settings import ParquetDatasetConnectorSettings
-
-if TYPE_CHECKING:
-    from retrievalbase.dataset import Dataset
-
-
-class ParquetDatasetConnector(DatasetConnector[ParquetDatasetConnectorSettings]):
-    def __init__(self, config: ParquetDatasetConnectorSettings):
-        super().__init__(config)
-
-    def _load(self) -> pl.DataFrame | pl.LazyFrame:
-        return pl.scan_parquet(self.config.path) if self.config.lazy else pl.read_parquet(self.config.path)
-
-    def to(self, ds: "Dataset[Any]") -> None:
-        ds.polars.write_parquet(self.config.path)
-
-    def target_exists(self, target: str) -> bool:
-        return Path(target).exists()