mplang-nightly 0.1.dev268__py3-none-any.whl → 0.1.dev270__py3-none-any.whl

mplang/v1/utils/spu_utils.py

@@ -1,130 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""SPU-related utility functions for protocol and field type conversion."""
-
-import spu.libspu as libspu
-
-# Global mappings for SPU protocol and field type conversion
-SPU_PROTOCOL_MAPPING = {
-    "REF2K": libspu.ProtocolKind.REF2K,
-    "SEMI2K": libspu.ProtocolKind.SEMI2K,
-    "ABY3": libspu.ProtocolKind.ABY3,
-    "CHEETAH": libspu.ProtocolKind.CHEETAH,
-    "SECURENN": libspu.ProtocolKind.SECURENN,
-}
-
-SPU_FIELD_MAPPING = {
-    "FM32": libspu.FieldType.FM32,
-    "FM64": libspu.FieldType.FM64,
-    "FM128": libspu.FieldType.FM128,
-}
-
-
-def parse_protocol(protocol: str | int) -> libspu.ProtocolKind:
-    """Parse SPU protocol from string or integer to ProtocolKind enum.
-
-    Args:
-        protocol: Protocol specification as string (e.g., "SEMI2K") or integer.
-
-    Returns:
-        libspu.ProtocolKind: The corresponding protocol enum.
-
-    Raises:
-        ValueError: If the protocol is invalid.
-
-    Examples:
-        >>> parse_spu_protocol("SEMI2K")
-        ProtocolKind.SEMI2K
-        >>> parse_spu_protocol(2)
-        ProtocolKind.SEMI2K
-    """
-    if isinstance(protocol, str):
-        if protocol not in SPU_PROTOCOL_MAPPING:
-            raise ValueError(
-                f"Invalid SPU protocol: {protocol}. "
-                f"Valid protocols are: {list(SPU_PROTOCOL_MAPPING.keys())}"
-            )
-        return SPU_PROTOCOL_MAPPING[protocol]
-    else:
-        # Assume it's an integer, validate it
-        try:
-            spu_protocol = libspu.ProtocolKind(protocol)
-            # Check if it's a valid enum value (not ???)
-            if spu_protocol.name == "???":
-                raise ValueError(f"Invalid SPU protocol value: {protocol}")
-            return spu_protocol
-        except TypeError as exc:
-            raise ValueError(
-                f"Invalid SPU protocol: {protocol}. "
-                f"Must be a valid protocol string or integer."
-            ) from exc
-
-
-def parse_field(field: str | int) -> libspu.FieldType:
-    """Parse SPU field type from string or integer to FieldType enum.
-
-    Args:
-        field: Field type specification as string (e.g., "FM64") or integer.
-
-    Returns:
-        libspu.FieldType: The corresponding field type enum.
-
-    Raises:
-        ValueError: If the field type is invalid.
-
-    Examples:
-        >>> parse_spu_field("FM64")
-        FieldType.FM64
-        >>> parse_spu_field(2)
-        FieldType.FM64
-    """
-    if isinstance(field, str):
-        if field not in SPU_FIELD_MAPPING:
-            raise ValueError(
-                f"Invalid SPU field type: {field}. "
-                f"Valid field types are: {list(SPU_FIELD_MAPPING.keys())}"
-            )
-        return SPU_FIELD_MAPPING[field]
-    else:
-        # Assume it's an integer, validate it
-        try:
-            spu_field = libspu.FieldType(field)
-            # Check if it's a valid enum value
-            if spu_field.name == "???":
-                raise ValueError(f"Invalid SPU field type value: {field}")
-            return spu_field
-        except TypeError as exc:
-            raise ValueError(
-                f"Invalid SPU field type: {field}. "
-                f"Must be a valid field type string or integer."
-            ) from exc
-
-
-def list_protocols() -> list[str]:
-    """Get list of valid SPU protocol names.
-
-    Returns:
-        List of valid protocol names as strings.
-    """
-    return list(SPU_PROTOCOL_MAPPING.keys())
-
-
-def list_fields() -> list[str]:
-    """Get list of valid SPU field type names.
-
-    Returns:
-        List of valid field type names as strings.
-    """
-    return list(SPU_FIELD_MAPPING.keys())

mplang/v1/utils/table_utils.py

@@ -1,185 +0,0 @@
-# Copyright 2025 Ant Group Co., Ltd.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from __future__ import annotations
-
-import io
-from typing import Any
-
-import pyarrow as pa
-import pyarrow.csv as pa_csv
-import pyarrow.orc as pa_orc
-import pyarrow.parquet as pa_pq
-
-from mplang.v1.core.table import TableLike
-
-__all__ = ["decode_table", "encode_table", "read_table", "write_table"]
-
-
-def _parse_kwargs(kwargs: dict[str, Any], keys: list[str]) -> dict[str, Any] | None:
-    if not kwargs:
-        return None
-
-    return {key: kwargs[key] for key in keys if key in kwargs}
-
-
-_csv_read_option_keys = [
-    "skip_rows",
-    "skip_rows_after_names",
-    "column_names",
-    "autogenerate_column_names",
-    "encoding",
-]
-_csv_parse_option_keys = [
-    "delimiter",
-    "quote_char",
-    "double_quote",
-    "escape_char",
-    "newlines_in_values",
-    "ignore_empty_lines",
-]
-_csv_convert_option_keys = [
-    "check_utf8",
-    "column_types",
-    "null_values",
-    "true_values",
-    "false_values",
-    "decimal_point",
-    "strings_can_be_null",
-    "quoted_strings_can_be_null",
-    "include_columns",
-    "include_missing_columns",
-    "auto_dict_encode",
-    "auto_dict_max_cardinality",
-    "timestamp_parsers",
-]
-
-
-def read_table(
-    source: Any,
-    format: str = "parquet",
-    columns: list[str] | None = None,
-    **kwargs: Any,
-) -> pa.Table:
-    """Read data from a file and return a PyArrow table.
-
-    Args:
-        source: The source to read data from (file path, file-like object, etc.)
-        format: The format of the data source ("parquet", "csv", or "orc")
-        columns: List of column names to read (None means all columns)
-        **kwargs: Additional keyword arguments passed to the underlying reader
-
-    Returns:
-        A PyArrow Table containing the data from the source
-
-    Raises:
-        ValueError: If an unsupported format is specified
-    """
-    match format:
-        case "csv":
-            if columns:
-                kwargs["include_columns"] = columns
-            read_args = _parse_kwargs(kwargs, _csv_read_option_keys)
-            parse_args = _parse_kwargs(kwargs, _csv_parse_option_keys)
-            convert_args = _parse_kwargs(kwargs, _csv_convert_option_keys)
-
-            read_opts = pa_csv.ReadOptions(**read_args) if read_args else None
-            parse_opts = pa_csv.ParseOptions(**parse_args) if parse_args else None
-            conv_opts = pa_csv.ConvertOptions(**convert_args) if convert_args else None
-            return pa_csv.read_csv(
-                source,
-                read_options=read_opts,
-                parse_options=parse_opts,
-                convert_options=conv_opts,
-            )
-        case "orc":
-            return pa_orc.read_table(source, columns=columns, **kwargs)
-        case "parquet":
-            return pa_pq.read_table(source, columns=columns, **kwargs)
-        case _:
-            raise ValueError(f"unsupported data format. {format}")
-
-
-def write_table(
-    data: TableLike,
-    where: Any,
-    format: str = "parquet",
-    **kwargs: Any,
-) -> None:
-    """Write a table-like object to a file in the specified format.
-
-    Args:
-        data: The table-like object to write (PyArrow Table or other compatible format)
-        where: The destination to write to (file path, file-like object, etc.)
-        format: The format to write the data in ("parquet", "csv", or "orc")
-        **kwargs: Additional keyword arguments passed to the underlying writer
-
-    Returns:
-        None
-
-    Raises:
-        ValueError: If the table has no columns or an unsupported format is specified
-    """
-    # Convert data to PyArrow Table if needed
-    table = data if isinstance(data, pa.Table) else pa.table(data)
-    if len(table.column_names) == 0:
-        raise ValueError("Cannot convert Table with no columns.")
-
-    match format:
-        case "csv":
-            options = pa_csv.WriteOptions(**kwargs) if kwargs else None
-            pa_csv.write_csv(table, where, write_options=options)
-        case "orc":
-            pa_orc.write_table(table, where, **kwargs)
-        case "parquet":
-            pa_pq.write_table(table, where, **kwargs)
-        case _:
-            raise ValueError(f"unsupported data format. {format}")
-
-
-def decode_table(
-    data: bytes,
-    format: str = "parquet",
-    columns: list[str] | None = None,
-    **kwargs: Any,
-) -> pa.Table:
-    """Decode a bytes object into a PyArrow table.
-
-    Args:
-        data: The bytes object containing the encoded table data
-        format: The format of the encoded data ("parquet", "csv", or "orc")
-        columns: List of column names to decode (None means all columns)
-        **kwargs: Additional keyword arguments passed to the underlying reader
-
-    Returns:
-        A PyArrow Table decoded from the bytes data
-    """
-    buffer = io.BytesIO(data)
-    return read_table(buffer, format=format, columns=columns, **kwargs)
-
-
-def encode_table(data: TableLike, format: str = "parquet", **kwargs: Any) -> bytes:
-    """Encode a table-like object into bytes.
-
-    Args:
-        data: The table-like object to encode (PyArrow Table or other compatible format)
-        format: The format to encode the data in ("parquet", "csv", or "orc")
-        **kwargs: Additional keyword arguments passed to the underlying writer
-
-    Returns:
-        Bytes object containing the encoded table data
-    """
-    buffer = io.BytesIO()
-    write_table(data, buffer, format, **kwargs)
-    return buffer.getvalue()