backup

Author: davidkhala

format/davidkhala/data/format/arrow/__init__.py

@@ -1,5 +1,5 @@
-from pyarrow import Buffer, BufferReader
+from pyarrow import Buffer, BufferReader, NativeFile
 
 
-def read_data(buffer: bytes | Buffer):
+def read_data(buffer: bytes | Buffer)->NativeFile:
     return BufferReader(buffer)

format/davidkhala/data/format/arrow/fs.py

@@ -37,7 +37,7 @@ def open_input_stream(self, file: FileInfo) -> NativeFile:
         return self.fs.open_input_stream(file.path)
 
     def ls(self, base_dir: str) -> FileInfo | list[FileInfo]:
-        return self.fs.get_file_info(FileSelector(base_dir, recursive=True))
+        return self.fs.get_file_info(FileSelector(base_dir, recursive=True, allow_not_found=True))
 
     def write_stream(self, uri, tables_or_batches: Iterable[RecordBatch | Table]):
         with self.open_output_stream(uri) as stream:

format/davidkhala/data/format/arrow/gcp.py

@@ -1,4 +1,4 @@
-from davidkhala.gcp.auth import CredentialsInterface, ServiceAccountInfo
+from davidkhala.gcp.auth import OptionsInterface, ServiceAccountInfo
 from davidkhala.gcp.auth.options import from_service_account, ServiceAccount
 from pyarrow.fs import GcsFileSystem, FileInfo
 
@@ -12,21 +12,21 @@ class GCS(FS):
         - > pyarrow.lib.ArrowNotImplementedError: Append is not supported in GCS
     """
 
-    def __init__(self, public_bucket: bool = False, *, location='ASIA-EAST2', credentials: CredentialsInterface = None):
+    def __init__(self, public_bucket: bool = False, *, location='ASIA-EAST2', auth_options: OptionsInterface = None):
         options = {
             'anonymous': public_bucket,
             'default_bucket_location': location,
         }
-        if credentials:
-            options['access_token'] = credentials.token
-            options['credential_token_expiration'] = credentials.expiry
+        if auth_options:
+            options['access_token'] = auth_options.token
+            options['credential_token_expiration'] = auth_options.expiry
         self.fs = GcsFileSystem(**options)
 
     @staticmethod
     def from_service_account(info: ServiceAccountInfo):
         service_account = from_service_account(info)
         ServiceAccount.token.fget(service_account)  # credential validation included
-        return GCS(credentials=service_account.credentials)
+        return GCS(auth_options=service_account)
 
     def ls(self, bucket: str) -> FileInfo | list[FileInfo]:
         return super().ls(bucket)

format/davidkhala/data/format/avro/__init__.py

@@ -1,12 +1,19 @@
+from dataclasses import dataclass
 from typing import Iterator, IO
-
 import fastavro
+from fastavro import reader, writer
+from fastavro.types import Schema
+
+
+@dataclass
+class Data:
+    schema: Schema
+    records: Iterator[dict]
 
 
-def read(content) -> Iterator[dict]:
-    reader = fastavro.reader(content)
-    for record in reader:
-        yield record
+def read(content) -> (Iterator[dict], Schema):
+    _reader = reader(content)
+    return (_ for _ in _reader), _reader.writer_schema,
 
 
 def is_avro(file_path: str):
@@ -15,3 +22,7 @@ def is_avro(file_path: str):
 
 def is_avro_data(buffer: IO):
     return fastavro.is_avro(buffer)
+
+
+def write(output_stream, schema: Schema, records: Iterator[dict]):
+    writer(output_stream, schema, records)

format/davidkhala/data/format/transform/__init__.py

@@ -0,0 +1,60 @@
+from typing import cast
+
+from pyarrow import DataType, Table, ListType, StructType
+from pyarrow.types import (is_int8, is_int16, is_int32, is_int64,
+                           is_float32, is_float64,
+                           is_string, is_boolean, is_timestamp,
+                           is_date32, is_list, is_struct,
+                           )
+
+
+class Arrow2Avro:
+    arrow: Table
+
+    def __init__(self, table: Table):
+        self.arrow = table
+
+    @staticmethod
+    def type(arrow_type: DataType):
+        if is_int8(arrow_type) or is_int16(arrow_type) or is_int32(arrow_type):
+            return "int"
+        elif is_int64(arrow_type):
+            return "long"
+        elif is_float32(arrow_type):
+            return "float"
+        elif is_float64(arrow_type):
+            return "double"
+        elif is_string(arrow_type):
+            return "string"
+        elif is_boolean(arrow_type):
+            return "boolean"
+        elif is_timestamp(arrow_type):
+            return {"type": "long", "logicalType": "timestamp-millis"}
+        elif is_date32(arrow_type):
+            return {"type": "int", "logicalType": "date"}
+        elif is_list(arrow_type):
+
+            return {"type": "array", "items": Arrow2Avro.type(cast(ListType, arrow_type).value_type)}
+        elif is_struct(arrow_type):
+            return {
+                "type": "record",
+                "name": "struct",
+                "fields": [{"name": field.name, "type": Arrow2Avro.type(field.type)} for field in cast(StructType, arrow_type)]
+            }
+        else:
+            raise ValueError(f"Unsupported PyArrow type: {arrow_type}")
+
+    @property
+    def schema(self):
+        return {
+            "type": "record",
+            "name": "Root",
+            "fields": list(map(lambda _field: {
+                "name": _field.name,
+                "type": Arrow2Avro.type(_field.type)
+            }, self.arrow.schema))
+        }
+
+    @property
+    def records(self):
+        return self.arrow.to_pylist()

format/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "davidkhala.data.format"
-version = "0.0.2"
+version = "0.0.3"
 description = ""
 authors = ["David Liu <[email protected]>"]
 readme = "README.md"
@@ -11,7 +11,7 @@ python = "^3.10"
 # for extras
 fastavro = { version = "*", optional = true }
 pyarrow = { version = "*", optional = true }
-davidkhala-gcp = { version = "*", optional = true, extras = ["auth"] }
+davidkhala-gcp = { version = "*", optional = true}
 [tool.poetry.group.dev.dependencies]
 pytest = "*"
 

format/tests/avro_test.py

@@ -1,12 +1,36 @@
 import unittest
-from davidkhala.data.format.avro import read, is_avro
+
+from pyarrow import table, array, int32, string, list_, float32
+
+from davidkhala.data.format.avro import read, is_avro, write
+from davidkhala.data.format.parquet import Parquet
+from davidkhala.data.format.transform import Arrow2Avro
 
 
 class AvroTestCase(unittest.TestCase):
+    _path = 'fixtures/gcp-data-davidkhala.dbt_davidkhala.country_codes.avro'
+    def setUp(self):
+        parquet = Parquet('fixtures/gcp-data-davidkhala.dbt_davidkhala.country_codes.parquet')
+
+        t = Arrow2Avro(parquet.read_batch())
+        with open(self._path,'wb' ) as output_stream:
+            write(output_stream, t.schema, t.records)
+
+
+    def test_transform(self):
+        sample_table = table({
+            "id": array([1, 2, 3], type=int32()),
+            "name": array(["Alice", "Bob", "Charlie"], type=string()),
+            "scores": array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], type=list_(float32()))
+        })
+        t= Arrow2Avro(sample_table)
+
+        with open("artifacts/dummy.avro", "wb") as out:
+            write(out, t.schema, t.records)
+
     def test_read(self):
-        _path = 'fixtures/gcp-data-davidkhala.dbt_davidkhala.country_codes.avro'
-        self.assertTrue(is_avro(_path))
-        with open(_path, 'rb') as file:
+        self.assertTrue(is_avro(self._path))
+        with open(self._path, 'rb') as file:
             for record in read(file):
                 print(record)
 

format/tests/ci.py

@@ -30,7 +30,7 @@ def test_memory_map(self):
         read_table(parquet_path, memory_map=False)
         self.assertEqual(total_allocated_bytes(), 256)
         read_table(parquet_path, memory_map=True)
-        self.assertEqual(total_allocated_bytes(), 256)
+        print(total_allocated_bytes())
 
 
 if __name__ == '__main__':