decentriq_util 0.1.0__tar.gz

decentriq_util-0.1.0/PKG-INFO

@@ -0,0 +1,16 @@
+Metadata-Version: 2.1
+Name: decentriq_util
+Version: 0.1.0
+Summary: Utility functions for Python code run in Decentriq workers
+Author: hafenr
+Author-email: robin.hafen@gmail.com
+Requires-Python: >=3.9,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: pandas (>=1.4.4)
+Requires-Dist: protobuf (>=4.21,<5.0)
+Requires-Dist: psutil (>=5.9.6,<6.0.0)
+Requires-Dist: pyspark (>=3.4.0,<4.0.0)

decentriq_util-0.1.0/decentriq_util/__init__.py

@@ -0,0 +1,29 @@
+"""
+The `decentriq_util` library provides helpers for common operations,
+not only when automating your workflows with the `decentriq_platform` package,
+but also when writing your own Python computations inside Data Clean Rooms.
+
+It is available inside the computation container for Python scripts
+and can be imported just like any other library. For example:
+```python
+import decentriq_util
+data = decentriq_util.sql.read_sql_data_from_dir("/path/to/computation")
+```
+"""
+from . import sql
+from . import proto
+from . import python
+from . import error
+from . import spark
+
+from .sql import (
+    read_tabular_data,
+    write_tabular_data
+)
+
+__pdoc__ = {
+    "sql": True,
+    "proto": False,
+    "python": False,
+    "error": False,
+}

decentriq_util-0.1.0/decentriq_util/error.py

@@ -0,0 +1,131 @@
+from typing import Callable, List
+
+import traceback
+import sys
+import os
+
+
+def _strip_away_nix_store(path: str) -> str:
+    """
+    Replace '/nix/store/hash-derivation-name' with '/system/' for better
+    readability of the stack trace.
+    """
+    if path.startswith("/nix/store"):
+        return "/system/" + "/".join(path.split("/")[4:])
+    else:
+        return path
+
+
+class SafeError(Exception):
+    """
+    Base class for errors that can be turned into safe error messages
+    that can be displayed to the user without risking the leakage of
+    sensitive information.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def safe_str(self) -> str:
+        """
+        Return a safe string representation of this exception
+        that does not include any sensitive data.
+
+        This method must be implemented explicitly.
+        """
+        raise Exception("Not implemented")
+
+
+class CatchSafeErrorContext:
+    """
+    Context manager to handle error messages and writing them to a pre-defined
+    location on the file system.
+    The constructed stack track is optionally cleaned from any sensitive
+    information that might be contained in the error message.
+    """
+    def __init__(
+            self,
+            should_clean_traceback: bool,
+            output_file: str,
+            should_exit: bool,
+            transform_library_path: List[Callable[[str], str]],
+    ):
+        self.should_clean_traceback = should_clean_traceback
+        self.output_file = output_file
+        self.should_exit = should_exit
+        self.transform_library_path = transform_library_path + [_strip_away_nix_store]
+
+
+    def __enter__(self):
+        return self
+
+
+    def __exit__(self, exception_type, exception, tb_):
+        if exception:
+            if exception_type == SystemExit and exception.code == 0:
+                # Don't produce a traceback when sys.exit(0) is used.
+                return
+
+            should_capture_local_variables = not self.should_clean_traceback
+            error_class_module = exception.__class__.__module__
+            ignored_modules = set(["builtins", "__main__"])
+            if error_class_module and error_class_module not in ignored_modules:
+                error_qualified_class_name =\
+                    error_class_module + "." + exception.__class__.__name__
+
+            else:
+                error_qualified_class_name = exception.__class__.__name__
+
+            # Skip the first frame as it is the one in the wrapper script
+            tbe = traceback.TracebackException.from_exception(
+                exception, capture_locals=should_capture_local_variables
+            )
+            lines = ["Traceback (most recent call last):"]
+            for frame in tbe.stack[1:]:
+                filename = frame.filename
+                # Apply all configured transformer functions to the file paths
+                # to make them more readable.
+                for transform_fn in self.transform_library_path:
+                    filename = transform_fn(filename)
+                lines.append(
+                    f"  File \"{filename}\","
+                    f" line {frame.lineno}, in {frame.name}\n    {frame.line}"
+                )
+
+            if self.should_clean_traceback:
+                # Even if we are supposed to provide a safe error, if it's one of our
+                # own errors, we can still construct an (informative) safe error
+                # message.
+                if isinstance(exception, SafeError):
+                    stringified_error =\
+                            error_qualified_class_name + ": " + exception.safe_str()
+                else:
+                    stringified_error = error_qualified_class_name +\
+                        ": (error message removed due to data privacy)"
+            else:
+                stringified_error = error_qualified_class_name + ": " + str(exception)
+
+            lines.append(stringified_error)
+            safe_traceback = "\n".join(lines)
+
+            with open(self.output_file, "w") as f:
+                f.write(safe_traceback)
+
+            if self.should_exit:
+                sys.exit(1)
+            else:
+                # Do not propagate error
+                return True
+
+
+def catch_safe_error(
+        transform_library_path: List[Callable[[str], str]] = []
+) -> CatchSafeErrorContext:
+    """Construct a CatchSafeErrorContext with default values"""
+    should_clean_traceback = \
+        "INCLUDE_CONTAINER_LOGS_ON_ERROR" not in os.environ or os.environ["INCLUDE_CONTAINER_LOGS_ON_ERROR"] != "true"
+    return CatchSafeErrorContext(
+        should_clean_traceback=should_clean_traceback,
+        should_exit=True,
+        output_file="/internal/safe_error",
+        transform_library_path=transform_library_path
+    )

decentriq_util-0.1.0/decentriq_util/proto/__init__.py

@@ -0,0 +1,24 @@
+from typing import Tuple
+from google.protobuf.message import Message
+from google.protobuf.internal.encoder import _VarintBytes # type: ignore
+from google.protobuf.internal.decoder import _DecodeVarint32 # type: ignore
+
+
+def parse_length_delimited(serialized_bytes: bytes, deserialized_object: Message) -> int:
+    """
+    Decodes a binary content and returns a parsed object of a given type
+    """
+    res: Tuple[int, int] = _DecodeVarint32(serialized_bytes, 0)
+    message_length, offset = res
+    end_offset = offset + message_length
+    deserialized_object.ParseFromString(bytes(serialized_bytes[offset:end_offset]))
+    return end_offset
+
+
+def serialize_length_delimited(message_object: Message) -> bytes:
+    """
+    Encodes an object of a given type to binary
+    """
+    serialized: bytes = _VarintBytes(message_object.ByteSize())
+    serialized += message_object.SerializeToString()
+    return serialized

decentriq_util-0.1.0/decentriq_util/python/__init__.py

@@ -0,0 +1,21 @@
+import traceback
+import datetime
+
+def log_entry(record: str):
+    """
+    Logs an entry in the user log file to be displayed at the end of the execution.
+    """
+    current_time = datetime.datetime.now(datetime.timezone.utc)
+    with open("/internal/logs", "a+") as f:
+        f.write(f"LOG {current_time}: {record}\n")
+        
+def log_stack_trace(e: Exception):
+    """
+    Allows to add a striped stack strace for an Exception e in the user log file.
+    The striped stack trace will not contain any environment variable or secrets
+    in it.
+    """
+    tb = traceback.TracebackException.from_exception(e, capture_locals=False)
+    safe_traceback = "".join([s for s in tb.format()][:-1])
+    error_qualified_class_name = e.__class__.__module__ + "." + e.__class__.__name__
+    log_entry(f"Encountered error: {error_qualified_class_name}, Traceback:\n{safe_traceback}")

decentriq_util-0.1.0/decentriq_util/spark/__init__.py

@@ -0,0 +1,219 @@
+import tempfile
+from typing import Optional, Tuple, Any
+from contextlib import contextmanager
+import multiprocessing
+import math
+import os
+import shutil
+import psutil
+
+from pyspark.sql import SparkSession
+
+
+def _get_cgroup_memory() -> int:
+    max_memory_str = None
+    if os.path.isfile("/sys/fs/cgroup/memory.max"):
+        with open("/sys/fs/cgroup/memory.max", "r") as f:
+            max_memory_str = f.read().strip()
+    elif os.path.isfile("/sys/fs/cgroup/memory/memory.limit_in_bytes"):
+        with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as f:
+            max_memory_str = f.read().strip()
+
+    max_memory = None
+    if max_memory_str == "max":
+        # Fallback to available virtual memory size
+        max_memory = psutil.virtual_memory().available
+    elif max_memory_str is not None:
+        try:
+            max_memory = int(max_memory_str)
+        except ValueError:
+            pass
+
+    if max_memory is not None:
+        return max_memory
+    else:
+        print("Unable to determine available memory from cgroup, assuming 4G")
+        return 4 * 1024 * 1024 * 1024
+
+
+def _get_num_partitions(
+        paths: list[str],
+        available_memory: int,
+        cpu_count: int,
+) -> int:
+    """Determine the number of partitions s.t. each partition is as big as possible
+    while still making sure all cores are occupied."""
+    def get_file_or_dir_size(path):
+        if os.path.isfile(path):
+            size = os.path.getsize(path)
+        elif os.path.isdir(path):
+            total_size = 0
+            for dirpath, dirnames, filenames in os.walk(path):
+                for f in filenames:
+                    fp = os.path.join(dirpath, f)
+                    total_size += os.path.getsize(fp)
+            size = total_size
+        else:
+            raise ValueError(f"The path {path} is neither a file nor a directory.")
+        return size
+
+    file_size = max([get_file_or_dir_size(path) for path in paths])
+
+    # Don't make the partitions smaller than 36MiB
+    megabyte_in_bytes = 1024 * 1024
+    min_partition_size = 32 * megabyte_in_bytes
+
+    # Try to make partitions as big as possible s.t. each CPU is saturated,
+    # while accounting for 20% overhead in spark.
+    target_partition_size = max(
+        int(0.8 * available_memory / cpu_count) // megabyte_in_bytes * megabyte_in_bytes,
+        min_partition_size
+    )
+
+    while (
+        # If the file is small enough that it cannot be split across all cores,
+        # divide the partition size further.
+        int(file_size / target_partition_size) < cpu_count
+    ):
+        # Don't go smaller than the minimum parition size
+        if target_partition_size / 2 < min_partition_size:
+            break
+        target_partition_size = target_partition_size / 2
+
+    print(f"File size: {file_size}")
+    print(f"Target partition size: {target_partition_size}")
+
+    num_partitions = int(math.ceil(file_size / target_partition_size))
+
+    return num_partitions
+
+
+def _determine_optimal_spark_settings(
+        files: list[str],
+        heap_size_extra_room: Optional[int] = 512 * 1024 * 1024
+) -> list[Tuple[str, str]]:
+    def get_spark_memory() -> int:
+        """Determine the amount of memory available to spark"""
+        cgroup_memory = _get_cgroup_memory()
+        spark_memory = cgroup_memory - (heap_size_extra_room or 512 * 1024 * 1024)
+        return spark_memory
+
+    cpu_count = multiprocessing.cpu_count()
+    spark_memory = get_spark_memory()
+    num_partitions = _get_num_partitions(files, spark_memory, cpu_count)
+
+    print(f"CPU count: {cpu_count}")
+    print(f"Spark memory: {spark_memory}")
+    print(f"Num partitions: {num_partitions}")
+
+    settings = [
+        ("spark.sql.shuffle.partitions", str(num_partitions)),
+        ("spark.default.parallelism", str(num_partitions)),
+        ("spark.driver.cores", str(cpu_count)),
+    ]
+    print("Spark settings:\n" + f"\n".join([str(x) for x in settings]))
+
+    return settings
+
+
+def _create_spark_session(
+        name: str = "local_spark_session",
+        parallelism: int = 8,
+        heap_size: Optional[int] = None,
+        heap_size_extra_room: int = 512 * 1024 * 1024,
+        config: list[Tuple[str, Any]] = [],
+        java_tmp_dir = "/output",
+) -> SparkSession:
+    """
+    :param name: The name of the spark session.
+    :param parallelism: The size of the executor pool computing internal spark tasks.
+    :param heap_size: If set then this determines the JVM heap size. Note that not all the heap size is used by spark.
+    :param heap_size_extra_room: If heap_size is unset then it is determined by the available memory in the current
+        cgroup. Then heap_size_extra_room is subtracted to arrive at the final heap size. This number should account for
+        the python runtime itself as well as other structures that may be needed during spark compute.
+    :param config: Additional config settings to set on the spark config.
+    :param java_tmp_dir: Location for the JVM to store temp files.
+    :return: The spark session.
+    """
+    os.environ["SPARK_EXECUTOR_POOL_SIZE"] = str(parallelism)
+    os.environ["JDK_JAVA_OPTIONS"] = f'-Duser.home=/output -Djava.io.tmpdir="{java_tmp_dir}"'
+    spark_memory = None
+    if heap_size is None:
+        cgroup_memory = _get_cgroup_memory()
+        spark_memory = cgroup_memory - heap_size_extra_room
+    else:
+        spark_memory = heap_size
+
+    spark_memory_4096 = (spark_memory // 4096) * 4096
+    if "spark.driver.memory" not in dict(config):
+        config.append(("spark.driver.memory", str(spark_memory_4096)))
+
+    ss = (
+        SparkSession.builder
+        .appName(name)
+        .master("local[*]")
+    )
+    for (key, value) in config:
+        ss = ss.config(key, value)
+    return ss.getOrCreate()
+
+
+@contextmanager
+def spark_session(temp_dir: str = "/scratch", input_files: list[str] = None, **kwargs):
+    """
+    Create a spark session and configure it according to the enclave environment.
+
+    **Parameters**:
+    - `temp_dir`: Where to store temporary data such as persisted data frames
+      or shuffle data.
+    - `input_files`: A list of input files on the basis of which the partition
+      size is determined.
+
+    **Example**:
+
+    ```python
+    import decentriq_util as dq
+
+    # Path to a potentially very large file
+    input_csv_path = "/input/my_file.csv"
+
+    # Automatically create and configure a spark session and
+    # make sure it's being stopped at the end.
+    with dq.spark.spark_session(input_files=[input_csv_path]) as ss:
+        # Read from a CSV file
+        df = ss.read.csv(input_csv_path, header=False).cache()
+
+        # Perform any pyspark transformations
+        print(f"Original number of rows: {df.count()}")
+        result_df = df.limit(100)
+
+        # Write the result to an output file
+        result_df.write.parquet("/output/my_file.parquet")
+    ```
+    """
+    with tempfile.TemporaryDirectory(dir=temp_dir, prefix="java-") as java_tmp:
+        with tempfile.TemporaryDirectory(dir=temp_dir, prefix="spark-") as spark_tmp:
+            config = kwargs.get("config", [])
+            config_dict = dict(config)
+            if input_files is not None and input_files:
+                optimal_settings = _determine_optimal_spark_settings(
+                    input_files, heap_size_extra_room=kwargs.get("heap_size_extra_room")
+                )
+                for key, value in optimal_settings:
+                    if key not in config_dict:
+                        config.append((key, value))
+            if "spark.local.dir" not in config_dict:
+                config.append(
+                    ("spark.local.dir", spark_tmp)
+                )
+            kwargs["config"] = config
+            if "java_tmp_dir" not in kwargs:
+                kwargs["java_tmp_dir"] = java_tmp
+            ss = _create_spark_session(**kwargs)
+            try:
+                yield ss
+            finally:
+                try:
+                    ss.stop()
+                except:
+                    pass

decentriq_util-0.1.0/decentriq_util/sql/__init__.py

@@ -0,0 +1,266 @@
+import typing
+from typing import Tuple, List
+import os
+import numpy
+import pandas as pd
+
+from .proto.compute_sql_pb2 import (
+    TableSchema,
+    PrimitiveType,
+    NamedColumn,
+    ColumnType
+)
+from ..proto import (
+    parse_length_delimited,
+    serialize_length_delimited
+)
+
+__pdoc__ = {
+    "proto": False,
+}
+
+
+datatype_mapping = {
+    PrimitiveType.INT64: pd.Int64Dtype(),
+    PrimitiveType.STRING: str,
+    PrimitiveType.FLOAT64: numpy.float64,
+}
+
+
+datatype_mapping_reverse = {
+    "Int64": PrimitiveType.INT64,
+    "int64": PrimitiveType.INT64,
+    "object": PrimitiveType.STRING,
+    "float64": PrimitiveType.FLOAT64
+}
+
+
+DATASET_FILE_NAME = "dataset.csv"
+SCHEMA_FILE_NAME = "types"
+
+
+def _get_numpy_dtypes_for_table_schema_object(
+        schema,
+        use_numpy_int_type: bool = False
+) -> List[Tuple[str, numpy.dtype]]:
+    ix = 1
+    columns = []
+    for col in schema.namedColumns:
+        # This matches R convention, unnamed columns get a position
+        name = col.name if col.name else f"V{ix}"
+        sql_type = col.columnType.primitiveType
+        if sql_type == PrimitiveType.INT64 and use_numpy_int_type:
+            dtype = numpy.int64
+        else:
+            dtype = datatype_mapping[sql_type]
+        columns.append((name, dtype))
+        ix += 1
+
+    return columns
+
+
+def _get_numpy_dtypes_for_table_schema(
+        schema_path: str,
+        use_numpy_int_type: bool = False
+) -> List[Tuple[str, numpy.dtype]]:
+    with open(schema_path, "rb") as f:
+        schema = TableSchema()
+        parse_length_delimited(f.read(), schema)
+
+    return _get_numpy_dtypes_for_table_schema_object(schema, use_numpy_int_type)
+
+
+def _get_table_schema_for_data_frame(df: pd.DataFrame) -> TableSchema:
+    dtypes_as_string = [str(dtype) for dtype in df.dtypes]
+    proto_columns = []
+    for name, dtype in zip(list(df.columns), dtypes_as_string):
+        primitive_type = datatype_mapping_reverse[dtype]
+        column = NamedColumn(
+            name=name,
+            columnType=ColumnType(
+                primitiveType=primitive_type,
+                nullable=True
+            )
+        )
+        proto_columns.append(column)
+    return TableSchema(
+        namedColumns=proto_columns
+    )
+
+
+def _assert_files_do_exist(dataset_file: str, schema_file: str) -> None:
+    if not os.path.exists(dataset_file):
+        raise Exception(
+            f"Cannot read data from {dataset_file}, file containing data"
+            f" {dataset_file} does not exist!"
+        )
+    if not os.path.exists(schema_file):
+        raise Exception(
+            f"Cannot read data from {dataset_file}, file encoding schema"
+            f" {schema_file} does not exist!"
+        )
+
+
+def _assert_files_dont_exist(dataset_file: str, schema_file: str):
+    if os.path.exists(dataset_file):
+        raise Exception(
+            "Cannot write data to {path}, file containing data"
+            f" {dataset_file} already exists!"
+        )
+    if os.path.exists(schema_file):
+        raise Exception(
+            "Cannot write data to {path}, file encoding schema"
+            f" {schema_file} already exists!"
+        )
+
+
+def read_tabular_data(path: str) -> pd.DataFrame:
+    """
+    Read data from a tabular input node or as written by an SQL computation.
+
+    The input file must not contain a header row!
+
+    Empty values for columns of type TEXT are replaced with empty strings.
+
+    Empty values in columns of type INT64 are read as the "pandas.NA" value.
+    For more information, please refer to the
+    [pandas documentation](https://pandas.pydata.org/pandas-docs/dev/user_guide/integer_na.html).
+    """
+    return read_sql_data_from_dir(path)
+
+
+def read_sql_data_from_dir(path: str) -> pd.DataFrame:
+    """
+    Note: this function is deprecated, use the function `read_tabular_data`.
+
+    This function uses the previous `read_sql_data` function to convert
+    the standard format of validated tabular input datasets and the output
+    of SQL computations into a pandas dataframe, automatically importing the
+     headers the are contained in the `types` file as well. One example below:
+    ```python
+    import decentriq_util.sql as dqu
+    import pandas as pd
+
+    dataframe_1 = dqu.read_sql_data_from_dir("/input/mySQLComputation")
+    dataframe_2 = dqu.read_sql_data_from_dir("/input/myTabularDataset")
+    ```
+    """
+    dataset_file = os.path.join(path, DATASET_FILE_NAME)
+    schema_file = os.path.join(path, SCHEMA_FILE_NAME)
+    _assert_files_do_exist(dataset_file, schema_file)
+    return read_sql_data(dataset_file, schema_file)
+
+
+def write_tabular_data(df: pd.DataFrame, path: str = "/output"):
+    """
+    Write a Pandas dataframe into the given directory in a format compatible
+    with downstream SQL computation. This function allows you to use the output of
+    a python computation as input to a SQL computation.
+    For example:
+
+    ```python
+    import decentriq_util.sql as dqu
+    import pandas as pd
+
+    my_dataframe = pd.DataFrame()
+    ...
+    dqu.write_tabular_data(my_dataframe)
+    ```
+    """
+    write_sql_data_to_dir(df, path)
+
+
+def write_sql_data_to_dir(df: pd.DataFrame, path: str):
+    """
+    Note: this function is deprecated, use the function `write_tabular_data`.
+
+    Write a Pandas dataframe into the given directory in a format compatible
+    with downstream SQL computations. This function allows you to use the output of
+    a python computation as input to a SQL computation.
+    To do so, you will need to use `/output` as path.
+
+    For example:
+    ```python
+    import decentriq_util.sql as dqu
+    import pandas as pd
+
+    my_dataframe = pd.DataFrame()
+    ...
+    dqu.write_sql_data_to_dir(my_dataframe, "/output")
+    ```
+    """
+    dataset_file = os.path.join(path, DATASET_FILE_NAME)
+    schema_file = os.path.join(path, SCHEMA_FILE_NAME)
+    _assert_files_dont_exist(dataset_file, schema_file)
+    return write_sql_data(df, dataset_file, schema_file)
+
+
+def read_sql_data(data_path: str, schema_path: str) -> pd.DataFrame:
+    """
+    Read data as output by the SQL validation node.
+
+    The input file must not contain a header row!
+
+    Empty values for columns of type TEXT are replaced with empty strings.
+
+    Empty values in columns of type INT64 are read as the "pandas.NA" value.
+    For more information, please refer to the
+    [pandas documentation](https://pandas.pydata.org/pandas-docs/dev/user_guide/integer_na.html).
+    """
+    _assert_files_do_exist(data_path, schema_path)
+    columns = _get_numpy_dtypes_for_table_schema(schema_path)
+    dtype_per_column = {
+        ix: column_type
+        for ix, (_, column_type) in enumerate(columns)
+    }
+    try:
+        # First read just the data.
+        # Tell the type linter that the value is in fact a DataFrame
+        # as it could be a TextParser object depending on whether either
+        # `chunksize` or `iterator` arguments are specified.
+        original_data = typing.cast(
+            pd.DataFrame,
+            pd.read_csv(
+                data_path,
+                index_col=None,
+                header=None, # sql verifier doesn't write headers
+                dtype=str    # treat everything as string at first
+            )
+        )
+        # Read the schema file that tells us about the data types.
+        # Replace NaN values with empty strings for string-type columns.
+        string_column_replacements =\
+            {ix: "" for ix, dtype in dtype_per_column.items() if dtype == str}
+        original_data.fillna(string_column_replacements, inplace=True)
+        # Check whether the DF is empty now (even if it wasn't in the beginning)
+        # If yes, return an empty DF.
+        if original_data.shape[0] == 0:
+            return pd.DataFrame(
+                { name: pd.Series(dtype=dtype) for name, dtype in columns }
+            )
+        else:
+            # Convert columns to the correct dtype
+            for ix, dtype in dtype_per_column.items():
+                if original_data[ix].dtype != dtype:
+                    original_data[ix] = original_data[ix].astype(dtype)
+    except pd.errors.EmptyDataError:
+        return pd.DataFrame(
+            { name: pd.Series(dtype=dtype) for name, dtype in columns }
+        )
+    header = [column_name for column_name, _ in columns]
+    original_data.columns = header
+
+    return original_data
+
+
+def write_sql_data(df: pd.DataFrame, data_path: str, schema_path: str):
+    """
+    Write a Pandas dataframe to the given path, and store its schema including column names at
+    `schema_path` with the protobuf standard used in the cross-workers interactions.
+    """
+    _assert_files_dont_exist(data_path, schema_path)
+    schema = _get_table_schema_for_data_frame(df)
+    schema_serialized = serialize_length_delimited(schema)
+    df.to_csv(data_path, index=False, header=False, sep=",")
+    with open(schema_path, "wb") as f:
+        f.write(schema_serialized)

decentriq_util-0.1.0/decentriq_util/sql/proto/__init__.py

@@ -0,0 +1,11 @@
+from .compute_sql_pb2 import (
+    TableSchema,
+    NamedColumn,
+    ColumnType,
+    PrimitiveType,
+    SqlWorkerConfiguration,
+    ValidationConfiguration,
+    ComputationConfiguration,
+    PrivacySettings
+)
+

decentriq_util-0.1.0/decentriq_util/sql/proto/compute_sql_pb2.py

@@ -0,0 +1,112 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: compute_sql.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import message as _message
+from google.protobuf import reflection as _reflection
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x11\x63ompute_sql.proto\x12\x0b\x63ompute_sql\"\xa3\x01\n\x16SqlWorkerConfiguration\x12:\n\nvalidation\x18\x01 \x01(\x0b\x32$.compute_sql.ValidationConfigurationH\x00\x12<\n\x0b\x63omputation\x18\x02 \x01(\x0b\x32%.compute_sql.ComputationConfigurationH\x00\x42\x0f\n\rconfiguration\"H\n\x17ValidationConfiguration\x12-\n\x0btableSchema\x18\x01 \x01(\x0b\x32\x18.compute_sql.TableSchema\"=\n\x0bTableSchema\x12.\n\x0cnamedColumns\x18\x01 \x03(\x0b\x32\x18.compute_sql.NamedColumn\"V\n\x0bNamedColumn\x12\x11\n\x04name\x18\x01 \x01(\tH\x00\x88\x01\x01\x12+\n\ncolumnType\x18\x02 \x01(\x0b\x32\x17.compute_sql.ColumnTypeB\x07\n\x05_name\"Q\n\nColumnType\x12\x31\n\rprimitiveType\x18\x01 \x01(\x0e\x32\x1a.compute_sql.PrimitiveType\x12\x10\n\x08nullable\x18\x02 \x01(\x08\"\x95\x01\n\x18\x43omputationConfiguration\x12\x14\n\x0csqlStatement\x18\x01 \x01(\t\x12\x35\n\x0fprivacySettings\x18\x02 \x01(\x0b\x32\x1c.compute_sql.PrivacySettings\x12,\n\x0b\x63onstraints\x18\x03 \x03(\x0b\x32\x17.compute_sql.Constraint\"2\n\x0fPrivacySettings\x12\x1f\n\x17minAggregationGroupSize\x18\x01 \x01(\x03\"!\n\nConstraint\x12\x13\n\x0b\x64\x65scription\x18\x01 \x01(\t*3\n\rPrimitiveType\x12\t\n\x05INT64\x10\x00\x12\n\n\x06STRING\x10\x01\x12\x0b\n\x07\x46LOAT64\x10\x02\x62\x06proto3')
+
+_PRIMITIVETYPE = DESCRIPTOR.enum_types_by_name['PrimitiveType']
+PrimitiveType = enum_type_wrapper.EnumTypeWrapper(_PRIMITIVETYPE)
+INT64 = 0
+STRING = 1
+FLOAT64 = 2
+
+
+_SQLWORKERCONFIGURATION = DESCRIPTOR.message_types_by_name['SqlWorkerConfiguration']
+_VALIDATIONCONFIGURATION = DESCRIPTOR.message_types_by_name['ValidationConfiguration']
+_TABLESCHEMA = DESCRIPTOR.message_types_by_name['TableSchema']
+_NAMEDCOLUMN = DESCRIPTOR.message_types_by_name['NamedColumn']
+_COLUMNTYPE = DESCRIPTOR.message_types_by_name['ColumnType']
+_COMPUTATIONCONFIGURATION = DESCRIPTOR.message_types_by_name['ComputationConfiguration']
+_PRIVACYSETTINGS = DESCRIPTOR.message_types_by_name['PrivacySettings']
+_CONSTRAINT = DESCRIPTOR.message_types_by_name['Constraint']
+SqlWorkerConfiguration = _reflection.GeneratedProtocolMessageType('SqlWorkerConfiguration', (_message.Message,), {
+  'DESCRIPTOR' : _SQLWORKERCONFIGURATION,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.SqlWorkerConfiguration)
+  })
+_sym_db.RegisterMessage(SqlWorkerConfiguration)
+
+ValidationConfiguration = _reflection.GeneratedProtocolMessageType('ValidationConfiguration', (_message.Message,), {
+  'DESCRIPTOR' : _VALIDATIONCONFIGURATION,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.ValidationConfiguration)
+  })
+_sym_db.RegisterMessage(ValidationConfiguration)
+
+TableSchema = _reflection.GeneratedProtocolMessageType('TableSchema', (_message.Message,), {
+  'DESCRIPTOR' : _TABLESCHEMA,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.TableSchema)
+  })
+_sym_db.RegisterMessage(TableSchema)
+
+NamedColumn = _reflection.GeneratedProtocolMessageType('NamedColumn', (_message.Message,), {
+  'DESCRIPTOR' : _NAMEDCOLUMN,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.NamedColumn)
+  })
+_sym_db.RegisterMessage(NamedColumn)
+
+ColumnType = _reflection.GeneratedProtocolMessageType('ColumnType', (_message.Message,), {
+  'DESCRIPTOR' : _COLUMNTYPE,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.ColumnType)
+  })
+_sym_db.RegisterMessage(ColumnType)
+
+ComputationConfiguration = _reflection.GeneratedProtocolMessageType('ComputationConfiguration', (_message.Message,), {
+  'DESCRIPTOR' : _COMPUTATIONCONFIGURATION,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.ComputationConfiguration)
+  })
+_sym_db.RegisterMessage(ComputationConfiguration)
+
+PrivacySettings = _reflection.GeneratedProtocolMessageType('PrivacySettings', (_message.Message,), {
+  'DESCRIPTOR' : _PRIVACYSETTINGS,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.PrivacySettings)
+  })
+_sym_db.RegisterMessage(PrivacySettings)
+
+Constraint = _reflection.GeneratedProtocolMessageType('Constraint', (_message.Message,), {
+  'DESCRIPTOR' : _CONSTRAINT,
+  '__module__' : 'compute_sql_pb2'
+  # @@protoc_insertion_point(class_scope:compute_sql.Constraint)
+  })
+_sym_db.RegisterMessage(Constraint)
+
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _PRIMITIVETYPE._serialized_start=747
+  _PRIMITIVETYPE._serialized_end=798
+  _SQLWORKERCONFIGURATION._serialized_start=35
+  _SQLWORKERCONFIGURATION._serialized_end=198
+  _VALIDATIONCONFIGURATION._serialized_start=200
+  _VALIDATIONCONFIGURATION._serialized_end=272
+  _TABLESCHEMA._serialized_start=274
+  _TABLESCHEMA._serialized_end=335
+  _NAMEDCOLUMN._serialized_start=337
+  _NAMEDCOLUMN._serialized_end=423
+  _COLUMNTYPE._serialized_start=425
+  _COLUMNTYPE._serialized_end=506
+  _COMPUTATIONCONFIGURATION._serialized_start=509
+  _COMPUTATIONCONFIGURATION._serialized_end=658
+  _PRIVACYSETTINGS._serialized_start=660
+  _PRIVACYSETTINGS._serialized_end=710
+  _CONSTRAINT._serialized_start=712
+  _CONSTRAINT._serialized_end=745
+# @@protoc_insertion_point(module_scope)

decentriq_util-0.1.0/decentriq_util/sql/proto/compute_sql_pb2.pyi

@@ -0,0 +1,146 @@
+"""
+@generated by mypy-protobuf.  Do not edit manually!
+isort:skip_file
+"""
+import builtins
+import google.protobuf.descriptor
+import google.protobuf.internal.containers
+import google.protobuf.internal.enum_type_wrapper
+import google.protobuf.message
+import typing
+import typing_extensions
+
+DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
+
+class _PrimitiveType:
+    ValueType = typing.NewType('ValueType', builtins.int)
+    V: typing_extensions.TypeAlias = ValueType
+class _PrimitiveTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_PrimitiveType.ValueType], builtins.type):
+    DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
+    INT64: _PrimitiveType.ValueType  # 0
+    STRING: _PrimitiveType.ValueType  # 1
+    FLOAT64: _PrimitiveType.ValueType  # 2
+class PrimitiveType(_PrimitiveType, metaclass=_PrimitiveTypeEnumTypeWrapper):
+    pass
+
+INT64: PrimitiveType.ValueType  # 0
+STRING: PrimitiveType.ValueType  # 1
+FLOAT64: PrimitiveType.ValueType  # 2
+global___PrimitiveType = PrimitiveType
+
+
+class SqlWorkerConfiguration(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    VALIDATION_FIELD_NUMBER: builtins.int
+    COMPUTATION_FIELD_NUMBER: builtins.int
+    @property
+    def validation(self) -> global___ValidationConfiguration: ...
+    @property
+    def computation(self) -> global___ComputationConfiguration: ...
+    def __init__(self,
+        *,
+        validation: typing.Optional[global___ValidationConfiguration] = ...,
+        computation: typing.Optional[global___ComputationConfiguration] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["computation",b"computation","configuration",b"configuration","validation",b"validation"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["computation",b"computation","configuration",b"configuration","validation",b"validation"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["configuration",b"configuration"]) -> typing.Optional[typing_extensions.Literal["validation","computation"]]: ...
+global___SqlWorkerConfiguration = SqlWorkerConfiguration
+
+class ValidationConfiguration(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    TABLESCHEMA_FIELD_NUMBER: builtins.int
+    @property
+    def tableSchema(self) -> global___TableSchema: ...
+    def __init__(self,
+        *,
+        tableSchema: typing.Optional[global___TableSchema] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["tableSchema",b"tableSchema"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["tableSchema",b"tableSchema"]) -> None: ...
+global___ValidationConfiguration = ValidationConfiguration
+
+class TableSchema(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    NAMEDCOLUMNS_FIELD_NUMBER: builtins.int
+    @property
+    def namedColumns(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___NamedColumn]: ...
+    def __init__(self,
+        *,
+        namedColumns: typing.Optional[typing.Iterable[global___NamedColumn]] = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["namedColumns",b"namedColumns"]) -> None: ...
+global___TableSchema = TableSchema
+
+class NamedColumn(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    NAME_FIELD_NUMBER: builtins.int
+    COLUMNTYPE_FIELD_NUMBER: builtins.int
+    name: typing.Text
+    @property
+    def columnType(self) -> global___ColumnType: ...
+    def __init__(self,
+        *,
+        name: typing.Optional[typing.Text] = ...,
+        columnType: typing.Optional[global___ColumnType] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["_name",b"_name","columnType",b"columnType","name",b"name"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["_name",b"_name","columnType",b"columnType","name",b"name"]) -> None: ...
+    def WhichOneof(self, oneof_group: typing_extensions.Literal["_name",b"_name"]) -> typing.Optional[typing_extensions.Literal["name"]]: ...
+global___NamedColumn = NamedColumn
+
+class ColumnType(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    PRIMITIVETYPE_FIELD_NUMBER: builtins.int
+    NULLABLE_FIELD_NUMBER: builtins.int
+    primitiveType: global___PrimitiveType.ValueType
+    nullable: builtins.bool
+    def __init__(self,
+        *,
+        primitiveType: global___PrimitiveType.ValueType = ...,
+        nullable: builtins.bool = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["nullable",b"nullable","primitiveType",b"primitiveType"]) -> None: ...
+global___ColumnType = ColumnType
+
+class ComputationConfiguration(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    SQLSTATEMENT_FIELD_NUMBER: builtins.int
+    PRIVACYSETTINGS_FIELD_NUMBER: builtins.int
+    CONSTRAINTS_FIELD_NUMBER: builtins.int
+    sqlStatement: typing.Text
+    @property
+    def privacySettings(self) -> global___PrivacySettings: ...
+    @property
+    def constraints(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Constraint]: ...
+    def __init__(self,
+        *,
+        sqlStatement: typing.Text = ...,
+        privacySettings: typing.Optional[global___PrivacySettings] = ...,
+        constraints: typing.Optional[typing.Iterable[global___Constraint]] = ...,
+        ) -> None: ...
+    def HasField(self, field_name: typing_extensions.Literal["privacySettings",b"privacySettings"]) -> builtins.bool: ...
+    def ClearField(self, field_name: typing_extensions.Literal["constraints",b"constraints","privacySettings",b"privacySettings","sqlStatement",b"sqlStatement"]) -> None: ...
+global___ComputationConfiguration = ComputationConfiguration
+
+class PrivacySettings(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    MINAGGREGATIONGROUPSIZE_FIELD_NUMBER: builtins.int
+    minAggregationGroupSize: builtins.int
+    def __init__(self,
+        *,
+        minAggregationGroupSize: builtins.int = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["minAggregationGroupSize",b"minAggregationGroupSize"]) -> None: ...
+global___PrivacySettings = PrivacySettings
+
+class Constraint(google.protobuf.message.Message):
+    DESCRIPTOR: google.protobuf.descriptor.Descriptor
+    DESCRIPTION_FIELD_NUMBER: builtins.int
+    description: typing.Text
+    def __init__(self,
+        *,
+        description: typing.Text = ...,
+        ) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["description",b"description"]) -> None: ...
+global___Constraint = Constraint

decentriq_util-0.1.0/pyproject.toml

@@ -0,0 +1,19 @@
+[tool.poetry]
+name = "decentriq_util"
+version = "0.1.0"
+description = "Utility functions for Python code run in Decentriq workers"
+authors = ["hafenr <robin.hafen@gmail.com>"]
+
+[tool.poetry.dependencies]
+python = "^3.9"
+protobuf = "^4.21"
+pandas = ">=1.4.4"
+pyspark = "^3.4.0"
+psutil = "^5.9.6"
+
+[tool.poetry.dev-dependencies]
+pytest = "^7.2.0"
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"