aptdata 0.0.2__py3-none-any.whl

aptdata/plugins/quality/validator.py

@@ -0,0 +1,139 @@
+"""Quality validator — wraps expectations into a workflow-compatible step."""
+
+from __future__ import annotations
+
+import logging
+import warnings
+from typing import Any
+
+from aptdata.plugins.quality.contract import EnforcementMode
+from aptdata.plugins.quality.expectations import BaseExpectation
+from aptdata.plugins.quality.report import CheckStatus, QualityReport
+from aptdata.telemetry.instrumentation import get_tracer
+
+logger = logging.getLogger(__name__)
+
+
+class QualityValidator:
+    """Runs a suite of :class:`~.expectations.BaseExpectation` objects against data.
+
+    Parameters
+    ----------
+    expectations:
+        List of expectations to evaluate.
+    enforcement:
+        How to react when an expectation fails
+        (see :class:`~.contract.EnforcementMode`).
+    name:
+        Human-readable identifier used in OTel span names.
+
+    Examples
+    --------
+    ::
+
+        from aptdata.plugins.quality import (
+            QualityValidator, ExpectColumnToNotBeNull, EnforcementMode
+        )
+
+        validator = QualityValidator(
+            expectations=[ExpectColumnToNotBeNull("age")],
+            enforcement=EnforcementMode.ABORT,
+        )
+        clean_df = validator.validate(df)
+    """
+
+    def __init__(
+        self,
+        expectations: list[BaseExpectation],
+        enforcement: EnforcementMode = EnforcementMode.ABORT,
+        name: str = "QualityValidator",
+    ) -> None:
+        self.expectations = expectations
+        self.enforcement = enforcement
+        self.name = name
+
+    def validate(self, data: Any) -> Any:
+        """Validate *data* against all expectations and return *data*.
+
+        This method is compatible with :meth:`~aptdata.core.workflow.Workflow.add_step`.
+        On success the original *data* object is returned unchanged.
+
+        Parameters
+        ----------
+        data:
+            A ``pd.DataFrame``, PySpark ``DataFrame``,
+            :class:`~aptdata.plugins.dataset.InMemoryDataset`, or
+            ``list[dict]``.
+
+        Returns
+        -------
+        Any
+            The original *data* object (pass-through).
+
+        Raises
+        ------
+        ValueError
+            When ``enforcement == ABORT`` and at least one expectation fails.
+        """
+        import pandas as pd  # type: ignore[import]
+
+        from aptdata.plugins.dataset import InMemoryDataset
+
+        # Resolve to a DataFrame for expectations.
+        resolved: Any
+        if isinstance(data, InMemoryDataset):
+            resolved = pd.DataFrame(data.read())
+        elif isinstance(data, list):
+            resolved = pd.DataFrame(data)
+        else:
+            resolved = data
+
+        dataset_uri = data.uri if isinstance(data, InMemoryDataset) else "unknown"
+        report = QualityReport(dataset_uri=dataset_uri)
+
+        tracer = get_tracer("aptdata.quality")
+        with tracer.start_as_current_span(self.name) as span:
+            span.set_attribute("aptdata.quality.validator_name", self.name)
+            span.set_attribute("aptdata.quality.enforcement", self.enforcement.value)
+            span.set_attribute(
+                "aptdata.quality.num_expectations", len(self.expectations)
+            )
+
+            for expectation in self.expectations:
+                result = expectation.validate(resolved)
+                report.checks.append(result)
+
+            passed = report.passed
+            span.set_attribute("aptdata.quality.passed", passed)
+            span.set_attribute("aptdata.quality.num_checks", len(report.checks))
+            failed_count = report.summary.get(CheckStatus.FAILED, 0)
+            span.set_attribute("aptdata.quality.failed_checks", failed_count)
+
+        if not passed:
+            failed_checks = [c for c in report.checks if c.status == CheckStatus.FAILED]
+            summary_msg = "; ".join(c.message for c in failed_checks)
+
+            if self.enforcement == EnforcementMode.ABORT:
+                raise ValueError(
+                    f"Data quality validation failed [{self.name}]: {summary_msg}"
+                )
+            elif self.enforcement == EnforcementMode.WARN:
+                warnings.warn(
+                    f"Data quality warning [{self.name}]: {summary_msg}",
+                    stacklevel=2,
+                )
+                logger.warning("Quality validation warning: %s", summary_msg)
+            else:
+                # TAG mode — attach quality metadata; pass-through
+                if hasattr(data, "schema_metadata") and isinstance(
+                    data.schema_metadata, dict
+                ):
+                    data.schema_metadata["quality_report"] = {
+                        "passed": False,
+                        "failed_checks": [c.expectation_name for c in failed_checks],
+                    }
+
+        return data
+
+
+__all__ = ["QualityValidator"]

aptdata/plugins/rest.py

@@ -0,0 +1,135 @@
+"""REST API reader plugin.
+
+Provides :class:`APIReader` — a generic reader that fetches data from a
+RESTful JSON endpoint using HTTP GET, with optional header injection
+(e.g. for bearer tokens) and simple offset/page-based pagination.
+
+Requires the optional ``httpx`` package.  A friendly
+:class:`~aptdata.plugins.manager.PluginDependencyError` is raised
+when it is not installed.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from aptdata.plugins.base import BaseReader
+from aptdata.plugins.dataset import InMemoryDataset
+from aptdata.plugins.manager import PluginDependencyError
+
+
+def _require_httpx() -> Any:
+    """Import and return the ``httpx`` module, or raise a friendly error."""
+    try:
+        import httpx  # noqa: WPS433
+    except ImportError:
+        raise PluginDependencyError("APIReader", "httpx") from None
+    return httpx
+
+
+class APIReader(BaseReader):
+    """Read JSON data from a REST API endpoint.
+
+    Parameters
+    ----------
+    endpoint:
+        The URL to send GET requests to.
+    headers:
+        Optional HTTP headers (e.g. ``{"Authorization": "Bearer xxx"}``).
+    params:
+        Optional query-string parameters merged into every request.
+    pagination_key:
+        When set, the reader fetches pages until the response JSON list is
+        empty.  The value is the query-string parameter name that carries the
+        page number (e.g. ``"page"``).  Pages start at ``1``.
+    max_pages:
+        Safety limit on the number of pages fetched (default ``100``).
+    timeout:
+        HTTP request timeout in seconds (default ``30``).
+    records_path:
+        Optional dot-separated path into the JSON response that contains
+        the list of records (e.g. ``"data.items"``).  When ``None`` the
+        response itself must be a JSON array.
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        *,
+        headers: dict[str, str] | None = None,
+        params: dict[str, Any] | None = None,
+        pagination_key: str | None = None,
+        max_pages: int = 100,
+        timeout: float = 30,
+        records_path: str | None = None,
+    ) -> None:
+        self.endpoint = endpoint
+        self.headers = headers or {}
+        self.params = params or {}
+        self.pagination_key = pagination_key
+        self.max_pages = max_pages
+        self.timeout = timeout
+        self.records_path = records_path
+
+    # -- helpers ------------------------------------------------------------
+
+    @staticmethod
+    def _extract_records(body: Any, path: str | None) -> list[dict[str, Any]]:
+        """Navigate *path* inside *body* and return the records list."""
+        if path is None:
+            if isinstance(body, list):
+                return body
+            raise ValueError(
+                "Expected a JSON array from the API response. "
+                "Set 'records_path' if records are nested."
+            )
+        current: Any = body
+        for key in path.split("."):
+            if isinstance(current, dict):
+                current = current[key]
+            else:
+                raise ValueError(f"Cannot traverse key '{key}' in a non-dict value.")
+        if not isinstance(current, list):
+            raise ValueError(
+                f"Expected a list at path '{path}', got {type(current).__name__}."
+            )
+        return current
+
+    # -- BaseReader ---------------------------------------------------------
+
+    def read(self, **kwargs: Any) -> InMemoryDataset:
+        httpx = _require_httpx()
+
+        all_records: list[dict[str, Any]] = []
+
+        with httpx.Client(timeout=self.timeout) as client:
+            if self.pagination_key is None:
+                response = client.get(
+                    self.endpoint,
+                    headers=self.headers,
+                    params=self.params,
+                )
+                response.raise_for_status()
+                all_records = self._extract_records(response.json(), self.records_path)
+            else:
+                for page in range(1, self.max_pages + 1):
+                    params = {**self.params, self.pagination_key: page}
+                    response = client.get(
+                        self.endpoint,
+                        headers=self.headers,
+                        params=params,
+                    )
+                    response.raise_for_status()
+                    page_records = self._extract_records(
+                        response.json(), self.records_path
+                    )
+                    if not page_records:
+                        break
+                    all_records.extend(page_records)
+
+        ds = InMemoryDataset(uri=self.endpoint)
+        ds.write(all_records)
+        return ds
+
+
+__all__ = ["APIReader"]

aptdata/plugins/transform/__init__.py

@@ -0,0 +1,14 @@
+"""Transform plugin package — engine-agnostic transformation wrappers.
+
+Provides :class:`PandasTransformer` and :class:`PySparkTransformer` as
+concrete :class:`~aptdata.plugins.base.BaseTransformer` implementations.
+Both use lazy imports so the framework core works without pandas or pyspark
+installed.
+"""
+
+from __future__ import annotations
+
+from aptdata.plugins.transform.pandas import PandasTransformer
+from aptdata.plugins.transform.spark import PySparkTransformer
+
+__all__ = ["PandasTransformer", "PySparkTransformer"]

aptdata/plugins/transform/pandas.py

@@ -0,0 +1,129 @@
+"""Pandas-based data transformer.
+
+Uses lazy imports so that the framework core works even when pandas is not
+installed.  A :class:`~aptdata.plugins.manager.PluginDependencyError` is
+raised at instantiation time if pandas is not available.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from aptdata.plugins.base import BaseTransformer
+from aptdata.plugins.manager import PluginDependencyError
+from aptdata.telemetry.instrumentation import get_tracer
+
+if TYPE_CHECKING:
+    import pandas as pd
+
+
+class PandasTransformer(BaseTransformer):
+    """Engine-agnostic transformer backed by a pandas callable.
+
+    Parameters
+    ----------
+    name:
+        Human-readable identifier for this transformer.
+    transform_func:
+        A callable ``(pd.DataFrame) -> pd.DataFrame`` that applies the
+        desired transformation.
+
+    Examples
+    --------
+    ::
+
+        from aptdata.plugins.transform import PandasTransformer
+
+        def double_values(df):
+            return df.assign(value=df["value"] * 2)
+
+        transformer = PandasTransformer("double_values", double_values)
+        result = transformer.transform(my_dataset)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        transform_func: Callable[[Any], Any],
+    ) -> None:
+        try:
+            import pandas  # noqa: F401
+        except ImportError as exc:
+            raise PluginDependencyError("PandasTransformer", "pandas") from exc
+        self._name = name
+        self._transform_func = transform_func
+
+    @property
+    def name(self) -> str:
+        """Human-readable name of this transformer."""
+        return self._name
+
+    def transform(self, data: Any) -> Any:
+        """Apply the transformation to *data*.
+
+        Parameters
+        ----------
+        data:
+            Accepts a :class:`~aptdata.plugins.dataset.InMemoryDataset`,
+            a ``pd.DataFrame``, or a ``list[dict]``.
+
+        Returns
+        -------
+        Any
+            The transformed result.  If the input was an
+            :class:`~aptdata.plugins.dataset.InMemoryDataset` the output
+            will also be one; otherwise a ``pd.DataFrame`` is returned.
+        """
+        import pandas as pd
+
+        from aptdata.plugins.dataset import InMemoryDataset
+
+        return_dataset = False
+        original_uri = "memory://transformed"
+        original_schema: dict[str, Any] = {}
+
+        if isinstance(data, InMemoryDataset):
+            return_dataset = True
+            original_uri = data.uri
+            original_schema = data.schema_metadata
+            df: pd.DataFrame = pd.DataFrame(data.read())
+        elif isinstance(data, pd.DataFrame):
+            df = data
+        elif isinstance(data, list):
+            df = pd.DataFrame(data)
+        else:
+            df = pd.DataFrame([data] if data is not None else [])
+
+        rows_in = len(df)
+
+        tracer = get_tracer("aptdata.transform")
+        with tracer.start_as_current_span(self._name) as span:
+            span.set_attribute("aptdata.transformer.name", self._name)
+            span.set_attribute("aptdata.transformer.engine", "pandas")
+            span.set_attribute("aptdata.transformer.rows_in", rows_in)
+
+            t0 = time.perf_counter()
+            result_df: pd.DataFrame = self._transform_func(df)
+            compute_time_ms = (time.perf_counter() - t0) * 1000.0
+
+            rows_out = len(result_df)
+            columns_out = list(result_df.columns)
+            span.set_attribute("aptdata.transformer.rows_out", rows_out)
+            span.set_attribute("aptdata.transformer.columns_out", str(columns_out))
+            span.set_attribute(
+                "aptdata.transformer.compute_time_ms", round(compute_time_ms, 3)
+            )
+
+        if return_dataset:
+            out_dataset = InMemoryDataset(
+                uri=original_uri, schema_metadata=original_schema
+            )
+            out_dataset.write(result_df.to_dict(orient="records"))
+            return out_dataset
+
+        return result_df
+
+
+__all__ = ["PandasTransformer"]

aptdata/plugins/transform/spark.py

@@ -0,0 +1,134 @@
+"""PySpark-based data transformer.
+
+Uses lazy imports so that the framework core works even when pyspark is not
+installed.  A :class:`~aptdata.plugins.manager.PluginDependencyError` is
+raised at instantiation time if pyspark is not available.
+"""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from typing import Any
+
+from aptdata.plugins.base import BaseTransformer
+from aptdata.plugins.manager import PluginDependencyError
+from aptdata.telemetry.instrumentation import get_tracer
+
+
+class PySparkTransformer(BaseTransformer):
+    """Engine-agnostic transformer backed by a PySpark callable.
+
+    Parameters
+    ----------
+    name:
+        Human-readable identifier for this transformer.
+    transform_func:
+        A callable ``(SparkSession, DataFrame) -> DataFrame`` that applies
+        the desired transformation.
+    app_name:
+        Spark application name passed to ``SparkSession.builder``.
+
+    Examples
+    --------
+    ::
+
+        from aptdata.plugins.transform import PySparkTransformer
+
+        def double_values(spark, df):
+            from pyspark.sql import functions as F
+            return df.withColumn("value", F.col("value") * 2)
+
+        transformer = PySparkTransformer("double_values", double_values)
+        result = transformer.transform(my_df)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        transform_func: Callable[[Any, Any], Any],
+        app_name: str = "SmartData",
+    ) -> None:
+        try:
+            import pyspark  # noqa: F401
+        except ImportError as exc:
+            raise PluginDependencyError("PySparkTransformer", "pyspark") from exc
+        self._name = name
+        self._transform_func = transform_func
+        self._app_name = app_name
+
+    @property
+    def name(self) -> str:
+        """Human-readable name of this transformer."""
+        return self._name
+
+    def transform(self, data: Any) -> Any:
+        """Apply the transformation to *data*.
+
+        Parameters
+        ----------
+        data:
+            Accepts a PySpark ``DataFrame``, a ``list[dict]``, or an
+            :class:`~aptdata.plugins.dataset.InMemoryDataset`.
+
+        Returns
+        -------
+        Any
+            The transformed PySpark ``DataFrame``.
+        """
+        from pyspark.sql import SparkSession
+        from pyspark.sql.types import StructType
+
+        from aptdata.plugins.dataset import InMemoryDataset
+
+        spark: Any = SparkSession.builder.appName(self._app_name).getOrCreate()
+
+        # Convert input to a PySpark DataFrame if needed.
+        if isinstance(data, InMemoryDataset):
+            records = data.read()
+            df = (
+                spark.createDataFrame(records)
+                if records
+                else spark.createDataFrame([], StructType([]))
+            )
+        elif isinstance(data, list):
+            df = (
+                spark.createDataFrame(data)
+                if data
+                else spark.createDataFrame([], StructType([]))
+            )
+        else:
+            # Assume it's already a PySpark DataFrame.
+            df = data
+
+        rows_in = df.count()
+
+        tracer = get_tracer("aptdata.transform")
+        with tracer.start_as_current_span(self._name) as span:
+            span.set_attribute("aptdata.transformer.name", self._name)
+            span.set_attribute("aptdata.transformer.engine", "pyspark")
+            span.set_attribute("aptdata.transformer.rows_in", rows_in)
+            span.set_attribute("aptdata.spark.app_name", self._app_name)
+
+            # Try to capture Spark UI URL.
+            try:
+                ui_url = spark.sparkContext.uiWebUrl
+                if ui_url:
+                    span.set_attribute("aptdata.spark.ui_url", ui_url)
+            except Exception:  # noqa: BLE001
+                pass
+
+            t0 = time.perf_counter()
+            result_df = self._transform_func(spark, df)
+            compute_time_ms = (time.perf_counter() - t0) * 1000.0
+
+            rows_out = result_df.count()
+            span.set_attribute("aptdata.transformer.rows_out", rows_out)
+            span.set_attribute(
+                "aptdata.transformer.compute_time_ms", round(compute_time_ms, 3)
+            )
+
+        return result_df
+
+
+__all__ = ["PySparkTransformer"]

aptdata/plugins/vector/__init__.py

@@ -0,0 +1,6 @@
+"""Vector database writers."""
+
+from aptdata.plugins.vector.base import VectorWriter
+from aptdata.plugins.vector.qdrant import QdrantWriter
+
+__all__ = ["VectorWriter", "QdrantWriter"]

aptdata/plugins/vector/base.py

@@ -0,0 +1,19 @@
+"""Base abstractions for vector DB writers."""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any
+
+from aptdata.core.dataset import BaseDataset
+from aptdata.plugins.base import BaseWriter
+
+
+class VectorWriter(BaseWriter):
+    """Base writer for vector databases."""
+
+    @abstractmethod
+    def write(
+        self, dataset: BaseDataset, **kwargs: Any
+    ) -> None:  # pragma: no cover - interface only
+        """Persist vectors from *dataset* into the destination vector store."""

aptdata/plugins/vector/qdrant.py

@@ -0,0 +1,41 @@
+"""Qdrant vector writer plugin."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from opentelemetry import trace
+
+from aptdata.core.dataset import BaseDataset
+from aptdata.plugins.vector.base import VectorWriter
+
+
+class QdrantWriter(VectorWriter):
+    """Write embeddings to an in-memory Qdrant-like collection buffer."""
+
+    def __init__(
+        self, *, collection: str, vector_column: str = "artigo_chunk_embedding"
+    ) -> None:
+        self.collection = collection
+        self.vector_column = vector_column
+        self.points: list[dict[str, Any]] = []
+
+    def write(self, dataset: BaseDataset, **kwargs: Any) -> None:
+        rows: list[dict[str, Any]] = dataset.read()
+        with trace.get_tracer("aptdata.plugins.vector").start_as_current_span(
+            "QdrantWriter.write"
+        ) as span:
+            for index, row in enumerate(rows):
+                vector = row.get(self.vector_column)
+                if vector is None:
+                    continue
+                point = {
+                    "id": row.get("document_id")
+                    or row.get("id")
+                    or f"{self.collection}-{index}",
+                    "vector": vector,
+                    "payload": row,
+                }
+                self.points.append(point)
+            span.set_attribute("aptdata.vector.collection", self.collection)
+            span.set_attribute("aptdata.vector.points", len(self.points))

aptdata/telemetry/__init__.py

@@ -0,0 +1,5 @@
+"""Telemetry helpers for aptdata."""
+
+from aptdata.telemetry.instrumentation import configure_telemetry, get_meter, get_tracer
+
+__all__ = ["configure_telemetry", "get_meter", "get_tracer"]