fundamental-client 0.2.3__py3-none-any.whl

fundamental/utils/data.py

@@ -0,0 +1,437 @@
+"""
+Data utilities for NEXUS client.
+
+This module provides data serialization and transformation utilities.
+"""
+
+from __future__ import annotations
+
+import logging
+import typing
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import pandas as pd
+import pyarrow.lib as palib
+from pydantic import BaseModel
+from sklearn.base import BaseEstimator
+
+from fundamental.clients.base import BaseClient
+from fundamental.constants import (
+    DEFAULT_COMPLETE_MULTIPART_UPLOAD_TIMEOUT_SECONDS,
+    DEFAULT_DOWNLOAD_RESULT_TIMEOUT_SECONDS,
+    DEFAULT_MODEL_METADATA_GENERATE_TIMEOUT_SECONDS,
+    DEFAULT_MODEL_METADATA_UPLOAD_TIMEOUT_SECONDS,
+)
+from fundamental.models import (
+    FeatureImportanceMultipartMetadataResponse,
+    FitMultipartMetadataResponse,
+    MultipartUploadInfo,
+    PredictMultipartMetadataResponse,
+)
+from fundamental.utils.http import api_call
+
+logger = logging.getLogger(__name__)
+
+XType = Union[pd.DataFrame, np.ndarray]
+YType = Union[np.ndarray, pd.Series]
+
+
+# Version-compatible _check_n_features
+# In sklearn 1.6+, _check_n_features was moved to sklearn.utils.validation
+# In sklearn 1.5.x, it's a method on BaseEstimator
+try:
+    from sklearn.utils.validation import (
+        _check_n_features,  # pyright: ignore[reportAttributeAccessIssue]
+    )
+
+    def check_n_features_compat(estimator: BaseEstimator, X: XType, *, reset: bool) -> None:
+        """Version-compatible wrapper for sklearn's _check_n_features."""
+        _check_n_features(estimator, X, reset=reset)
+
+except ImportError:
+    # sklearn < 1.6.0: _check_n_features is a method on BaseEstimator
+    def check_n_features_compat(estimator: BaseEstimator, X: XType, *, reset: bool) -> None:
+        """Version-compatible wrapper for sklearn's _check_n_features."""
+        estimator._check_n_features(X, reset=reset)  # type: ignore[attr-defined]
+
+
+class FileUpload(BaseModel):
+    """File upload with metadata for multipart HTTP requests."""
+
+    file_name: str
+    file_type: str
+    content: bytes
+
+
+class PredictMetadata(BaseModel):
+    request_id: str
+    x_test_upload: MultipartUploadInfo
+
+
+def serialize_df_to_parquet_bytes(
+    data: Union[pd.DataFrame, np.ndarray, pd.Series],
+) -> bytes:
+    """Serialize data to Parquet-formatted bytes for HTTP transmission.
+
+    Raises
+    ------
+    ValueError
+        If data contains types that cannot be serialized to Parquet.
+    """
+    if isinstance(data, np.ndarray):
+        data = pd.DataFrame(data)
+    elif isinstance(data, pd.Series):
+        data = data.to_frame()
+
+    try:
+        buffer = BytesIO()
+        data.to_parquet(
+            buffer,
+            index=False,
+            engine="pyarrow",
+            write_statistics=False,
+        )
+        return buffer.getvalue()
+    except (palib.ArrowInvalid, palib.ArrowNotImplementedError) as e:
+        error_msg = str(e)
+        # Provide specific error messages for common issues
+        if "cannot mix struct and non-struct" in error_msg:
+            raise ValueError(
+                "Data contains mixed types that cannot be serialized. "
+                "Please ensure all columns have consistent types."
+            ) from e
+        if "Unsupported numpy type" in error_msg or "complex" in error_msg.lower():
+            raise ValueError(
+                "Data contains unsupported types (e.g., complex numbers). "
+                "Please convert to numeric or string types."
+            ) from e
+        # Generic message for other PyArrow errors
+        raise ValueError(
+            "Data cannot be serialized. Please ensure data contains only numeric or string types."
+        ) from e
+
+
+def to_httpx_post_file_format(
+    file_uploads: List[FileUpload],
+) -> Dict[str, tuple[str, bytes]]:
+    """Convert file uploads to httpx multipart format."""
+    files = {}
+    for f in file_uploads:
+        file = f"{f.file_name}.{f.file_type}"
+        files[f.file_name] = (file, f.content)
+    return files
+
+
+def validate_inputs_type(X: Any, y: Optional[Any] = None) -> None:  # noqa: ANN401
+    """
+    Validate X and y inputs match expected types.
+
+    Parameters
+    ----------
+    X : XType
+        Input features.
+    y : YType, optional
+        Target values.
+
+    Raises
+    ------
+    TypeError
+        If inputs don't match expected types.
+    """
+
+    def _validate_param(var: Any, expected_type: Any, param_name: str) -> None:  # noqa: ANN401
+        if var is not None:
+            if hasattr(expected_type, "__args__"):  # Union Type
+                valid_types = typing.get_args(expected_type)
+            else:
+                valid_types = (expected_type,)
+
+            if not isinstance(var, valid_types):
+                type_names = " or ".join(t.__name__ for t in valid_types)
+                raise TypeError(
+                    f"{param_name} must be {type_names}, got {type(var).__name__}. "
+                    f"Please convert your data to a supported format."
+                )
+
+    _validate_param(var=X, expected_type=XType, param_name="X")
+    _validate_param(var=y, expected_type=YType, param_name="y")
+
+
+def validate_data(
+    X: XType,
+    y: Optional[YType] = None,
+) -> None:
+    """
+    Validate data shapes and alignment.
+
+    Parameters
+    ----------
+    X : XType
+        Input features.
+    y : YType, optional
+        Target values.
+
+    Raises
+    ------
+    ValueError
+        If data is empty or misaligned.
+    """
+
+    if X.shape[0] == 0:
+        raise ValueError(
+            "Training data X cannot be empty. Please provide feature data with at least one sample."
+        )
+
+    if y is not None and len(y) == 0:
+        raise ValueError(
+            "Target data y cannot be empty. Please provide target values for training."
+        )
+
+    if y is not None and X.shape[0] != y.shape[0]:
+        raise ValueError(
+            f"Mismatched sample counts: X has {X.shape[0]} samples but y has "
+            f"{y.shape[0]} samples. Both must have the same number of samples."
+        )
+
+
+def _upload_data_in_parts(
+    data: bytes, upload_urls: List[str], part_size: int, client: BaseClient
+) -> List[Dict[str, Any]]:
+    parts = []
+    num_parts = len(upload_urls)
+
+    for part_number, upload_url in enumerate(upload_urls, start=1):
+        # Calculate the start and end positions for this part
+        start = (part_number - 1) * part_size
+        end = min(start + part_size, len(data))
+        part_data = data[start:end]
+
+        logger.info(f"Uploading part {part_number}/{num_parts} ({len(part_data)} bytes)")
+
+        response = api_call(
+            method="PUT",
+            full_url=upload_url,
+            client=client,
+            content=part_data,
+            timeout=DEFAULT_MODEL_METADATA_UPLOAD_TIMEOUT_SECONDS,
+        )
+
+        etag = response.headers.get("ETag", "").strip('"')
+        parts.append({"PartNumber": part_number, "ETag": etag})
+
+    return parts
+
+
+def _complete_multipart_upload(
+    trained_model_id: str,
+    file_type: str,
+    upload_id: str,
+    parts: List[Dict[str, Any]],
+    client: BaseClient,
+    request_id: Optional[str] = None,
+) -> None:
+    json_data = {
+        "trained_model_id": trained_model_id,
+        "file_type": file_type,
+        "upload_id": upload_id,
+        "parts": parts,
+    }
+
+    if request_id:
+        json_data["request_id"] = request_id
+
+    api_call(
+        method="POST",
+        full_url=client.config.get_full_complete_multipart_upload_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_COMPLETE_MULTIPART_UPLOAD_TIMEOUT_SECONDS,
+    )
+
+
+def _upload_and_complete_multipart(
+    data: bytes,
+    upload_info: MultipartUploadInfo,
+    trained_model_id: str,
+    file_type: str,
+    client: BaseClient,
+    request_id: Optional[str] = None,
+) -> None:
+    parts = _upload_data_in_parts(
+        data=data,
+        upload_urls=upload_info.upload_urls,
+        part_size=upload_info.part_size,
+        client=client,
+    )
+    _complete_multipart_upload(
+        trained_model_id=trained_model_id,
+        file_type=file_type,
+        upload_id=upload_info.upload_id,
+        parts=parts,
+        client=client,
+        request_id=request_id,
+    )
+
+
+def _set_client_trace_dict(data: Dict[str, Any], client: BaseClient) -> None:
+    trace_dict = {}
+    for k in ["trace_id", "span_id"]:
+        if k in data:
+            trace_dict[k] = data[k]
+    if trace_dict:
+        client._set_trace_dict(trace_dict)
+
+
+def create_fit_task_metadata(
+    x_train_size: int,
+    y_train_size: int,
+    client: BaseClient,
+) -> FitMultipartMetadataResponse:
+    json_data = {
+        "x_train_size": x_train_size,
+        "y_train_size": y_train_size,
+    }
+
+    client._set_trace_dict(None)
+    response = api_call(
+        method="POST",
+        full_url=client.config.get_full_fit_model_metadata_generate_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_MODEL_METADATA_GENERATE_TIMEOUT_SECONDS,
+    )
+    data = response.json()
+
+    _set_client_trace_dict(data, client)
+
+    return FitMultipartMetadataResponse(**data)
+
+
+def create_predict_task_metadata(
+    trained_model_id: str,
+    x_test_size: int,
+    client: BaseClient,
+) -> PredictMetadata:
+    json_data = {
+        "trained_model_id": trained_model_id,
+        "x_test_size": x_test_size,
+    }
+
+    client._set_trace_dict(None)
+    response = api_call(
+        method="POST",
+        full_url=client.config.get_full_predict_model_metadata_generate_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_MODEL_METADATA_GENERATE_TIMEOUT_SECONDS,
+    )
+    data = response.json()
+
+    _set_client_trace_dict(data, client)
+
+    # Adapt PredictMultipartMetadataResponse to PredictMetadata interface
+    api_response = PredictMultipartMetadataResponse(**data)
+    return PredictMetadata(
+        request_id=api_response.request_id,
+        x_test_upload=MultipartUploadInfo(
+            upload_id=api_response.upload_id,
+            upload_urls=api_response.upload_urls,
+            num_parts=api_response.num_parts,
+            part_size=api_response.part_size,
+        ),
+    )
+
+
+def create_feature_importance_task_metadata(
+    trained_model_id: str,
+    x_size: int,
+    client: BaseClient,
+) -> FeatureImportanceMultipartMetadataResponse:
+    json_data: Dict[str, Any] = {
+        "trained_model_id": trained_model_id,
+        "x_size": x_size,
+    }
+
+    client._set_trace_dict(None)
+    response = api_call(
+        method="POST",
+        full_url=client.config.get_full_feature_importance_model_metadata_generate_url(),
+        client=client,
+        json=json_data,
+        timeout=DEFAULT_MODEL_METADATA_GENERATE_TIMEOUT_SECONDS,
+    )
+    data = response.json()
+
+    _set_client_trace_dict(data, client)
+
+    return FeatureImportanceMultipartMetadataResponse(**data)
+
+
+def upload_fit_data(
+    X_serialized: bytes,
+    y_serialized: bytes,
+    metadata: FitMultipartMetadataResponse,
+    client: BaseClient,
+) -> None:
+    _upload_and_complete_multipart(
+        data=X_serialized,
+        upload_info=metadata.x_train_upload,
+        trained_model_id=metadata.trained_model_id,
+        file_type="x_train",
+        client=client,
+    )
+    _upload_and_complete_multipart(
+        data=y_serialized,
+        upload_info=metadata.y_train_upload,
+        trained_model_id=metadata.trained_model_id,
+        file_type="y_train",
+        client=client,
+    )
+
+
+def upload_predict_data(
+    X_serialized: bytes,
+    metadata: PredictMetadata,
+    trained_model_id: str,
+    client: BaseClient,
+) -> None:
+    _upload_and_complete_multipart(
+        data=X_serialized,
+        upload_info=metadata.x_test_upload,
+        trained_model_id=trained_model_id,
+        file_type="x_test",
+        client=client,
+        request_id=metadata.request_id,
+    )
+
+
+def upload_feature_importance_data(
+    X_serialized: bytes,
+    metadata: FeatureImportanceMultipartMetadataResponse,
+    trained_model_id: str,
+    client: BaseClient,
+) -> None:
+    _upload_and_complete_multipart(
+        data=X_serialized,
+        upload_info=metadata.x_upload,
+        trained_model_id=trained_model_id,
+        file_type="x_feature_importance",
+        client=client,
+        request_id=metadata.request_id,
+    )
+
+
+def download_result_from_url(
+    download_url: str,
+    client: BaseClient,
+    timeout: int = DEFAULT_DOWNLOAD_RESULT_TIMEOUT_SECONDS,
+) -> Any:  # noqa: ANN401
+    response = api_call(
+        method="GET",
+        full_url=download_url,
+        client=client,
+        timeout=timeout,
+    )
+    return response.json()