dataflux-core 0.1.0__py3-none-any.whl

dataflux/exceptions.py

@@ -0,0 +1,150 @@
+class DataFluxError(Exception):
+    """Base exception for all DataFlux errors."""
+
+
+# ============================================================================
+# Provider Errors
+# ============================================================================
+
+class InvalidProviderError(DataFluxError):
+    """Raised when an invalid provider is supplied."""
+
+
+class ProviderAlreadyRegisteredError(DataFluxError):
+    """Raised when a provider with the same name is already registered."""
+
+    def __init__(self, provider_name: str):
+        super().__init__(
+            f"Provider '{provider_name}' is already registered."
+        )
+
+
+class ProviderNotFoundError(DataFluxError):
+    """Raised when a requested provider does not exist."""
+
+    def __init__(self, provider_name: str):
+        super().__init__(
+            f"Provider '{provider_name}' was not found."
+        )
+
+
+# ============================================================================
+# Dataset Errors
+# ============================================================================
+
+class DatasetNotFoundError(DataFluxError):
+    """Raised when no dataset matches the query."""
+
+    def __init__(self, dataset: str | int):
+        super().__init__(
+            f"Dataset '{dataset}' was not found."
+        )
+
+
+class InvalidDatasetError(DataFluxError):
+    """Raised when a dataset identifier is invalid."""
+
+    def __init__(self, dataset: str | int):
+        super().__init__(
+            f"'{dataset}' is not a valid dataset identifier."
+        )
+
+
+class DatasetLoadError(DataFluxError):
+    """Raised when a dataset cannot be loaded."""
+
+    def __init__(self, dataset: str | int):
+        super().__init__(
+            f"Failed to load dataset '{dataset}'."
+        )
+
+
+# ============================================================================
+# Search Errors
+# ============================================================================
+
+class SearchError(DataFluxError):
+    """Raised when a search operation fails."""
+
+
+class EmptySearchQueryError(SearchError):
+    """Raised when an empty search query is provided."""
+
+    def __init__(self):
+        super().__init__(
+            "Search query cannot be empty."
+        )
+
+
+# ============================================================================
+# Pull Errors
+# ============================================================================
+
+class PullError(DataFluxError):
+    """Raised when pulling a dataset fails."""
+
+
+# ============================================================================
+# Download Errors
+# ============================================================================
+
+class DownloadError(DataFluxError):
+    """Raised when a download operation fails."""
+
+    def __init__(self, url: str):
+        super().__init__(
+            f"Failed to download resource from '{url}'."
+        )
+
+
+# ============================================================================
+# Cache Errors
+# ============================================================================
+
+class CacheError(DataFluxError):
+    """Raised when a cache operation fails."""
+
+
+# ============================================================================
+# Filesystem Errors
+# ============================================================================
+
+class FileSystemError(DataFluxError):
+    """Raised when a filesystem operation fails."""
+
+
+# ============================================================================
+# Export Errors
+# ============================================================================
+
+class ExportError(DataFluxError):
+    """Raised when exporting a dataset fails."""
+
+
+class UnsupportedExportFormatError(ExportError):
+    """Raised when an unsupported export format is requested."""
+
+    def __init__(self, extension: str):
+        super().__init__(
+            f"Export format '{extension}' is not supported."
+        )
+
+
+class ExportFileExistsError(ExportError):
+    """Raised when the export file already exists."""
+
+    def __init__(self, path: str):
+        super().__init__(
+            f"'{path}' already exists. "
+            "Use overwrite=True to replace it."
+        )
+
+
+class InvalidExportDataError(ExportError):
+    """Raised when the object being exported is invalid."""
+
+    def __init__(
+        self,
+        message: str = "export() expects a Polars DataFrame.",
+    ):
+        super().__init__(message)

dataflux/export.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+from pathlib import Path
+import polars as pl
+from dataflux.exceptions import ExportFileExistsError,UnsupportedExportFormatError,InvalidExportDataError
+SUPPORTED_EXPORTS = {
+    ".csv",
+    ".parquet",
+    ".json",
+    ".ndjson",
+    ".ipc",
+    ".feather",
+    ".arrow",
+}
+
+
+def export(df:pl.DataFrame,path:str | Path,*,overwrite:bool = False,mkdir:bool=True,)->Path:
+    if not isinstance(df,pl.DataFrame):
+        raise InvalidExportDataError(
+            "export() expect a Polars DataFrame"
+        )
+    path = Path(path)
+
+    if mkdir:
+        path.parent.mkdir(
+            parents=True,
+            exist_ok=True,
+        )
+    if path.exists() and not overwrite:
+        raise ExportFileExistsError(
+            f"'{path}' already exists."
+        )
+
+    suffix = path.suffix.lower()
+
+    if suffix not in SUPPORTED_EXPORTS:
+        raise UnsupportedExportFormatError(
+            f"Unsupported export format '{suffix}'. "
+            f"Supported formats:{', '.join(sorted(SUPPORTED_EXPORTS))} "
+        )
+    match suffix:
+        case ".csv":
+            df.write_csv(path)
+
+        case ".parquet":
+            df.write_parquet(path)
+
+        case ".json":
+            df.write_json(path)
+
+        case ".ndjson":
+            df.write_ndjson(path)
+
+        case ".ipc":
+            df.write_ipc(path)
+
+        case ".feather":
+            df.write_ipc(path)
+
+        case ".arrow":
+            df.write_ipc(path)
+
+    return path

dataflux/flux.py

@@ -0,0 +1,55 @@
+from dataflux.models.search_result import SearchResult
+from dataflux.registry import ProviderRegistry
+from dataflux.resolver import Resolver
+from dataflux.providers.uci import UCIProvider
+from dataflux.models.dataset import DatasetInfo
+from dataflux.providers.sklearn import SKLearnProvider
+from dataflux.providers.kaggle import KaggleProvider
+from dataflux.providers.huggingface import HuggingFaceProvider
+from dataflux.providers.torchvision import TorchVisionProvider
+from dataflux.providers.seaborn import SeabornProvider
+from dataflux.providers.statsmodels import StatsModelsProvider
+from dataflux.providers.vega_datasets import VegaDatasetsProvider
+from dataflux.providers.worldbank import WorldBankProvider
+from dataflux.display.search import display_search
+from dataflux.display.info import display_info 
+from dataflux.export import export
+import polars as pl
+
+class Flux:
+
+    def __init__(self):
+        self._registry = ProviderRegistry()
+        self._resolver = Resolver(self._registry)
+
+        self._registry.register(UCIProvider())
+        self._registry.register(SKLearnProvider())
+        self._registry.register(KaggleProvider())
+        self._registry.register(HuggingFaceProvider())
+        self._registry.register(TorchVisionProvider())
+        self._registry.register(SeabornProvider())
+        self._registry.register(StatsModelsProvider())
+        self._registry.register(VegaDatasetsProvider())
+        self._registry.register(WorldBankProvider())
+
+    def search(self, query: str,*,raw:bool=False,display:bool=True,limit:int | None=10) -> list[SearchResult]:
+        results = self._resolver.resolve_dataset(query)
+        if display and not raw:
+            display_search(results,limit)
+        return results
+
+
+    def info(self,result:SearchResult,*,raw:bool=False,display:bool=True)-> DatasetInfo:
+        provider = self._resolver.resolve_provider(result.provider)
+        if display and not raw:
+            display_info(provider.info(result.id))
+        return provider.info(result.id)
+
+    def pull(self,result:SearchResult)->pl.DataFrame:
+        provider = self._resolver.resolve_provider(result.provider)
+        return provider.pull(result.id)
+
+    def export(self,df:pl.DataFrame,path:str,**kwargs,):
+        return export(df,path,**kwargs)
+
+flux = Flux()

dataflux/models/dataset.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from typing import Any
+@dataclass(slots=True)
+class DatasetInfo:
+    id: str | int
+    name: str
+    description: str | None
+    instances: int | None
+    features: int | None
+    tasks: list[str]
+    target: list[str]
+    has_missing_values: bool | None
+    provider: str
+    url: str | None
+    extra: dict[str, Any]

dataflux/models/search_result.py

@@ -0,0 +1,10 @@
+from dataclasses import dataclass ,field
+
+@dataclass(slots=True)
+class SearchResult:
+    id:str 
+    name:str
+    provider:str
+    description:str | None = None
+
+    relevance: int = field(default=0, repr=False)

dataflux/providers/base.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+
+class BaseProvider(ABC):
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Unique provider name."""
+
+    @abstractmethod
+    def search(self, query: str) -> Any:
+        """Search datasets."""
+
+    @abstractmethod
+    def pull(self, dataset: str) -> Path:
+        """Download the requested dataset and return its local path."""
+
+    @abstractmethod
+    def info(self, dataset: str) -> Any:
+        """Return dataset metadata."""

dataflux/providers/huggingface.py

@@ -0,0 +1,91 @@
+from dataflux.providers.base import BaseProvider
+from dataflux.models.search_result import SearchResult
+from dataflux.models.dataset import DatasetInfo
+from huggingface_hub import HfApi
+from datasets import load_dataset
+from dataflux.utils.search import search_score
+import polars as pl
+
+
+class HuggingFaceProvider(BaseProvider):
+
+    def __init__(self):
+        self.api = HfApi()
+
+    @property
+    def name(self) -> str:
+        return "huggingface"
+
+    def search(self, query: str) -> list[SearchResult]:
+        matches:list[tuple[int,SearchResult]] =[]
+
+        datasets = self.api.list_datasets(
+            search=query,
+            limit=20
+        )
+        for dataset in datasets:
+            score = search_score(
+                query,
+                dataset.id
+            )
+            if score == -1:
+                continue
+            matches.append(
+                ( score,
+                SearchResult(
+                    id=dataset.id,
+                    name=dataset.id.split("/")[-1].replace("-"," ").title(),
+                    provider=self.name,
+                    description=None,
+                    relevance=score,
+                )
+            )
+        )
+        matches.sort(key=lambda x:x[0],reverse=True)
+        return [result for _, result in matches]
+
+
+    def info(self, dataset_id: str) -> DatasetInfo:
+        dataset = self.api.dataset_info(dataset_id)
+        return DatasetInfo(
+            id=dataset_id,
+            name=dataset.id.split("/")[-1].replace("-"," ").title(),
+            description=dataset.description,
+            instances=(
+                dataset.card_data["dataset_info"]["splits"][0]["num_examples"]
+                if dataset.card_data and "dataset_info" in dataset.card_data
+                else None
+            ),
+            features=(
+                len(dataset.card_data["dataset_info"]["features"])
+                if dataset.card_data and "dataset_info" in dataset.card_data
+                else None
+            ),
+            tasks=(
+                dataset.card_data.get("task_categories")
+                if dataset.card_data
+                else None
+            ),
+            target=None,
+            has_missing_values=None,
+            provider=self.name,
+            url=f"https://huggingface.co/datasets/{dataset_id}",
+            extra={
+                "author":dataset.author,
+                "tags":dataset.tags,
+                "downloads":dataset.downloads,
+                "likes":dataset.likes,
+                "created_at":str(dataset.created_at),
+                "last_modified": str(dataset.last_modified),
+            }
+        )
+
+        
+
+    def pull(self, dataset_id: str):
+        dataset = load_dataset(dataset_id)
+        if hasattr(dataset,"keys"):
+            dataset = dataset[list(dataset.keys())[0]]
+
+        return dataset.to_polars()
+

dataflux/providers/kaggle.py

@@ -0,0 +1,85 @@
+import json
+from pathlib import Path
+
+import kagglehub
+import polars as pl
+from kagglehub import KaggleDatasetAdapter
+
+from dataflux.models.dataset import DatasetInfo
+from dataflux.models.search_result import SearchResult
+from dataflux.providers.base import BaseProvider
+from dataflux.utils.search import search_score
+
+
+class KaggleProvider(BaseProvider):
+
+    @property
+    def name(self) -> str:
+        return "kaggle"
+
+    def _load_index(self) -> list[dict]:
+        with open(
+            Path(__file__).parent.parent / "resources" / "kaggle_index.json",
+            "r",
+            encoding="utf-8",
+        ) as f:
+            return json.load(f)
+
+    def search(self, query: str) -> list[SearchResult]:
+        matches: list[tuple[int, SearchResult]] = []
+
+        for dataset in self._load_index():
+            score = search_score(
+                query,
+                dataset["title"],
+                dataset["subtitle"],
+            )
+
+            if score == -1:
+                continue
+
+            matches.append(
+                (
+                    score,
+                    SearchResult(
+                        id=f"{dataset['owner_slug']}/{dataset['dataset_slug']}",
+                        name=dataset["title"],
+                        provider=self.name,
+                        description=dataset["subtitle"],
+                        relevance=score
+                    ),
+                )
+            )
+
+        matches.sort(key=lambda x: x[0], reverse=True)
+
+        return [result for _, result in matches]
+
+    def info(self, dataset_id: str) -> DatasetInfo:
+        for dataset in self._load_index():
+            if f"{dataset['owner_slug']}/{dataset['dataset_slug']}" == dataset_id:
+                return DatasetInfo(
+                    id=dataset_id,
+                    name=dataset["title"],
+                    description=dataset["subtitle"],
+                    instances=None,
+                    features=None,
+                    tasks=None,
+                    target=None,
+                    has_missing_values=None,
+                    provider=self.name,
+                    url=f"https://www.kaggle.com/datasets/{dataset_id}",
+                    extra=dataset,
+                )
+
+        raise ValueError(f"Dataset '{dataset_id}' not found.")
+
+    def pull(self, dataset_id: str) -> pl.DataFrame:
+        dataset_path = Path(kagglehub.dataset_download(dataset_id))
+        csv_file = next(dataset_path.rglob("*.csv"))
+
+        return kagglehub.dataset_load(
+            KaggleDatasetAdapter.POLARS,
+            handle=dataset_id,
+            path=csv_file.name,
+        ).collect()

dataflux/providers/seaborn.py

@@ -0,0 +1,66 @@
+from dataflux.providers.base import BaseProvider
+from dataflux.models.search_result import SearchResult
+from dataflux.models.dataset import DatasetInfo
+from dataflux.resources.seaborn_dataset import SEABORN_DATASETS
+from dataflux.utils.search import search_score
+import seaborn as sns
+import polars as pl
+
+
+class SeabornProvider(BaseProvider):
+
+    @property
+    def name(self) -> str:
+        return "seaborn"
+
+    def _load_df(self, dataset_id: str):
+        return sns.load_dataset(dataset_id)
+
+    def search(self, query: str) -> list[SearchResult]:
+        matches: list[tuple[int, SearchResult]] = []
+
+        for dataset in SEABORN_DATASETS:
+            score = search_score(query, dataset)
+
+            if score == -1:
+                continue
+
+            matches.append(
+                (
+                    score,
+                    SearchResult(
+                        id=dataset,
+                        name=dataset.replace("_", " ").title(),
+                        provider=self.name,
+                        description=None,
+                        relevance=score
+                    ),
+                )
+            )
+
+        matches.sort(key=lambda x: x[0], reverse=True)
+
+        return [result for _, result in matches]
+
+    def info(self, dataset_id: str) -> DatasetInfo:
+        dataset = self._load_df(dataset_id)
+
+        return DatasetInfo(
+            id=dataset_id,
+            name=dataset_id.replace("_", " ").title(),
+            description=None,
+            instances=dataset.shape[0],
+            features=dataset.shape[1],
+            tasks=None,
+            target=None,
+            has_missing_values=dataset.isnull().values.any(),
+            provider=self.name,
+            url=None,
+            extra={
+                "columns": list(dataset.columns),
+                "dtypes": {k: str(v) for k, v in dataset.dtypes.items()},
+            },
+        )
+
+    def pull(self, dataset_id: str) -> pl.DataFrame:
+        return pl.from_pandas(self._load_df(dataset_id))

dataflux/providers/sklearn.py

@@ -0,0 +1,67 @@
+from dataflux.providers.base import BaseProvider
+from dataflux.models.search_result import SearchResult
+from dataflux.models.dataset import DatasetInfo
+from dataflux.resources.sklearn_dataset import SKLEARN_DATASETS
+from dataflux.utils.search import search_score
+import polars as pl
+
+
+class SKLearnProvider(BaseProvider):
+
+    @property
+    def name(self) -> str:
+        return "sklearn"
+
+    def _load_dataset(self, dataset_id: str, *, as_frame: bool = False):
+        loader = SKLEARN_DATASETS[dataset_id]
+        return loader(as_frame=as_frame) if as_frame else loader()
+
+    def search(self, query: str) -> list[SearchResult]:
+        matches: list[tuple[int, SearchResult]] = []
+
+        for dataset in SKLEARN_DATASETS:
+            score = search_score(query, dataset)
+
+            if score == -1:
+                continue
+
+            matches.append(
+                (
+                    score,
+                    SearchResult(
+                        id=dataset,
+                        name=dataset.replace("_", " ").title(),
+                        provider=self.name,
+                        relevance=score
+                    ),
+                )
+            )
+
+        matches.sort(key=lambda x: x[0], reverse=True)
+
+        return [result for _, result in matches]
+
+    def info(self, dataset_id: str) -> DatasetInfo:
+        dataset = self._load_dataset(dataset_id)
+
+        return DatasetInfo(
+            id=dataset_id,
+            name=dataset_id.replace("_", " ").title(),
+            description=dataset.DESCR,
+            instances=dataset.data.shape[0],
+            features=dataset.data.shape[1],
+            tasks=None,
+            target=list(dataset.target_names),
+            has_missing_values=False,
+            provider=self.name,
+            url=None,
+            extra={
+                "feature_names": dataset.feature_names,
+                "filename": dataset.filename,
+                "data_module": dataset.data_module,
+            },
+        )
+
+    def pull(self, dataset_id: str) -> pl.DataFrame:
+        dataset = self._load_dataset(dataset_id, as_frame=True)
+        return pl.from_pandas(dataset.frame)