decentriq_util 0.1.0.dev0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,29 @@
1
+ """
2
+ The `decentriq_util` library provides helpers for common operations,
3
+ not only when automating your workflows with the `decentriq_platform` package,
4
+ but also when writing your own Python computations inside Data Clean Rooms.
5
+
6
+ It is available inside the computation container for Python scripts
7
+ and can be imported just like any other library. For example:
8
+ ```python
9
+ import decentriq_util
10
+ data = decentriq_util.sql.read_sql_data_from_dir("/path/to/computation")
11
+ ```
12
+ """
13
+ from . import sql
14
+ from . import proto
15
+ from . import python
16
+ from . import error
17
+ from . import spark
18
+
19
+ from .sql import (
20
+ read_tabular_data,
21
+ write_tabular_data
22
+ )
23
+
24
+ __pdoc__ = {
25
+ "sql": True,
26
+ "proto": False,
27
+ "python": False,
28
+ "error": False,
29
+ }
@@ -0,0 +1,131 @@
1
+ from typing import Callable, List
2
+
3
+ import traceback
4
+ import sys
5
+ import os
6
+
7
+
8
+ def _strip_away_nix_store(path: str) -> str:
9
+ """
10
+ Replace '/nix/store/hash-derivation-name' with '/system/' for better
11
+ readability of the stack trace.
12
+ """
13
+ if path.startswith("/nix/store"):
14
+ return "/system/" + "/".join(path.split("/")[4:])
15
+ else:
16
+ return path
17
+
18
+
19
+ class SafeError(Exception):
20
+ """
21
+ Base class for errors that can be turned into safe error messages
22
+ that can be displayed to the user without risking the leakage of
23
+ sensitive information.
24
+ """
25
+ def __init__(self, *args, **kwargs):
26
+ super().__init__(*args, **kwargs)
27
+
28
+ def safe_str(self) -> str:
29
+ """
30
+ Return a safe string representation of this exception
31
+ that does not include any sensitive data.
32
+
33
+ This method must be implemented explicitly.
34
+ """
35
+ raise Exception("Not implemented")
36
+
37
+
38
+ class CatchSafeErrorContext:
39
+ """
40
+ Context manager to handle error messages and writing them to a pre-defined
41
+ location on the file system.
42
+ The constructed stack track is optionally cleaned from any sensitive
43
+ information that might be contained in the error message.
44
+ """
45
+ def __init__(
46
+ self,
47
+ should_clean_traceback: bool,
48
+ output_file: str,
49
+ should_exit: bool,
50
+ transform_library_path: List[Callable[[str], str]],
51
+ ):
52
+ self.should_clean_traceback = should_clean_traceback
53
+ self.output_file = output_file
54
+ self.should_exit = should_exit
55
+ self.transform_library_path = transform_library_path + [_strip_away_nix_store]
56
+
57
+
58
+ def __enter__(self):
59
+ return self
60
+
61
+
62
+ def __exit__(self, exception_type, exception, tb_):
63
+ if exception:
64
+ if exception_type == SystemExit and exception.code == 0:
65
+ # Don't produce a traceback when sys.exit(0) is used.
66
+ return
67
+
68
+ should_capture_local_variables = not self.should_clean_traceback
69
+ error_class_module = exception.__class__.__module__
70
+ ignored_modules = set(["builtins", "__main__"])
71
+ if error_class_module and error_class_module not in ignored_modules:
72
+ error_qualified_class_name =\
73
+ error_class_module + "." + exception.__class__.__name__
74
+
75
+ else:
76
+ error_qualified_class_name = exception.__class__.__name__
77
+
78
+ # Skip the first frame as it is the one in the wrapper script
79
+ tbe = traceback.TracebackException.from_exception(
80
+ exception, capture_locals=should_capture_local_variables
81
+ )
82
+ lines = ["Traceback (most recent call last):"]
83
+ for frame in tbe.stack[1:]:
84
+ filename = frame.filename
85
+ # Apply all configured transformer functions to the file paths
86
+ # to make them more readable.
87
+ for transform_fn in self.transform_library_path:
88
+ filename = transform_fn(filename)
89
+ lines.append(
90
+ f" File \"{filename}\","
91
+ f" line {frame.lineno}, in {frame.name}\n {frame.line}"
92
+ )
93
+
94
+ if self.should_clean_traceback:
95
+ # Even if we are supposed to provide a safe error, if it's one of our
96
+ # own errors, we can still construct an (informative) safe error
97
+ # message.
98
+ if isinstance(exception, SafeError):
99
+ stringified_error =\
100
+ error_qualified_class_name + ": " + exception.safe_str()
101
+ else:
102
+ stringified_error = error_qualified_class_name +\
103
+ ": (error message removed due to data privacy)"
104
+ else:
105
+ stringified_error = error_qualified_class_name + ": " + str(exception)
106
+
107
+ lines.append(stringified_error)
108
+ safe_traceback = "\n".join(lines)
109
+
110
+ with open(self.output_file, "w") as f:
111
+ f.write(safe_traceback)
112
+
113
+ if self.should_exit:
114
+ sys.exit(1)
115
+ else:
116
+ # Do not propagate error
117
+ return True
118
+
119
+
120
+ def catch_safe_error(
121
+ transform_library_path: List[Callable[[str], str]] = []
122
+ ) -> CatchSafeErrorContext:
123
+ """Construct a CatchSafeErrorContext with default values"""
124
+ should_clean_traceback = \
125
+ "INCLUDE_CONTAINER_LOGS_ON_ERROR" not in os.environ or os.environ["INCLUDE_CONTAINER_LOGS_ON_ERROR"] != "true"
126
+ return CatchSafeErrorContext(
127
+ should_clean_traceback=should_clean_traceback,
128
+ should_exit=True,
129
+ output_file="/internal/safe_error",
130
+ transform_library_path=transform_library_path
131
+ )
@@ -0,0 +1,24 @@
1
+ from typing import Tuple
2
+ from google.protobuf.message import Message
3
+ from google.protobuf.internal.encoder import _VarintBytes # type: ignore
4
+ from google.protobuf.internal.decoder import _DecodeVarint32 # type: ignore
5
+
6
+
7
+ def parse_length_delimited(serialized_bytes: bytes, deserialized_object: Message) -> int:
8
+ """
9
+ Decodes a binary content and returns a parsed object of a given type
10
+ """
11
+ res: Tuple[int, int] = _DecodeVarint32(serialized_bytes, 0)
12
+ message_length, offset = res
13
+ end_offset = offset + message_length
14
+ deserialized_object.ParseFromString(bytes(serialized_bytes[offset:end_offset]))
15
+ return end_offset
16
+
17
+
18
+ def serialize_length_delimited(message_object: Message) -> bytes:
19
+ """
20
+ Encodes an object of a given type to binary
21
+ """
22
+ serialized: bytes = _VarintBytes(message_object.ByteSize())
23
+ serialized += message_object.SerializeToString()
24
+ return serialized
@@ -0,0 +1,21 @@
1
+ import traceback
2
+ import datetime
3
+
4
+ def log_entry(record: str):
5
+ """
6
+ Logs an entry in the user log file to be displayed at the end of the execution.
7
+ """
8
+ current_time = datetime.datetime.now(datetime.timezone.utc)
9
+ with open("/internal/logs", "a+") as f:
10
+ f.write(f"LOG {current_time}: {record}\n")
11
+
12
+ def log_stack_trace(e: Exception):
13
+ """
14
+ Allows to add a striped stack strace for an Exception e in the user log file.
15
+ The striped stack trace will not contain any environment variable or secrets
16
+ in it.
17
+ """
18
+ tb = traceback.TracebackException.from_exception(e, capture_locals=False)
19
+ safe_traceback = "".join([s for s in tb.format()][:-1])
20
+ error_qualified_class_name = e.__class__.__module__ + "." + e.__class__.__name__
21
+ log_entry(f"Encountered error: {error_qualified_class_name}, Traceback:\n{safe_traceback}")
@@ -0,0 +1,215 @@
1
+ import tempfile
2
+ from typing import Optional, Tuple, Any
3
+ from contextlib import contextmanager
4
+ import multiprocessing
5
+ import math
6
+ import os
7
+ import shutil
8
+ import psutil
9
+
10
+ from pyspark.sql import SparkSession
11
+
12
+
13
+ def _get_cgroup_memory() -> int:
14
+ max_memory_str = None
15
+ if os.path.isfile("/sys/fs/cgroup/memory.max"):
16
+ with open("/sys/fs/cgroup/memory.max", "r") as f:
17
+ max_memory_str = f.read().strip()
18
+ elif os.path.isfile("/sys/fs/cgroup/memory/memory.limit_in_bytes"):
19
+ with open("/sys/fs/cgroup/memory/memory.limit_in_bytes", "r") as f:
20
+ max_memory_str = f.read().strip()
21
+
22
+ if max_memory_str == "max":
23
+ # Fallback to available virtual memory size
24
+ max_memory = psutil.virtual_memory().available
25
+ else:
26
+ max_memory = int(max_memory_str)
27
+
28
+ if max_memory is not None:
29
+ return max_memory
30
+ else:
31
+ print("Unable to determine available memory from cgroup, assuming 4G")
32
+ return 4 * 1024 * 1024
33
+
34
+
35
+ def _get_num_partitions(
36
+ paths: list[str],
37
+ available_memory: int,
38
+ cpu_count: int,
39
+ ) -> int:
40
+ """Determine the number of partitions s.t. each partition is as big as possible
41
+ while still making sure all cores are occupied."""
42
+ def get_file_or_dir_size(path):
43
+ if os.path.isfile(path):
44
+ size = os.path.getsize(path)
45
+ elif os.path.isdir(path):
46
+ total_size = 0
47
+ for dirpath, dirnames, filenames in os.walk(path):
48
+ for f in filenames:
49
+ fp = os.path.join(dirpath, f)
50
+ total_size += os.path.getsize(fp)
51
+ size = total_size
52
+ else:
53
+ raise ValueError(f"The path {path} is neither a file nor a directory.")
54
+ return size
55
+
56
+ file_size = max([get_file_or_dir_size(path) for path in paths])
57
+
58
+ # Don't make the partitions smaller than 36MiB
59
+ megabyte_in_bytes = 1024 * 1024
60
+ min_partition_size = 32 * megabyte_in_bytes
61
+
62
+ # Try to make partitions as big as possible s.t. each CPU is saturated,
63
+ # while accounting for 20% overhead in spark.
64
+ target_partition_size = max(
65
+ int(0.8 * available_memory / cpu_count) // megabyte_in_bytes * megabyte_in_bytes,
66
+ min_partition_size
67
+ )
68
+
69
+ while (
70
+ # If the file is small enough that it cannot be split across all cores,
71
+ # divide the partition size further.
72
+ int(file_size / target_partition_size) < cpu_count
73
+ ):
74
+ # Don't go smaller than the minimum parition size
75
+ if target_partition_size / 2 < min_partition_size:
76
+ break
77
+ target_partition_size = target_partition_size / 2
78
+
79
+ print(f"File size: {file_size}")
80
+ print(f"Target partition size: {target_partition_size}")
81
+
82
+ num_partitions = int(math.ceil(file_size / target_partition_size))
83
+
84
+ return num_partitions
85
+
86
+
87
+ def _determine_optimal_spark_settings(
88
+ files: list[str],
89
+ heap_size_extra_room: Optional[int] = 512 * 1024 * 1024
90
+ ) -> list[Tuple[str, str]]:
91
+ def get_spark_memory() -> int:
92
+ """Determine the amount of memory available to spark"""
93
+ cgroup_memory = _get_cgroup_memory()
94
+ spark_memory = cgroup_memory - (heap_size_extra_room or 512 * 1024 * 1024)
95
+ return spark_memory
96
+
97
+ cpu_count = multiprocessing.cpu_count()
98
+ spark_memory = get_spark_memory()
99
+ num_partitions = _get_num_partitions(files, spark_memory, cpu_count)
100
+
101
+ print(f"CPU count: {cpu_count}")
102
+ print(f"Spark memory: {spark_memory}")
103
+ print(f"Num partitions: {num_partitions}")
104
+
105
+ settings = [
106
+ ("spark.sql.shuffle.partitions", str(num_partitions)),
107
+ ("spark.default.parallelism", str(num_partitions)),
108
+ ("spark.driver.cores", str(cpu_count)),
109
+ ]
110
+ print("Spark settings:\n" + f"\n".join([str(x) for x in settings]))
111
+
112
+ return settings
113
+
114
+
115
+ def _create_spark_session(
116
+ name: str = "local_spark_session",
117
+ parallelism: int = 8,
118
+ heap_size: Optional[int] = None,
119
+ heap_size_extra_room: int = 512 * 1024 * 1024,
120
+ config: list[Tuple[str, Any]] = [],
121
+ java_tmp_dir = "/output",
122
+ ) -> SparkSession:
123
+ """
124
+ :param name: The name of the spark session.
125
+ :param parallelism: The size of the executor pool computing internal spark tasks.
126
+ :param heap_size: If set then this determines the JVM heap size. Note that not all the heap size is used by spark.
127
+ :param heap_size_extra_room: If heap_size is unset then it is determined by the available memory in the current
128
+ cgroup. Then heap_size_extra_room is subtracted to arrive at the final heap size. This number should account for
129
+ the python runtime itself as well as other structures that may be needed during spark compute.
130
+ :param config: Additional config settings to set on the spark config.
131
+ :param java_tmp_dir: Location for the JVM to store temp files.
132
+ :return: The spark session.
133
+ """
134
+ os.environ["SPARK_EXECUTOR_POOL_SIZE"] = str(parallelism)
135
+ os.environ["JDK_JAVA_OPTIONS"] = f'-Duser.home=/output -Djava.io.tmpdir="{java_tmp_dir}"'
136
+ spark_memory = None
137
+ if heap_size is None:
138
+ cgroup_memory = _get_cgroup_memory()
139
+ spark_memory = cgroup_memory - heap_size_extra_room
140
+ else:
141
+ spark_memory = heap_size
142
+
143
+ spark_memory_4096 = (spark_memory // 4096) * 4096
144
+ if "spark.driver.memory" not in dict(config):
145
+ config.append(("spark.driver.memory", str(spark_memory_4096)))
146
+
147
+ ss = (
148
+ SparkSession.builder
149
+ .appName(name)
150
+ .master("local[*]")
151
+ )
152
+ for (key, value) in config:
153
+ ss = ss.config(key, value)
154
+ return ss.getOrCreate()
155
+
156
+
157
+ @contextmanager
158
+ def spark_session(temp_dir: str = "/scratch", input_files: list[str] = None, **kwargs):
159
+ """
160
+ Create a spark session and configure it according to the enclave environment.
161
+
162
+ **Parameters**:
163
+ - `temp_dir`: Where to store temporary data such as persisted data frames
164
+ or shuffle data.
165
+ - `input_files`: A list of input files on the basis of which the partition
166
+ size is determined.
167
+
168
+ **Example**:
169
+
170
+ ```python
171
+ import decentriq_util as dq
172
+
173
+ # Path to a potentially very large file
174
+ input_csv_path = "/input/my_file.csv"
175
+
176
+ # Automatically create and configure a spark session and
177
+ # make sure it's being stopped at the end.
178
+ with dq.spark.spark_session(input_files=[input_csv_path]) as ss:
179
+ # Read from a CSV file
180
+ df = ss.read.csv(input_csv_path, header=False).cache()
181
+
182
+ # Perform any pyspark transformations
183
+ print(f"Original number of rows: {df.count()}")
184
+ result_df = df.limit(100)
185
+
186
+ # Write the result to an output file
187
+ result_df.write.parquet("/output/my_file.parquet")
188
+ ```
189
+ """
190
+ with tempfile.TemporaryDirectory(dir=temp_dir, prefix="java-") as java_tmp:
191
+ with tempfile.TemporaryDirectory(dir=temp_dir, prefix="spark-") as spark_tmp:
192
+ config = kwargs.get("config", [])
193
+ config_dict = dict(config)
194
+ if input_files is not None and input_files:
195
+ optimal_settings = _determine_optimal_spark_settings(
196
+ input_files, heap_size_extra_room=kwargs.get("heap_size_extra_room")
197
+ )
198
+ for key, value in optimal_settings:
199
+ if key not in config_dict:
200
+ config.append((key, value))
201
+ if "spark.local.dir" not in config_dict:
202
+ config.append(
203
+ ("spark.local.dir", spark_tmp)
204
+ )
205
+ kwargs["config"] = config
206
+ if "java_tmp_dir" not in kwargs:
207
+ kwargs["java_tmp_dir"] = java_tmp
208
+ ss = _create_spark_session(**kwargs)
209
+ try:
210
+ yield ss
211
+ finally:
212
+ try:
213
+ ss.stop()
214
+ except:
215
+ pass
@@ -0,0 +1,266 @@
1
+ import typing
2
+ from typing import Tuple, List
3
+ import os
4
+ import numpy
5
+ import pandas as pd
6
+
7
+ from .proto.compute_sql_pb2 import (
8
+ TableSchema,
9
+ PrimitiveType,
10
+ NamedColumn,
11
+ ColumnType
12
+ )
13
+ from ..proto import (
14
+ parse_length_delimited,
15
+ serialize_length_delimited
16
+ )
17
+
18
+ __pdoc__ = {
19
+ "proto": False,
20
+ }
21
+
22
+
23
+ datatype_mapping = {
24
+ PrimitiveType.INT64: pd.Int64Dtype(),
25
+ PrimitiveType.STRING: str,
26
+ PrimitiveType.FLOAT64: numpy.float64,
27
+ }
28
+
29
+
30
+ datatype_mapping_reverse = {
31
+ "Int64": PrimitiveType.INT64,
32
+ "int64": PrimitiveType.INT64,
33
+ "object": PrimitiveType.STRING,
34
+ "float64": PrimitiveType.FLOAT64
35
+ }
36
+
37
+
38
+ DATASET_FILE_NAME = "dataset.csv"
39
+ SCHEMA_FILE_NAME = "types"
40
+
41
+
42
+ def _get_numpy_dtypes_for_table_schema_object(
43
+ schema,
44
+ use_numpy_int_type: bool = False
45
+ ) -> List[Tuple[str, numpy.dtype]]:
46
+ ix = 1
47
+ columns = []
48
+ for col in schema.namedColumns:
49
+ # This matches R convention, unnamed columns get a position
50
+ name = col.name if col.name else f"V{ix}"
51
+ sql_type = col.columnType.primitiveType
52
+ if sql_type == PrimitiveType.INT64 and use_numpy_int_type:
53
+ dtype = numpy.int64
54
+ else:
55
+ dtype = datatype_mapping[sql_type]
56
+ columns.append((name, dtype))
57
+ ix += 1
58
+
59
+ return columns
60
+
61
+
62
+ def _get_numpy_dtypes_for_table_schema(
63
+ schema_path: str,
64
+ use_numpy_int_type: bool = False
65
+ ) -> List[Tuple[str, numpy.dtype]]:
66
+ with open(schema_path, "rb") as f:
67
+ schema = TableSchema()
68
+ parse_length_delimited(f.read(), schema)
69
+
70
+ return _get_numpy_dtypes_for_table_schema_object(schema, use_numpy_int_type)
71
+
72
+
73
+ def _get_table_schema_for_data_frame(df: pd.DataFrame) -> TableSchema:
74
+ dtypes_as_string = [str(dtype) for dtype in df.dtypes]
75
+ proto_columns = []
76
+ for name, dtype in zip(list(df.columns), dtypes_as_string):
77
+ primitive_type = datatype_mapping_reverse[dtype]
78
+ column = NamedColumn(
79
+ name=name,
80
+ columnType=ColumnType(
81
+ primitiveType=primitive_type,
82
+ nullable=True
83
+ )
84
+ )
85
+ proto_columns.append(column)
86
+ return TableSchema(
87
+ namedColumns=proto_columns
88
+ )
89
+
90
+
91
+ def _assert_files_do_exist(dataset_file: str, schema_file: str) -> None:
92
+ if not os.path.exists(dataset_file):
93
+ raise Exception(
94
+ f"Cannot read data from {dataset_file}, file containing data"
95
+ f" {dataset_file} does not exist!"
96
+ )
97
+ if not os.path.exists(schema_file):
98
+ raise Exception(
99
+ f"Cannot read data from {dataset_file}, file encoding schema"
100
+ f" {schema_file} does not exist!"
101
+ )
102
+
103
+
104
+ def _assert_files_dont_exist(dataset_file: str, schema_file: str):
105
+ if os.path.exists(dataset_file):
106
+ raise Exception(
107
+ "Cannot write data to {path}, file containing data"
108
+ f" {dataset_file} already exists!"
109
+ )
110
+ if os.path.exists(schema_file):
111
+ raise Exception(
112
+ "Cannot write data to {path}, file encoding schema"
113
+ f" {schema_file} already exists!"
114
+ )
115
+
116
+
117
+ def read_tabular_data(path: str) -> pd.DataFrame:
118
+ """
119
+ Read data from a tabular input node or as written by an SQL computation.
120
+
121
+ The input file must not contain a header row!
122
+
123
+ Empty values for columns of type TEXT are replaced with empty strings.
124
+
125
+ Empty values in columns of type INT64 are read as the "pandas.NA" value.
126
+ For more information, please refer to the
127
+ [pandas documentation](https://pandas.pydata.org/pandas-docs/dev/user_guide/integer_na.html).
128
+ """
129
+ return read_sql_data_from_dir(path)
130
+
131
+
132
+ def read_sql_data_from_dir(path: str) -> pd.DataFrame:
133
+ """
134
+ Note: this function is deprecated, use the function `read_tabular_data`.
135
+
136
+ This function uses the previous `read_sql_data` function to convert
137
+ the standard format of validated tabular input datasets and the output
138
+ of SQL computations into a pandas dataframe, automatically importing the
139
+ headers the are contained in the `types` file as well. One example below:
140
+ ```python
141
+ import decentriq_util.sql as dqu
142
+ import pandas as pd
143
+
144
+ dataframe_1 = dqu.read_sql_data_from_dir("/input/mySQLComputation")
145
+ dataframe_2 = dqu.read_sql_data_from_dir("/input/myTabularDataset")
146
+ ```
147
+ """
148
+ dataset_file = os.path.join(path, DATASET_FILE_NAME)
149
+ schema_file = os.path.join(path, SCHEMA_FILE_NAME)
150
+ _assert_files_do_exist(dataset_file, schema_file)
151
+ return read_sql_data(dataset_file, schema_file)
152
+
153
+
154
+ def write_tabular_data(df: pd.DataFrame, path: str = "/output"):
155
+ """
156
+ Write a Pandas dataframe into the given directory in a format compatible
157
+ with downstream SQL computation. This function allows you to use the output of
158
+ a python computation as input to a SQL computation.
159
+ For example:
160
+
161
+ ```python
162
+ import decentriq_util.sql as dqu
163
+ import pandas as pd
164
+
165
+ my_dataframe = pd.DataFrame()
166
+ ...
167
+ dqu.write_tabular_data(my_dataframe)
168
+ ```
169
+ """
170
+ write_sql_data_to_dir(df, path)
171
+
172
+
173
+ def write_sql_data_to_dir(df: pd.DataFrame, path: str):
174
+ """
175
+ Note: this function is deprecated, use the function `write_tabular_data`.
176
+
177
+ Write a Pandas dataframe into the given directory in a format compatible
178
+ with downstream SQL computations. This function allows you to use the output of
179
+ a python computation as input to a SQL computation.
180
+ To do so, you will need to use `/output` as path.
181
+
182
+ For example:
183
+ ```python
184
+ import decentriq_util.sql as dqu
185
+ import pandas as pd
186
+
187
+ my_dataframe = pd.DataFrame()
188
+ ...
189
+ dqu.write_sql_data_to_dir(my_dataframe, "/output")
190
+ ```
191
+ """
192
+ dataset_file = os.path.join(path, DATASET_FILE_NAME)
193
+ schema_file = os.path.join(path, SCHEMA_FILE_NAME)
194
+ _assert_files_dont_exist(dataset_file, schema_file)
195
+ return write_sql_data(df, dataset_file, schema_file)
196
+
197
+
198
+ def read_sql_data(data_path: str, schema_path: str) -> pd.DataFrame:
199
+ """
200
+ Read data as output by the SQL validation node.
201
+
202
+ The input file must not contain a header row!
203
+
204
+ Empty values for columns of type TEXT are replaced with empty strings.
205
+
206
+ Empty values in columns of type INT64 are read as the "pandas.NA" value.
207
+ For more information, please refer to the
208
+ [pandas documentation](https://pandas.pydata.org/pandas-docs/dev/user_guide/integer_na.html).
209
+ """
210
+ _assert_files_do_exist(data_path, schema_path)
211
+ columns = _get_numpy_dtypes_for_table_schema(schema_path)
212
+ dtype_per_column = {
213
+ ix: column_type
214
+ for ix, (_, column_type) in enumerate(columns)
215
+ }
216
+ try:
217
+ # First read just the data.
218
+ # Tell the type linter that the value is in fact a DataFrame
219
+ # as it could be a TextParser object depending on whether either
220
+ # `chunksize` or `iterator` arguments are specified.
221
+ original_data = typing.cast(
222
+ pd.DataFrame,
223
+ pd.read_csv(
224
+ data_path,
225
+ index_col=None,
226
+ header=None, # sql verifier doesn't write headers
227
+ dtype=str # treat everything as string at first
228
+ )
229
+ )
230
+ # Read the schema file that tells us about the data types.
231
+ # Replace NaN values with empty strings for string-type columns.
232
+ string_column_replacements =\
233
+ {ix: "" for ix, dtype in dtype_per_column.items() if dtype == str}
234
+ original_data.fillna(string_column_replacements, inplace=True)
235
+ # Check whether the DF is empty now (even if it wasn't in the beginning)
236
+ # If yes, return an empty DF.
237
+ if original_data.shape[0] == 0:
238
+ return pd.DataFrame(
239
+ { name: pd.Series(dtype=dtype) for name, dtype in columns }
240
+ )
241
+ else:
242
+ # Convert columns to the correct dtype
243
+ for ix, dtype in dtype_per_column.items():
244
+ if original_data[ix].dtype != dtype:
245
+ original_data[ix] = original_data[ix].astype(dtype)
246
+ except pd.errors.EmptyDataError:
247
+ return pd.DataFrame(
248
+ { name: pd.Series(dtype=dtype) for name, dtype in columns }
249
+ )
250
+ header = [column_name for column_name, _ in columns]
251
+ original_data.columns = header
252
+
253
+ return original_data
254
+
255
+
256
+ def write_sql_data(df: pd.DataFrame, data_path: str, schema_path: str):
257
+ """
258
+ Write a Pandas dataframe to the given path, and store its schema including column names at
259
+ `schema_path` with the protobuf standard used in the cross-workers interactions.
260
+ """
261
+ _assert_files_dont_exist(data_path, schema_path)
262
+ schema = _get_table_schema_for_data_frame(df)
263
+ schema_serialized = serialize_length_delimited(schema)
264
+ df.to_csv(data_path, index=False, header=False, sep=",")
265
+ with open(schema_path, "wb") as f:
266
+ f.write(schema_serialized)
@@ -0,0 +1,11 @@
1
+ from .compute_sql_pb2 import (
2
+ TableSchema,
3
+ NamedColumn,
4
+ ColumnType,
5
+ PrimitiveType,
6
+ SqlWorkerConfiguration,
7
+ ValidationConfiguration,
8
+ ComputationConfiguration,
9
+ PrivacySettings
10
+ )
11
+
@@ -0,0 +1,112 @@
1
+ # -*- coding: utf-8 -*-
2
+ # Generated by the protocol buffer compiler. DO NOT EDIT!
3
+ # source: compute_sql.proto
4
+ """Generated protocol buffer code."""
5
+ from google.protobuf.internal import enum_type_wrapper
6
+ from google.protobuf import descriptor as _descriptor
7
+ from google.protobuf import descriptor_pool as _descriptor_pool
8
+ from google.protobuf import message as _message
9
+ from google.protobuf import reflection as _reflection
10
+ from google.protobuf import symbol_database as _symbol_database
11
+ # @@protoc_insertion_point(imports)
12
+
13
+ _sym_db = _symbol_database.Default()
14
+
15
+
16
+
17
+
18
+ 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')
19
+
20
+ _PRIMITIVETYPE = DESCRIPTOR.enum_types_by_name['PrimitiveType']
21
+ PrimitiveType = enum_type_wrapper.EnumTypeWrapper(_PRIMITIVETYPE)
22
+ INT64 = 0
23
+ STRING = 1
24
+ FLOAT64 = 2
25
+
26
+
27
+ _SQLWORKERCONFIGURATION = DESCRIPTOR.message_types_by_name['SqlWorkerConfiguration']
28
+ _VALIDATIONCONFIGURATION = DESCRIPTOR.message_types_by_name['ValidationConfiguration']
29
+ _TABLESCHEMA = DESCRIPTOR.message_types_by_name['TableSchema']
30
+ _NAMEDCOLUMN = DESCRIPTOR.message_types_by_name['NamedColumn']
31
+ _COLUMNTYPE = DESCRIPTOR.message_types_by_name['ColumnType']
32
+ _COMPUTATIONCONFIGURATION = DESCRIPTOR.message_types_by_name['ComputationConfiguration']
33
+ _PRIVACYSETTINGS = DESCRIPTOR.message_types_by_name['PrivacySettings']
34
+ _CONSTRAINT = DESCRIPTOR.message_types_by_name['Constraint']
35
+ SqlWorkerConfiguration = _reflection.GeneratedProtocolMessageType('SqlWorkerConfiguration', (_message.Message,), {
36
+ 'DESCRIPTOR' : _SQLWORKERCONFIGURATION,
37
+ '__module__' : 'compute_sql_pb2'
38
+ # @@protoc_insertion_point(class_scope:compute_sql.SqlWorkerConfiguration)
39
+ })
40
+ _sym_db.RegisterMessage(SqlWorkerConfiguration)
41
+
42
+ ValidationConfiguration = _reflection.GeneratedProtocolMessageType('ValidationConfiguration', (_message.Message,), {
43
+ 'DESCRIPTOR' : _VALIDATIONCONFIGURATION,
44
+ '__module__' : 'compute_sql_pb2'
45
+ # @@protoc_insertion_point(class_scope:compute_sql.ValidationConfiguration)
46
+ })
47
+ _sym_db.RegisterMessage(ValidationConfiguration)
48
+
49
+ TableSchema = _reflection.GeneratedProtocolMessageType('TableSchema', (_message.Message,), {
50
+ 'DESCRIPTOR' : _TABLESCHEMA,
51
+ '__module__' : 'compute_sql_pb2'
52
+ # @@protoc_insertion_point(class_scope:compute_sql.TableSchema)
53
+ })
54
+ _sym_db.RegisterMessage(TableSchema)
55
+
56
+ NamedColumn = _reflection.GeneratedProtocolMessageType('NamedColumn', (_message.Message,), {
57
+ 'DESCRIPTOR' : _NAMEDCOLUMN,
58
+ '__module__' : 'compute_sql_pb2'
59
+ # @@protoc_insertion_point(class_scope:compute_sql.NamedColumn)
60
+ })
61
+ _sym_db.RegisterMessage(NamedColumn)
62
+
63
+ ColumnType = _reflection.GeneratedProtocolMessageType('ColumnType', (_message.Message,), {
64
+ 'DESCRIPTOR' : _COLUMNTYPE,
65
+ '__module__' : 'compute_sql_pb2'
66
+ # @@protoc_insertion_point(class_scope:compute_sql.ColumnType)
67
+ })
68
+ _sym_db.RegisterMessage(ColumnType)
69
+
70
+ ComputationConfiguration = _reflection.GeneratedProtocolMessageType('ComputationConfiguration', (_message.Message,), {
71
+ 'DESCRIPTOR' : _COMPUTATIONCONFIGURATION,
72
+ '__module__' : 'compute_sql_pb2'
73
+ # @@protoc_insertion_point(class_scope:compute_sql.ComputationConfiguration)
74
+ })
75
+ _sym_db.RegisterMessage(ComputationConfiguration)
76
+
77
+ PrivacySettings = _reflection.GeneratedProtocolMessageType('PrivacySettings', (_message.Message,), {
78
+ 'DESCRIPTOR' : _PRIVACYSETTINGS,
79
+ '__module__' : 'compute_sql_pb2'
80
+ # @@protoc_insertion_point(class_scope:compute_sql.PrivacySettings)
81
+ })
82
+ _sym_db.RegisterMessage(PrivacySettings)
83
+
84
+ Constraint = _reflection.GeneratedProtocolMessageType('Constraint', (_message.Message,), {
85
+ 'DESCRIPTOR' : _CONSTRAINT,
86
+ '__module__' : 'compute_sql_pb2'
87
+ # @@protoc_insertion_point(class_scope:compute_sql.Constraint)
88
+ })
89
+ _sym_db.RegisterMessage(Constraint)
90
+
91
+ if _descriptor._USE_C_DESCRIPTORS == False:
92
+
93
+ DESCRIPTOR._options = None
94
+ _PRIMITIVETYPE._serialized_start=747
95
+ _PRIMITIVETYPE._serialized_end=798
96
+ _SQLWORKERCONFIGURATION._serialized_start=35
97
+ _SQLWORKERCONFIGURATION._serialized_end=198
98
+ _VALIDATIONCONFIGURATION._serialized_start=200
99
+ _VALIDATIONCONFIGURATION._serialized_end=272
100
+ _TABLESCHEMA._serialized_start=274
101
+ _TABLESCHEMA._serialized_end=335
102
+ _NAMEDCOLUMN._serialized_start=337
103
+ _NAMEDCOLUMN._serialized_end=423
104
+ _COLUMNTYPE._serialized_start=425
105
+ _COLUMNTYPE._serialized_end=506
106
+ _COMPUTATIONCONFIGURATION._serialized_start=509
107
+ _COMPUTATIONCONFIGURATION._serialized_end=658
108
+ _PRIVACYSETTINGS._serialized_start=660
109
+ _PRIVACYSETTINGS._serialized_end=710
110
+ _CONSTRAINT._serialized_start=712
111
+ _CONSTRAINT._serialized_end=745
112
+ # @@protoc_insertion_point(module_scope)
@@ -0,0 +1,146 @@
1
+ """
2
+ @generated by mypy-protobuf. Do not edit manually!
3
+ isort:skip_file
4
+ """
5
+ import builtins
6
+ import google.protobuf.descriptor
7
+ import google.protobuf.internal.containers
8
+ import google.protobuf.internal.enum_type_wrapper
9
+ import google.protobuf.message
10
+ import typing
11
+ import typing_extensions
12
+
13
+ DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
14
+
15
+ class _PrimitiveType:
16
+ ValueType = typing.NewType('ValueType', builtins.int)
17
+ V: typing_extensions.TypeAlias = ValueType
18
+ class _PrimitiveTypeEnumTypeWrapper(google.protobuf.internal.enum_type_wrapper._EnumTypeWrapper[_PrimitiveType.ValueType], builtins.type):
19
+ DESCRIPTOR: google.protobuf.descriptor.EnumDescriptor
20
+ INT64: _PrimitiveType.ValueType # 0
21
+ STRING: _PrimitiveType.ValueType # 1
22
+ FLOAT64: _PrimitiveType.ValueType # 2
23
+ class PrimitiveType(_PrimitiveType, metaclass=_PrimitiveTypeEnumTypeWrapper):
24
+ pass
25
+
26
+ INT64: PrimitiveType.ValueType # 0
27
+ STRING: PrimitiveType.ValueType # 1
28
+ FLOAT64: PrimitiveType.ValueType # 2
29
+ global___PrimitiveType = PrimitiveType
30
+
31
+
32
+ class SqlWorkerConfiguration(google.protobuf.message.Message):
33
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
34
+ VALIDATION_FIELD_NUMBER: builtins.int
35
+ COMPUTATION_FIELD_NUMBER: builtins.int
36
+ @property
37
+ def validation(self) -> global___ValidationConfiguration: ...
38
+ @property
39
+ def computation(self) -> global___ComputationConfiguration: ...
40
+ def __init__(self,
41
+ *,
42
+ validation: typing.Optional[global___ValidationConfiguration] = ...,
43
+ computation: typing.Optional[global___ComputationConfiguration] = ...,
44
+ ) -> None: ...
45
+ def HasField(self, field_name: typing_extensions.Literal["computation",b"computation","configuration",b"configuration","validation",b"validation"]) -> builtins.bool: ...
46
+ def ClearField(self, field_name: typing_extensions.Literal["computation",b"computation","configuration",b"configuration","validation",b"validation"]) -> None: ...
47
+ def WhichOneof(self, oneof_group: typing_extensions.Literal["configuration",b"configuration"]) -> typing.Optional[typing_extensions.Literal["validation","computation"]]: ...
48
+ global___SqlWorkerConfiguration = SqlWorkerConfiguration
49
+
50
+ class ValidationConfiguration(google.protobuf.message.Message):
51
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
52
+ TABLESCHEMA_FIELD_NUMBER: builtins.int
53
+ @property
54
+ def tableSchema(self) -> global___TableSchema: ...
55
+ def __init__(self,
56
+ *,
57
+ tableSchema: typing.Optional[global___TableSchema] = ...,
58
+ ) -> None: ...
59
+ def HasField(self, field_name: typing_extensions.Literal["tableSchema",b"tableSchema"]) -> builtins.bool: ...
60
+ def ClearField(self, field_name: typing_extensions.Literal["tableSchema",b"tableSchema"]) -> None: ...
61
+ global___ValidationConfiguration = ValidationConfiguration
62
+
63
+ class TableSchema(google.protobuf.message.Message):
64
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
65
+ NAMEDCOLUMNS_FIELD_NUMBER: builtins.int
66
+ @property
67
+ def namedColumns(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___NamedColumn]: ...
68
+ def __init__(self,
69
+ *,
70
+ namedColumns: typing.Optional[typing.Iterable[global___NamedColumn]] = ...,
71
+ ) -> None: ...
72
+ def ClearField(self, field_name: typing_extensions.Literal["namedColumns",b"namedColumns"]) -> None: ...
73
+ global___TableSchema = TableSchema
74
+
75
+ class NamedColumn(google.protobuf.message.Message):
76
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
77
+ NAME_FIELD_NUMBER: builtins.int
78
+ COLUMNTYPE_FIELD_NUMBER: builtins.int
79
+ name: typing.Text
80
+ @property
81
+ def columnType(self) -> global___ColumnType: ...
82
+ def __init__(self,
83
+ *,
84
+ name: typing.Optional[typing.Text] = ...,
85
+ columnType: typing.Optional[global___ColumnType] = ...,
86
+ ) -> None: ...
87
+ def HasField(self, field_name: typing_extensions.Literal["_name",b"_name","columnType",b"columnType","name",b"name"]) -> builtins.bool: ...
88
+ def ClearField(self, field_name: typing_extensions.Literal["_name",b"_name","columnType",b"columnType","name",b"name"]) -> None: ...
89
+ def WhichOneof(self, oneof_group: typing_extensions.Literal["_name",b"_name"]) -> typing.Optional[typing_extensions.Literal["name"]]: ...
90
+ global___NamedColumn = NamedColumn
91
+
92
+ class ColumnType(google.protobuf.message.Message):
93
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
94
+ PRIMITIVETYPE_FIELD_NUMBER: builtins.int
95
+ NULLABLE_FIELD_NUMBER: builtins.int
96
+ primitiveType: global___PrimitiveType.ValueType
97
+ nullable: builtins.bool
98
+ def __init__(self,
99
+ *,
100
+ primitiveType: global___PrimitiveType.ValueType = ...,
101
+ nullable: builtins.bool = ...,
102
+ ) -> None: ...
103
+ def ClearField(self, field_name: typing_extensions.Literal["nullable",b"nullable","primitiveType",b"primitiveType"]) -> None: ...
104
+ global___ColumnType = ColumnType
105
+
106
+ class ComputationConfiguration(google.protobuf.message.Message):
107
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
108
+ SQLSTATEMENT_FIELD_NUMBER: builtins.int
109
+ PRIVACYSETTINGS_FIELD_NUMBER: builtins.int
110
+ CONSTRAINTS_FIELD_NUMBER: builtins.int
111
+ sqlStatement: typing.Text
112
+ @property
113
+ def privacySettings(self) -> global___PrivacySettings: ...
114
+ @property
115
+ def constraints(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Constraint]: ...
116
+ def __init__(self,
117
+ *,
118
+ sqlStatement: typing.Text = ...,
119
+ privacySettings: typing.Optional[global___PrivacySettings] = ...,
120
+ constraints: typing.Optional[typing.Iterable[global___Constraint]] = ...,
121
+ ) -> None: ...
122
+ def HasField(self, field_name: typing_extensions.Literal["privacySettings",b"privacySettings"]) -> builtins.bool: ...
123
+ def ClearField(self, field_name: typing_extensions.Literal["constraints",b"constraints","privacySettings",b"privacySettings","sqlStatement",b"sqlStatement"]) -> None: ...
124
+ global___ComputationConfiguration = ComputationConfiguration
125
+
126
+ class PrivacySettings(google.protobuf.message.Message):
127
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
128
+ MINAGGREGATIONGROUPSIZE_FIELD_NUMBER: builtins.int
129
+ minAggregationGroupSize: builtins.int
130
+ def __init__(self,
131
+ *,
132
+ minAggregationGroupSize: builtins.int = ...,
133
+ ) -> None: ...
134
+ def ClearField(self, field_name: typing_extensions.Literal["minAggregationGroupSize",b"minAggregationGroupSize"]) -> None: ...
135
+ global___PrivacySettings = PrivacySettings
136
+
137
+ class Constraint(google.protobuf.message.Message):
138
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
139
+ DESCRIPTION_FIELD_NUMBER: builtins.int
140
+ description: typing.Text
141
+ def __init__(self,
142
+ *,
143
+ description: typing.Text = ...,
144
+ ) -> None: ...
145
+ def ClearField(self, field_name: typing_extensions.Literal["description",b"description"]) -> None: ...
146
+ global___Constraint = Constraint
@@ -0,0 +1,16 @@
1
+ Metadata-Version: 2.1
2
+ Name: decentriq_util
3
+ Version: 0.1.0.dev0
4
+ Summary: Utility functions for Python code run in Decentriq workers
5
+ Author: hafenr
6
+ Author-email: robin.hafen@gmail.com
7
+ Requires-Python: >=3.9,<4.0
8
+ Classifier: Programming Language :: Python :: 3
9
+ Classifier: Programming Language :: Python :: 3.9
10
+ Classifier: Programming Language :: Python :: 3.10
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Requires-Dist: pandas (>=1.4.4)
14
+ Requires-Dist: protobuf (>=4.21,<5.0)
15
+ Requires-Dist: psutil (>=5.9.6,<6.0.0)
16
+ Requires-Dist: pyspark (>=3.4.0,<4.0.0)
@@ -0,0 +1,12 @@
1
+ decentriq_util/__init__.py,sha256=wjHRy7lrpmioI3iRT8pKUPTD0kqHMVKGb1p33U-vMb4,733
2
+ decentriq_util/error.py,sha256=MQn5zdLGzV31cy7KNtHMU1kVNUq1j9z3mzbMDBvfXvk,4798
3
+ decentriq_util/proto/__init__.py,sha256=KRWAX1OV9ZT041U4HotBJhgnGx265FQ8BoCYMMNfTQc,914
4
+ decentriq_util/python/__init__.py,sha256=NxsROEFbKKBy_WodSp0_vKNAnzr36FXlHfItid21knQ,873
5
+ decentriq_util/spark/__init__.py,sha256=J8GB_0rc7xpCXctJxQhGT2VEGT9YZQXvdc7ucrZ5X1c,7871
6
+ decentriq_util/sql/__init__.py,sha256=y8mjAuwtXoFufXwYHOazAX7XM5H_39LVOSELpDYYzB4,8863
7
+ decentriq_util/sql/proto/__init__.py,sha256=6AXw4_0BzJDiR8GClJvu9ztLJGif84jVeBieXyPsRJA,210
8
+ decentriq_util/sql/proto/compute_sql_pb2.py,sha256=gWYGaxUj1wMxSK7PMIAfpbl015AsNY9IFgc8kMGpnaA,5822
9
+ decentriq_util/sql/proto/compute_sql_pb2.pyi,sha256=ubGrxmUsCxO0rhozqBH8OLZ2HvolYGFnYIcazbnNQBo,6914
10
+ decentriq_util-0.1.0.dev0.dist-info/METADATA,sha256=96EdcOnvgjUEkq4IT-NdtWEWzAek-XvJzQ-Ghb0JDxs,610
11
+ decentriq_util-0.1.0.dev0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
12
+ decentriq_util-0.1.0.dev0.dist-info/RECORD,,
@@ -0,0 +1,4 @@
1
+ Wheel-Version: 1.0
2
+ Generator: poetry-core 1.8.1
3
+ Root-Is-Purelib: true
4
+ Tag: py3-none-any