Unzip Transformer, DID Finder, Pandas Python Code Generators

.github/workflows/deploy-config.json

@@ -19,6 +19,11 @@
     "image_name": "servicex-did-finder",
     "test_required": false
   },
+  {
+    "dir_name": "did_finder_unzip",
+    "image_name": "servicex-did-finder-unzip",
+    "test_required": false
+  },
   {
     "dir_name": "code_generator_funcadl_uproot",
     "image_name": "servicex_code_gen_func_adl_uproot",

code_generator_python/Dockerfile

@@ -20,6 +20,7 @@ COPY boot.sh ./
 COPY transformer_capabilities.json ./
 COPY servicex/ ./servicex
 COPY scripts/from_text_to_zip.py .
+COPY transformer_capabilities.json ./
 RUN chmod +x boot.sh
 
 USER servicex

code_generator_python/servicex/python_code_generator/python_translator.py

@@ -38,6 +38,7 @@ class PythonTranslator(CodeGenerator):
     def generate_code(self, query, cache_path: str):
 
         src = base64.b64decode(query).decode('ascii')
+        print("SRC", src)
         hash = "no-hash"
         query_file_path = os.path.join(cache_path, hash)
 
@@ -59,6 +60,10 @@ def generate_code(self, query, cache_path: str):
         shutil.copyfile(capabilities_path, os.path.join(query_file_path,
                                                         "transformer_capabilities.json"))
 
+        with open(os.path.join(query_file_path, 'generated_transformer.py'), 'w') as python_file:
+            python_file.write(src)
+
         os.system("ls -lht " + query_file_path)
         os.system(f"cat {query_file_path}/generated_transformer.py")
+
         return GeneratedFileResult(hash, query_file_path)

code_generator_python/servicex/python_code_generator/unzip_translator.py

@@ -0,0 +1,13 @@
+def run_query(input_filenames=None, tree_name=None):
+    from stream_unzip import stream_unzip
+    import httpx
+
+    def zipped_chunks(input_path):
+        # Iterable that yields the bytes of a zip file
+        with httpx.stream('GET', input_path) as r:
+            yield from r.iter_bytes(chunk_size=65536)
+
+    for file_name, file_size, unzipped_chunks in stream_unzip(zipped_chunks(input_filenames[0])):
+        # unzipped_chunks must be iterated to completion or UnfinishedIterationError will be raised
+        for chunk in unzipped_chunks:
+            yield chunk, file_name

code_generator_python/servicex/templates/transform_single_file.py

@@ -8,6 +8,7 @@
 import uproot
 instance = os.environ.get('INSTANCE_NAME', 'Unknown')
 default_tree_name = "servicex"
+codegen_type = os.environ.get('CODEGEN_TYPE', 'default')
 
 
 def transform_single_file(file_path: str, output_path: Path, output_format: str):
@@ -20,41 +21,58 @@ def transform_single_file(file_path: str, output_path: Path, output_format: str)
     try:
         stime = time.time()
 
-        awkward_array = generated_transformer.run_query(file_path)
-        total_events = ak.num(awkward_array, axis=0)
+        if codegen_type == 'default':
+            awkward_array = generated_transformer.run_query(file_path)
+            total_events = ak.num(awkward_array, axis=0)
 
-        ttime = time.time()
+            ttime = time.time()
 
-        if output_format == 'root-file':
-            etime = time.time()
-            with uproot.recreate(output_path) as writer:
-                writer[default_tree_name] = {field: awkward_array[field] for field in
-                                             awkward_array.fields} if awkward_array.fields \
-                    else awkward_array
-            wtime = time.time()
+            if output_format == 'root-file':
+                etime = time.time()
+                with uproot.recreate(output_path) as writer:
+                    writer[default_tree_name] = {field: awkward_array[field] for field in
+                                                 awkward_array.fields} if awkward_array.fields \
+                        else awkward_array
+                wtime = time.time()
 
-        else:
-            explode_records = bool(awkward_array.fields)
-            try:
-                arrow = ak.to_arrow_table(awkward_array, explode_records=explode_records)
-            except TypeError:
-                arrow = ak.to_arrow_table(ak.repartition(awkward_array, None),
-                                          explode_records=explode_records)
+            else:
+                explode_records = bool(awkward_array.fields)
+                try:
+                    arrow = ak.to_arrow_table(awkward_array, explode_records=explode_records)
+                except TypeError:
+                    arrow = ak.to_arrow_table(ak.repartition(awkward_array, None),
+                                              explode_records=explode_records)
 
-            etime = time.time()
+                etime = time.time()
 
-            writer = pq.ParquetWriter(output_path, arrow.schema)
-            writer.write_table(table=arrow)
-            writer.close()
+                writer = pq.ParquetWriter(output_path, arrow.schema)
+                writer.write_table(table=arrow)
+                writer.close()
 
-            wtime = time.time()
+                wtime = time.time()
 
-        output_size = os.stat(output_path).st_size
-        print(f'Detailed transformer times. query_time:{round(ttime - stime, 3)} '
-              f'serialization: {round(etime - ttime, 3)} '
-              f'writing: {round(wtime - etime, 3)}')
+            output_size = os.stat(output_path).st_size
+            print(f'Detailed transformer times. query_time:{round(ttime - stime, 3)} '
+                  f'serialization: {round(etime - ttime, 3)} '
+                  f'writing: {round(wtime - etime, 3)}')
 
-        print(f"Transform stats: Total Events: {total_events}, resulting file size {output_size}")
+            print(f"Transform stats: Total Events: {total_events}, \
+                   resulting file size {output_size}")
+        elif codegen_type == 'unzip':
+            folder_output_path = os.path.dirname(output_path)
+            for bytes, file_name in generated_transformer.run_query(file_path):
+                file_output_path = os.path.join(folder_output_path, file_name.decode('utf-8'))
+                print('File: ', file_output_path)
+                with open(file_output_path, 'ab') as f:
+                    f.write(bytes)
+            total_events = 0
+            output_size = os.stat(folder_output_path).st_size
+        elif codegen_type == 'pandas':
+            folder_output_path = os.path.dirname(output_path)
+            pd_data = generated_transformer.run_query(file_path)
+            pd_data.to_parquet(output_path)
+            total_events = 0
+            output_size = os.stat(folder_output_path).st_size
     except Exception as error:
         mesg = f"Failed to transform input file {file_path}: {error}"
         print(mesg)

code_generator_python/tests/test_python_translator.py

@@ -31,12 +31,13 @@
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are met:
 #
+
 import base64
 import os
 import tempfile
 
 from servicex.python_code_generator.python_translator import \
-    PythonTranslator
+     PythonTranslator
 
 
 def test_generate_code():

code_generator_python/transformer_capabilities.json

@@ -4,6 +4,6 @@
   "limitations": "Would be good to note what isn't implemented",
   "file-formats": ["parquet", "root"],
   "stats-parser": "UprootStats",
-  "language": "python",
+  "language": "python3",
   "command": "/generated/transform_single_file.py"
 }

did_finder_rucio/scripts/did_finder.py

@@ -69,6 +69,7 @@ async def callback(did_name, info):
                 request_id=info['request-id']
             )
             for file in lookup_request.lookup_files():
+                logger.info(f"File: {file}")
                 yield file
 
         start_did_finder('rucio',

did_finder_unzip/Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3
+
+# Create app directory
+WORKDIR /usr/src/app
+
+# for CA certificates
+USER root
+RUN mkdir -p /etc/grid-security/certificates /etc/grid-security/vomsdir
+
+ENV POETRY_VERSION=1.2.2
+RUN python3 -m pip install --upgrade pip
+
+RUN pip install poetry==$POETRY_VERSION
+COPY pyproject.toml pyproject.toml
+COPY poetry.lock poetry.lock
+
+RUN poetry config virtualenvs.create false && \
+    poetry install --no-root --no-interaction --no-ansi
+
+COPY . .
+
+ENV X509_USER_PROXY /tmp/grid-security/x509up
+ENV X509_CERT_DIR /etc/grid-security/certificates
+