|
| 1 | +from typing import Dict, List, Any, NamedTuple, Optional, final, Literal |
| 2 | +import pyarrow.csv as csv |
| 3 | +from dataclasses import dataclass |
| 4 | +import dataclasses |
| 5 | +from abc import ABC |
| 6 | +import pyarrow.parquet as pq |
| 7 | +import pyarrow as pa |
| 8 | +import json |
| 9 | +from pathlib import Path |
| 10 | +from cloud2sql.arrow.model import ArrowModel |
| 11 | +from cloud2sql.schema_utils import insert_node |
| 12 | +from resotoclient.models import JsObject |
| 13 | + |
| 14 | + |
| 15 | +class WriteResult(NamedTuple): |
| 16 | + table_name: str |
| 17 | + |
| 18 | + |
| 19 | +class FileWriter(ABC): |
| 20 | + pass |
| 21 | + |
| 22 | + |
| 23 | +@final |
| 24 | +@dataclass(frozen=True) |
| 25 | +class Parquet(FileWriter): |
| 26 | + parquet_writer: pq.ParquetWriter |
| 27 | + |
| 28 | + |
| 29 | +@final |
| 30 | +@dataclass(frozen=True) |
| 31 | +class CSV(FileWriter): |
| 32 | + csv_writer: csv.CSVWriter |
| 33 | + |
| 34 | + |
| 35 | +@final |
| 36 | +@dataclass |
| 37 | +class ArrowBatch: |
| 38 | + table_name: str |
| 39 | + rows: List[Dict[str, Any]] |
| 40 | + schema: pa.Schema |
| 41 | + writer: FileWriter |
| 42 | + |
| 43 | + |
| 44 | +class ConversionTarget(ABC): |
| 45 | + pass |
| 46 | + |
| 47 | + |
| 48 | +@final |
| 49 | +@dataclass(frozen=True) |
| 50 | +class ParquetMap(ConversionTarget): |
| 51 | + convert_values_to_str: bool |
| 52 | + |
| 53 | + |
| 54 | +@final |
| 55 | +@dataclass(frozen=True) |
| 56 | +class ParquetString(ConversionTarget): |
| 57 | + pass |
| 58 | + |
| 59 | + |
| 60 | +@dataclass |
| 61 | +class NormalizationPath: |
| 62 | + path: List[Optional[str]] |
| 63 | + convert_to: ConversionTarget |
| 64 | + |
| 65 | + |
| 66 | +# workaround until fix is merged https://issues.apache.org/jira/browse/ARROW-17832 |
| 67 | +# |
| 68 | +# here we collect the paths to the JSON object fields that we want to convert to arrow types |
| 69 | +# so that later we will do the transformations |
| 70 | +# |
| 71 | +# currently we do dict -> list[(key, value)] and converting values which are defined as strings |
| 72 | +# in the scheme to strings/json strings |
| 73 | +def collect_normalization_paths(schema: pa.Schema) -> List[NormalizationPath]: |
| 74 | + paths: List[NormalizationPath] = [] |
| 75 | + |
| 76 | + def collect_paths_to_maps_helper(path: List[Optional[str]], typ: pa.DataType) -> None: |
| 77 | + # if we see a map, then full stop. we add the path to the list |
| 78 | + # if the value type is string, we remember that too |
| 79 | + if isinstance(typ, pa.lib.MapType): |
| 80 | + stringify_items = pa.types.is_string(typ.item_type) |
| 81 | + normalization_path = NormalizationPath(path, ParquetMap(stringify_items)) |
| 82 | + paths.append(normalization_path) |
| 83 | + # structs are traversed but they have no interest for us |
| 84 | + elif isinstance(typ, pa.lib.StructType): |
| 85 | + for field_idx in range(0, typ.num_fields): |
| 86 | + field = typ.field(field_idx) |
| 87 | + collect_paths_to_maps_helper(path + [field.name], field.type) |
| 88 | + # the lists traversed too. None will is added to the path to be consumed by the recursion |
| 89 | + # in order to reach the correct level |
| 90 | + elif isinstance(typ, pa.lib.ListType): |
| 91 | + collect_paths_to_maps_helper(path + [None], typ.value_type) |
| 92 | + # if we see a string, then we stop and add a path to the list |
| 93 | + elif pa.types.is_string(typ): |
| 94 | + normalization_path = NormalizationPath(path, ParquetString()) |
| 95 | + paths.append(normalization_path) |
| 96 | + |
| 97 | + # bootstrap the recursion |
| 98 | + for idx, typ in enumerate(schema.types): |
| 99 | + collect_paths_to_maps_helper([schema.names[idx]], typ) |
| 100 | + |
| 101 | + return paths |
| 102 | + |
| 103 | + |
| 104 | +def normalize(npath: NormalizationPath, obj: Any) -> Any: |
| 105 | + path = npath.path |
| 106 | + reached_target = len(path) == 0 |
| 107 | + |
| 108 | + # we're on the correct node, time to convert it into something |
| 109 | + if reached_target: |
| 110 | + if isinstance(npath.convert_to, ParquetString): |
| 111 | + # everything that should be a string is either a string or a json string |
| 112 | + return obj if isinstance(obj, str) else json.dumps(obj) |
| 113 | + elif isinstance(npath.convert_to, ParquetMap): |
| 114 | + # we can only convert dicts to maps. if it is not the case then it is a bug |
| 115 | + if not isinstance(obj, dict): |
| 116 | + raise Exception(f"Expected dict, got {type(obj)}. path: {npath}") |
| 117 | + |
| 118 | + def value_to_string(v: Any) -> str: |
| 119 | + if isinstance(v, str): |
| 120 | + return v |
| 121 | + else: |
| 122 | + return json.dumps(v) |
| 123 | + |
| 124 | + # in case the map should contain string values, we convert them to strings |
| 125 | + if npath.convert_to.convert_values_to_str: |
| 126 | + return [(k, value_to_string(v)) for k, v in obj.items()] |
| 127 | + else: |
| 128 | + return list(obj.items()) |
| 129 | + # we're not at the target node yet, so we traverse the tree deeper |
| 130 | + else: |
| 131 | + # if we see a dict, we try to go deeper in case it contains the key we are looking for |
| 132 | + # otherwise we return the object as is. This is valid because the fields are optional |
| 133 | + if isinstance(obj, dict): |
| 134 | + key = path[0] |
| 135 | + if key in obj: |
| 136 | + # consume the current element of the path |
| 137 | + new_npath = dataclasses.replace(npath, path=path[1:]) |
| 138 | + obj[key] = normalize(new_npath, obj[key]) |
| 139 | + return obj |
| 140 | + # in case of a list, we process all its elements |
| 141 | + elif isinstance(obj, list): |
| 142 | + # check that the path is correct |
| 143 | + assert path[0] is None |
| 144 | + # consume the current element of the path |
| 145 | + new_npath = dataclasses.replace(npath, path=path[1:]) |
| 146 | + return [normalize(new_npath, v) for v in obj] |
| 147 | + else: |
| 148 | + raise Exception(f"Unexpected object type {type(obj)}, path: {npath}") |
| 149 | + |
| 150 | + |
| 151 | +def write_batch_to_file(batch: ArrowBatch) -> ArrowBatch: |
| 152 | + |
| 153 | + to_normalize = collect_normalization_paths(batch.schema) |
| 154 | + |
| 155 | + for row in batch.rows: |
| 156 | + for path in to_normalize: |
| 157 | + normalize(path, row) |
| 158 | + |
| 159 | + pa_table = pa.Table.from_pylist(batch.rows, batch.schema) |
| 160 | + if isinstance(batch.writer, Parquet): |
| 161 | + batch.writer.parquet_writer.write_table(pa_table) |
| 162 | + elif isinstance(batch.writer, CSV): |
| 163 | + batch.writer.csv_writer.write_table(pa_table) |
| 164 | + else: |
| 165 | + raise ValueError(f"Unknown format {batch.writer}") |
| 166 | + return ArrowBatch(table_name=batch.table_name, rows=[], schema=batch.schema, writer=batch.writer) |
| 167 | + |
| 168 | + |
| 169 | +def close_writer(batch: ArrowBatch) -> None: |
| 170 | + if isinstance(batch.writer, Parquet): |
| 171 | + batch.writer.parquet_writer.close() |
| 172 | + elif isinstance(batch.writer, CSV): |
| 173 | + batch.writer.csv_writer.close() |
| 174 | + else: |
| 175 | + raise ValueError(f"Unknown format {batch.writer}") |
| 176 | + |
| 177 | + |
| 178 | +def new_writer(format: Literal["parquet", "csv"], table_name: str, schema: pa.Schema, result_dir: Path) -> FileWriter: |
| 179 | + def ensure_path(path: Path) -> Path: |
| 180 | + path.mkdir(parents=True, exist_ok=True) |
| 181 | + return path |
| 182 | + |
| 183 | + if format == "parquet": |
| 184 | + return Parquet(pq.ParquetWriter(Path(ensure_path(result_dir), f"{table_name}.parquet"), schema=schema)) |
| 185 | + elif format == "csv": |
| 186 | + return CSV(csv.CSVWriter(Path(ensure_path(result_dir), f"{table_name}.csv"), schema=schema)) |
| 187 | + else: |
| 188 | + raise ValueError(f"Unknown format {format}") |
| 189 | + |
| 190 | + |
| 191 | +class ArrowWriter: |
| 192 | + def __init__( |
| 193 | + self, model: ArrowModel, result_directory: Path, rows_per_batch: int, output_format: Literal["parquet", "csv"] |
| 194 | + ): |
| 195 | + self.model = model |
| 196 | + self.kind_by_id: Dict[str, str] = {} |
| 197 | + self.batches: Dict[str, ArrowBatch] = {} |
| 198 | + self.rows_per_batch: int = rows_per_batch |
| 199 | + self.result_directory: Path = result_directory |
| 200 | + self.output_format: Literal["parquet", "csv"] = output_format |
| 201 | + |
| 202 | + def insert_value(self, table_name: str, values: Any) -> Optional[WriteResult]: |
| 203 | + if self.model.schemas.get(table_name): |
| 204 | + schema = self.model.schemas[table_name] |
| 205 | + batch = self.batches.get( |
| 206 | + table_name, |
| 207 | + ArrowBatch( |
| 208 | + table_name, |
| 209 | + [], |
| 210 | + schema, |
| 211 | + new_writer(self.output_format, table_name, schema, self.result_directory), |
| 212 | + ), |
| 213 | + ) |
| 214 | + |
| 215 | + batch.rows.append(values) |
| 216 | + self.batches[table_name] = batch |
| 217 | + return WriteResult(table_name) |
| 218 | + return None |
| 219 | + |
| 220 | + def insert_node(self, node: JsObject) -> None: |
| 221 | + result = insert_node( |
| 222 | + node, |
| 223 | + self.kind_by_id, |
| 224 | + self.insert_value, |
| 225 | + with_tmp_prefix=False, |
| 226 | + ) |
| 227 | + should_write_batch = result and len(self.batches[result.table_name].rows) > self.rows_per_batch |
| 228 | + if result and should_write_batch: |
| 229 | + batch = self.batches[result.table_name] |
| 230 | + self.batches[result.table_name] = write_batch_to_file(batch) |
| 231 | + |
| 232 | + def close(self) -> None: |
| 233 | + for table_name, batch in self.batches.items(): |
| 234 | + batch = write_batch_to_file(batch) |
| 235 | + self.batches[table_name] = batch |
| 236 | + close_writer(batch) |
0 commit comments