pyspark-fluvius 0.1.0__py3-none-any.whl

pyspark_fluvius/__init__.py

@@ -0,0 +1,71 @@
+"""PySpark custom data sources for Fluvius Energy API.
+
+This package provides PySpark data sources for reading energy and mandate data
+from the Fluvius Energy API directly into Spark DataFrames.
+
+Example:
+    ```python
+    from pyspark.sql import SparkSession
+    import pyspark_fluvius  # Registers data sources
+
+    spark = SparkSession.builder.getOrCreate()
+
+    # Read mandates
+    mandates_df = spark.read.format("fluvius.mandates") \\
+        .option("status", "Approved") \\
+        .load()
+
+    # Read energy data
+    energy_df = spark.read.format("fluvius.energy") \\
+        .option("ean", "541234567890123456") \\
+        .option("period_type", "readTime") \\
+        .option("granularity", "daily") \\
+        .option("from_date", "2024-01-01") \\
+        .option("to_date", "2024-01-31") \\
+        .load()
+    ```
+"""
+
+from __future__ import annotations
+
+from .datasources import FluviusEnergyDataSource, FluviusMandatesDataSource
+from .schemas import ENERGY_SCHEMA, MANDATES_SCHEMA
+
+__version__ = "0.1.0"
+__all__ = [
+    "FluviusEnergyDataSource",
+    "FluviusMandatesDataSource",
+    "ENERGY_SCHEMA",
+    "MANDATES_SCHEMA",
+    "register_datasources",
+]
+
+
+def register_datasources() -> None:
+    """Register Fluvius data sources with the active SparkSession.
+
+    This function registers both fluvius.energy and fluvius.mandates data sources
+    with the current SparkSession. Call this after creating your SparkSession.
+
+    Example:
+        ```python
+        from pyspark.sql import SparkSession
+        from pyspark_fluvius import register_datasources
+
+        spark = SparkSession.builder.getOrCreate()
+        register_datasources()
+
+        df = spark.read.format("fluvius.mandates").load()
+        ```
+    """
+    from pyspark.sql import SparkSession
+
+    spark = SparkSession.getActiveSession()
+    if spark is None:
+        raise RuntimeError(
+            "No active SparkSession found. "
+            "Create a SparkSession before calling register_datasources()."
+        )
+
+    spark.dataSource.register(FluviusEnergyDataSource)
+    spark.dataSource.register(FluviusMandatesDataSource)

pyspark_fluvius/converters/__init__.py

@@ -0,0 +1,6 @@
+"""Converters from Pydantic models to Spark Rows."""
+
+from .energy_converter import convert_energy_response
+from .mandates_converter import convert_mandate
+
+__all__ = ["convert_energy_response", "convert_mandate"]

pyspark_fluvius/converters/energy_converter.py

@@ -0,0 +1,427 @@
+"""Convert energy models to Spark-compatible tuples."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+
+def _get_enum_value(value: object) -> str | None:
+    """Extract string value from an enum or return string as-is."""
+    if value is None:
+        return None
+    if hasattr(value, "value"):
+        return value.value
+    return str(value)
+
+
+if TYPE_CHECKING:
+    from fluvius_energy_api.models.energy import (
+        GetEnergyResponseApiDataResponse,
+        MeasurementDirection,
+        MeasurementTimeSlice,
+        MeasurementValue,
+        MeasurementValueSet,
+        PhysicalMeter,
+        SubHeadpoint,
+    )
+
+# Tuple structure matches ENERGY_SCHEMA field order:
+# 11 header fields + 4 directions * (4 fields for total + 5 registers * 3 fields) = 11 + 4*19 = 87 fields
+EnergyTuple = tuple[
+    str | None,  # ean
+    str | None,  # energy_type
+    str | None,  # metering_type
+    datetime | None,  # measurement_start
+    datetime | None,  # measurement_end
+    str | None,  # granularity
+    str | None,  # meter_seq_number
+    str | None,  # meter_id
+    str | None,  # subheadpoint_ean
+    str | None,  # subheadpoint_type
+    str | None,  # subheadpoint_seq_number
+    # Offtake: total (4), day (3), night (3), reactive (3), inductive (3), capacitive (3) = 19 fields
+    float | None, str | None, str | None, str | None,  # offtake_total
+    float | None, str | None, str | None,  # offtake_day
+    float | None, str | None, str | None,  # offtake_night
+    float | None, str | None, str | None,  # offtake_reactive
+    float | None, str | None, str | None,  # offtake_inductive
+    float | None, str | None, str | None,  # offtake_capacitive
+    # Injection: 19 fields
+    float | None, str | None, str | None, str | None,  # injection_total
+    float | None, str | None, str | None,  # injection_day
+    float | None, str | None, str | None,  # injection_night
+    float | None, str | None, str | None,  # injection_reactive
+    float | None, str | None, str | None,  # injection_inductive
+    float | None, str | None, str | None,  # injection_capacitive
+    # Production: 19 fields
+    float | None, str | None, str | None, str | None,  # production_total
+    float | None, str | None, str | None,  # production_day
+    float | None, str | None, str | None,  # production_night
+    float | None, str | None, str | None,  # production_reactive
+    float | None, str | None, str | None,  # production_inductive
+    float | None, str | None, str | None,  # production_capacitive
+    # Auxiliary: 19 fields
+    float | None, str | None, str | None, str | None,  # auxiliary_total
+    float | None, str | None, str | None,  # auxiliary_day
+    float | None, str | None, str | None,  # auxiliary_night
+    float | None, str | None, str | None,  # auxiliary_reactive
+    float | None, str | None, str | None,  # auxiliary_inductive
+    float | None, str | None, str | None,  # auxiliary_capacitive
+]
+
+
+def _extract_measurement_value_with_gas(
+    mv: MeasurementValue | None,
+) -> tuple[float | None, str | None, str | None, str | None]:
+    """Extract fields from a MeasurementValue including gas conversion factor."""
+    if mv is None:
+        return None, None, None, None
+    return (
+        mv.value,
+        _get_enum_value(mv.unit),
+        _get_enum_value(mv.validation_state),
+        _get_enum_value(mv.gas_conversion_factor),
+    )
+
+
+def _extract_measurement_value(
+    mv: MeasurementValue | None,
+) -> tuple[float | None, str | None, str | None]:
+    """Extract fields from a MeasurementValue without gas conversion factor."""
+    if mv is None:
+        return None, None, None
+    return (
+        mv.value,
+        _get_enum_value(mv.unit),
+        _get_enum_value(mv.validation_state),
+    )
+
+
+def _extract_value_set(vs: MeasurementValueSet | None) -> tuple:
+    """Extract all measurement values from a MeasurementValueSet.
+
+    Returns 19 fields: total (4) + day (3) + night (3) + reactive (3) + inductive (3) + capacitive (3)
+
+    If total is not provided but day and night are, computes total = day + night.
+    """
+    if vs is None:
+        return (None,) * 19
+
+    # Extract raw values
+    total_fields = _extract_measurement_value_with_gas(vs.total)
+    day_fields = _extract_measurement_value(vs.day)
+    night_fields = _extract_measurement_value(vs.night)
+
+    # Compute total from day + night if total is missing
+    total_value, total_unit, total_validation, total_gas = total_fields
+    day_value, day_unit, day_validation = day_fields
+    night_value, night_unit, night_validation = night_fields
+
+    if total_value is None and day_value is not None and night_value is not None:
+        total_value = day_value + night_value
+        total_unit = day_unit  # Use day's unit (should be same as night)
+        total_validation = day_validation  # Use day's validation state
+        total_fields = (total_value, total_unit, total_validation, total_gas)
+
+    return (
+        *total_fields,
+        *day_fields,
+        *night_fields,
+        *_extract_measurement_value(vs.reactive),
+        *_extract_measurement_value(vs.inductive),
+        *_extract_measurement_value(vs.capacitive),
+    )
+
+
+def _extract_measurements(
+    directions: list[MeasurementDirection] | None,
+) -> tuple:
+    """Extract all measurement values from directions.
+
+    Returns 76 fields: 4 directions * 19 fields each.
+    """
+    if not directions:
+        return (None,) * 76
+
+    # Take the first direction (typical case has one)
+    direction = directions[0]
+
+    return (
+        *_extract_value_set(direction.offtake),
+        *_extract_value_set(direction.injection),
+        *_extract_value_set(direction.production),
+        *_extract_value_set(direction.auxiliary),
+    )
+
+
+def _process_time_slice(
+    time_slice: MeasurementTimeSlice,
+    ean: str | None,
+    energy_type: str | None,
+    metering_type: str,
+    granularity: str,
+    meter_seq_number: str | None = None,
+    meter_id: str | None = None,
+    subheadpoint_ean: str | None = None,
+    subheadpoint_type: str | None = None,
+    subheadpoint_seq_number: str | None = None,
+) -> EnergyTuple:
+    """Convert a single time slice to a tuple."""
+    measurements = _extract_measurements(time_slice.measurements)
+
+    return (
+        ean,
+        energy_type,
+        metering_type,
+        time_slice.start,
+        time_slice.end,
+        granularity,
+        meter_seq_number,
+        meter_id,
+        subheadpoint_ean,
+        subheadpoint_type,
+        subheadpoint_seq_number,
+        *measurements,
+    )
+
+
+def _process_energy_list(
+    energy_list: list[MeasurementTimeSlice] | None,
+    ean: str | None,
+    energy_type: str | None,
+    metering_type: str,
+    granularity: str,
+    meter_seq_number: str | None = None,
+    meter_id: str | None = None,
+    subheadpoint_ean: str | None = None,
+    subheadpoint_type: str | None = None,
+    subheadpoint_seq_number: str | None = None,
+) -> list[EnergyTuple]:
+    """Process a list of time slices."""
+    if not energy_list:
+        return []
+
+    return [
+        _process_time_slice(
+            ts,
+            ean,
+            energy_type,
+            metering_type,
+            granularity,
+            meter_seq_number,
+            meter_id,
+            subheadpoint_ean,
+            subheadpoint_type,
+            subheadpoint_seq_number,
+        )
+        for ts in energy_list
+    ]
+
+
+def _process_physical_meter(
+    meter: PhysicalMeter,
+    ean: str | None,
+    energy_type: str | None,
+    metering_type: str,
+) -> list[EnergyTuple]:
+    """Process energy data from a physical meter."""
+    results: list[EnergyTuple] = []
+
+    results.extend(
+        _process_energy_list(
+            meter.daily_energy,
+            ean,
+            energy_type,
+            metering_type,
+            "daily",
+            meter.seq_number,
+            meter.meter_id,
+        )
+    )
+    results.extend(
+        _process_energy_list(
+            meter.hourly_energy,
+            ean,
+            energy_type,
+            metering_type,
+            "hourly",
+            meter.seq_number,
+            meter.meter_id,
+        )
+    )
+    results.extend(
+        _process_energy_list(
+            meter.quarter_hourly_energy,
+            ean,
+            energy_type,
+            metering_type,
+            "quarter_hourly",
+            meter.seq_number,
+            meter.meter_id,
+        )
+    )
+
+    return results
+
+
+def _process_subheadpoint(
+    sub: SubHeadpoint,
+    ean: str | None,
+    energy_type: str | None,
+    metering_type: str,
+) -> list[EnergyTuple]:
+    """Process energy data from a subheadpoint."""
+    results: list[EnergyTuple] = []
+
+    subheadpoint_type = sub.type_discriminator.replace("submetering-", "")
+
+    results.extend(
+        _process_energy_list(
+            sub.daily_energy,
+            ean,
+            energy_type,
+            metering_type,
+            "daily",
+            subheadpoint_ean=sub.ean,
+            subheadpoint_type=subheadpoint_type,
+            subheadpoint_seq_number=sub.seq_number,
+        )
+    )
+    results.extend(
+        _process_energy_list(
+            sub.hourly_energy,
+            ean,
+            energy_type,
+            metering_type,
+            "hourly",
+            subheadpoint_ean=sub.ean,
+            subheadpoint_type=subheadpoint_type,
+            subheadpoint_seq_number=sub.seq_number,
+        )
+    )
+    results.extend(
+        _process_energy_list(
+            sub.quarter_hourly_energy,
+            ean,
+            energy_type,
+            metering_type,
+            "quarter_hourly",
+            subheadpoint_ean=sub.ean,
+            subheadpoint_type=subheadpoint_type,
+            subheadpoint_seq_number=sub.seq_number,
+        )
+    )
+
+    return results
+
+
+def convert_energy_response(response: GetEnergyResponseApiDataResponse) -> list[EnergyTuple]:
+    """Convert an energy API response to a list of tuples for Spark Rows.
+
+    This function flattens the nested energy response structure into rows
+    suitable for a Spark DataFrame.
+
+    Args:
+        response: The energy API response from fluvius-energy-api.
+
+    Returns:
+        A list of tuples matching the ENERGY_SCHEMA field order.
+    """
+    results: list[EnergyTuple] = []
+
+    if not response.data or not response.data.headpoint:
+        return results
+
+    headpoint = response.data.headpoint
+    ean = headpoint.ean
+    energy_type = _get_enum_value(headpoint.energy_type)
+    metering_type = headpoint.type_discriminator
+
+    # Process based on metering type
+    if metering_type == "metering-on-headpoint":
+        # Direct energy data on headpoint
+        results.extend(
+            _process_energy_list(
+                headpoint.daily_energy,  # type: ignore[attr-defined]
+                ean,
+                energy_type,
+                metering_type,
+                "daily",
+            )
+        )
+        results.extend(
+            _process_energy_list(
+                headpoint.hourly_energy,  # type: ignore[attr-defined]
+                ean,
+                energy_type,
+                metering_type,
+                "hourly",
+            )
+        )
+        results.extend(
+            _process_energy_list(
+                headpoint.quarter_hourly_energy,  # type: ignore[attr-defined]
+                ean,
+                energy_type,
+                metering_type,
+                "quarter_hourly",
+            )
+        )
+
+        # Process subheadpoints if present
+        sub_headpoints = getattr(headpoint, "sub_headpoints", None)
+        if sub_headpoints:
+            for sub in sub_headpoints:
+                results.extend(_process_subheadpoint(sub, ean, energy_type, metering_type))
+
+    elif metering_type == "metering-on-meter":
+        # Energy data on physical meters
+        physical_meters = getattr(headpoint, "physical_meters", None)
+        if physical_meters:
+            for meter in physical_meters:
+                results.extend(_process_physical_meter(meter, ean, energy_type, metering_type))
+
+    elif metering_type == "metering-on-headpoint-and-meter":
+        # Both headpoint and meter level data
+        results.extend(
+            _process_energy_list(
+                headpoint.daily_energy,  # type: ignore[attr-defined]
+                ean,
+                energy_type,
+                metering_type,
+                "daily",
+            )
+        )
+        results.extend(
+            _process_energy_list(
+                headpoint.hourly_energy,  # type: ignore[attr-defined]
+                ean,
+                energy_type,
+                metering_type,
+                "hourly",
+            )
+        )
+        results.extend(
+            _process_energy_list(
+                headpoint.quarter_hourly_energy,  # type: ignore[attr-defined]
+                ean,
+                energy_type,
+                metering_type,
+                "quarter_hourly",
+            )
+        )
+
+        # Process physical meters
+        physical_meters = getattr(headpoint, "physical_meters", None)
+        if physical_meters:
+            for meter in physical_meters:
+                results.extend(_process_physical_meter(meter, ean, energy_type, metering_type))
+
+        # Process subheadpoints
+        sub_headpoints = getattr(headpoint, "sub_headpoints", None)
+        if sub_headpoints:
+            for sub in sub_headpoints:
+                results.extend(_process_subheadpoint(sub, ean, energy_type, metering_type))
+
+    return results

pyspark_fluvius/converters/mandates_converter.py

@@ -0,0 +1,52 @@
+"""Convert mandate models to Spark-compatible tuples."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from fluvius_energy_api.models.mandate import Mandate
+
+MandateTuple = tuple[
+    str | None,  # reference_number
+    str | None,  # status
+    str | None,  # ean
+    str | None,  # energy_type
+    datetime | None,  # data_period_from
+    datetime | None,  # data_period_to
+    str | None,  # data_service_type
+    datetime | None,  # mandate_expiration_date
+    str | None,  # renewal_status
+]
+
+
+def _get_enum_value(value: object) -> str | None:
+    """Extract string value from an enum or return string as-is."""
+    if value is None:
+        return None
+    if hasattr(value, "value"):
+        return value.value
+    return str(value)
+
+
+def convert_mandate(mandate: Mandate) -> MandateTuple:
+    """Convert a Mandate Pydantic model to a tuple for Spark Row.
+
+    Args:
+        mandate: The Mandate model from fluvius-energy-api.
+
+    Returns:
+        A tuple matching the MANDATES_SCHEMA field order.
+    """
+    return (
+        mandate.reference_number,
+        _get_enum_value(mandate.status),
+        mandate.ean,
+        _get_enum_value(mandate.energy_type),
+        mandate.data_period_from,
+        mandate.data_period_to,
+        _get_enum_value(mandate.data_service_type),
+        mandate.mandate_expiration_date,
+        _get_enum_value(mandate.renewal_status),
+    )

pyspark_fluvius/datasources/__init__.py

@@ -0,0 +1,6 @@
+"""Fluvius PySpark data sources."""
+
+from .energy import FluviusEnergyDataSource
+from .mandates import FluviusMandatesDataSource
+
+__all__ = ["FluviusEnergyDataSource", "FluviusMandatesDataSource"]

pyspark_fluvius/datasources/energy.py

@@ -0,0 +1,100 @@
+"""Fluvius Energy data source for PySpark."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from pyspark.sql.datasource import DataSource, DataSourceReader
+from pyspark.sql.types import StructType
+
+from ..readers.energy_reader import FluviusEnergyReader
+from ..schemas.energy_schema import ENERGY_SCHEMA
+
+if TYPE_CHECKING:
+    pass
+
+
+class FluviusEnergyDataSource(DataSource):
+    """PySpark data source for reading Fluvius energy measurements.
+
+    This data source allows you to read energy measurement data from the
+    Fluvius Energy API directly into a Spark DataFrame.
+
+    Required Options:
+        - ean: GSRN EAN-code that identifies the installation
+        - period_type: Type of period ("readTime" or "insertTime")
+
+    Optional Options:
+        Credential options (if not using environment variables):
+            - subscription_key: Azure API Management subscription key
+            - client_id: Azure AD application (client) ID
+            - tenant_id: Azure AD tenant ID
+            - scope: OAuth2 scope
+            - data_access_contract_number: Data access contract number
+            - certificate_thumbprint: Certificate thumbprint (for cert auth)
+            - private_key: Private key in PEM format (for cert auth)
+            - client_secret: Client secret (for secret auth)
+            - credentials_prefix: Environment variable prefix (default: "FLUVIUS")
+
+        Environment options:
+            - environment: "sandbox" (default) or "production"
+
+        Filter options:
+            - reference_number: Custom reference number
+            - granularity: Granularity filter (e.g., "daily", "hourly_quarterhourly")
+            - complex_energy_types: Types of complex energy (e.g., "active,reactive")
+            - from_date: Start date (ISO format, e.g., "2024-01-01")
+            - to_date: End date (ISO format, e.g., "2024-01-31")
+
+    Example:
+        ```python
+        df = spark.read.format("fluvius.energy") \\
+            .option("ean", "541234567890123456") \\
+            .option("period_type", "readTime") \\
+            .option("granularity", "daily") \\
+            .option("from_date", "2024-01-01") \\
+            .option("to_date", "2024-01-31") \\
+            .load()
+        ```
+
+    Schema:
+        The returned DataFrame has the following columns:
+        - ean: EAN code of the installation
+        - energy_type: "E" (electricity) or "G" (gas)
+        - metering_type: Type of metering installation
+        - measurement_start: Start time of the measurement period
+        - measurement_end: End time of the measurement period
+        - granularity: Measurement granularity (daily, hourly, quarter_hourly)
+        - meter_seq_number: Physical meter sequence number (if applicable)
+        - meter_id: Physical meter ID (if applicable)
+        - subheadpoint_ean: Subheadpoint EAN (for submetering)
+        - subheadpoint_type: Type of subheadpoint
+        - subheadpoint_seq_number: Subheadpoint sequence number
+        - offtake_total_value/unit/validation_state/gas_conversion_factor
+        - offtake_day_value/unit/validation_state
+        - offtake_night_value/unit/validation_state
+        - injection_total_value/unit/validation_state
+        - injection_day_value/unit/validation_state
+        - injection_night_value/unit/validation_state
+        - production_total_value/unit/validation_state
+    """
+
+    @classmethod
+    def name(cls) -> str:
+        """Return the short name of this data source."""
+        return "fluvius.energy"
+
+    def schema(self) -> StructType:
+        """Return the schema for energy data."""
+        return ENERGY_SCHEMA
+
+    def reader(self, schema: StructType) -> DataSourceReader:
+        """Return a reader for energy data.
+
+        Args:
+            schema: The schema to use (typically the default ENERGY_SCHEMA).
+
+        Returns:
+            A FluviusEnergyReader instance.
+        """
+        return FluviusEnergyReader(schema, self.options)