0
\$\begingroup\$

This is my first Apache pipeline. It takes a JSON file and saves the correctly formatted rows in one table, and the misformed rows in another.

My biggest worry is the efficiency. I have tested the code with the tiny (cr. 700 rows) file. The real input file is huge, it contains tens of millions of rows.

The error messages, PEP8 compliance etc. could have been better, but it's the least worry, for now at least.

The JSON format of the file is a bit unusual, the data is stored as attributes, thus the use of xml.etree.ElementTree.

from datetime import datetime
import os
import apache_beam as beam
import xml.etree.ElementTree as element_tree

temp_bucket = os.environ.get('google_temp_bucket', None)

wtbq_dict = {"method":"STREAMING_INSERTS"} if temp_bucket is None else {"custom_gcs_temp_location" : temp_bucket}
wtbq_dict["write_disposition"] = beam.io.BigQueryDisposition.WRITE_APPEND if temp_bucket is None else beam.io.BigQueryDisposition.WRITE_TRUNCATE
wtbq_dict["create_disposition"] = beam.io.BigQueryDisposition.CREATE_IF_NEEDED
wtbq_dict["ignore_unknown_columns"] = True

def to_datetime(s):
  # input example: '2023-07-17T19:24:01.893'
  result = datetime.strptime(s[:s.find('.')], '%Y-%m-%dT%H:%M:%S')
  return result

types_map = { int : 'INTEGER', str : 'STRING', datetime : 'TIMESTAMP'}

results_structure = [('Id', int, 'id'), ('ViewCount', int, 'view_count'), ('CreationDate', datetime, 'creation_date') ]
results_table_schema = {'fields': [{'name' : col_name, 'type': types_map[col_type], 'mode': 'Required'} for _, col_type, col_name in results_structure] }

errors_structure = [(int, 'i'), (str, 'id'), (str, 'missing_keys'), (str, 'wrong_type_keys') ]
errors_table_schema = {'fields': [{'name' : col_name, 'type': types_map[col_type], 'mode': 'Required'} for col_type, col_name in errors_structure] }

def result_or_error(inp, _):
  return 0 if 'i' not in inp else 1

def tree_elem_to_dict(i, el):
  result = {}
  missing_keys = []
  wrong_type_keys = []
  for tree_key, transform_type, result_key in results_structure:
    if tree_key not in el.attrib:
      missing_keys.append(tree_key)
      continue
    try:
      transform_func = to_datetime if transform_type is datetime else transform_type
      result[result_key] = transform_func(el.attrib[tree_key])
    except:
      wrong_type_keys.append([el.attrib[tree_key], transform_type, type(el.attrib[tree_key])])
  if (missing_keys or wrong_type_keys):
    return {'i' : i, 
            'id' : el.attrib['Id'] if 'Id' in el.attrib else 'unknown', 
            'missing_keys' : str(missing_keys), 
            'wrong_type_keys' : str(wrong_type_keys)}
  else:
    return result

def parse_into_dict(filename):
  tree = element_tree.parse(filename)
  results = [tree_elem_to_dict(i, el) for i, el in enumerate(tree.iterfind('row'))]
  return results

def run_a_posts_pipeline(project_id, filename, wtbq_dict):
  beam_options = beam.options.pipeline_options.PipelineOptions()

  with beam.Pipeline(options=beam_options) as pipeline:
    results, errors = ( pipeline | 'file_to_dict' >> beam.Create(parse_into_dict(filename))
                                 | 'Partition' >> beam.Partition(result_or_error, 2))
    
    results | 'tobq_results' >> beam.io.WriteToBigQuery(table=f'{project_id}:dummy2_dataset.dummy_posts_results',
                                                   schema=results_table_schema,                  
                                                   **wtbq_dict)

    errors  | 'tobq_errors'  >> beam.io.WriteToBigQuery(table=f'{project_id}:dummy2_dataset.dummy_posts_errors',
                                                             schema=errors_table_schema,                  
                                                             **wtbq_dict)

run_a_posts_pipeline(project_id="project_id", filename='filename.xml', wtbq_dict=wtbq_dict)
\$\endgroup\$

1 Answer 1

1
\$\begingroup\$

imports

It would be convenient to use from ... import ..., so we can write WRITE_APPEND rather than beam.io.BigQueryDisposition.WRITE_APPEND.

I don't know what "wt" in "wtbq" big query means, but that's fine.

Using two-space indent is normal for some languages, but absolutely not in python. It makes the source code harder to read than necessary. Use black -S *.py to fix up your source every now and again.

validate args

def to_datetime(s):
  # input example: '2023-07-17T19:24:01.893'
  result = datetime.strptime(s[:s.find('.')], '%Y-%m-%dT%H:%M:%S')

Thank you for the example in the comment; that's helpful.

Directly returning the strptime expression would have been fine -- no need to name it result.

The .find() call should be .index(). Or we should assert '.' in s. Or we should document that "dot is optional", and behave gracefully when it's absent.

I feel it is mostly reasonable for caller to pass in
'2023-07-17T18:24:19', despite that helpful comment. It is entirely unreasonable that caller gets
datetime(2023, 7, 17, 18, 24, 1), which is eighteen seconds off, with no indication that things went badly. Either raise fatal error, or compute a more plausible result. As things stand this is minimally a documentation defect, and would likely be viewed as a code defect.

state things in the positive

def result_or_error(inp, _):
    return 0 if 'i' not in inp else 1

This would be slightly easier for humans to read if it were phrased:

def result_or_error(inp, _):
    return 1 if 'i' in inp else 0

and much easier if phrased:

def result_or_error(inp, _) -> int:
    return int('i' in inp)

The name is misleading -- on my initial reading of the signature I thought we would get back either a proper result or some failure object like None. Eventually I read that it's responsible for two-bin partitioning.

design of public API

def tree_elem_to_dict(i, el):

It's pretty clear that 2nd parameter is a tree element. But declaring i: int would have been helpful. Or even better, offer a """docstring""" that mentions the meaning of i.

            'id' : el.attrib['Id'] if 'Id' in el.attrib else 'unknown', 

This defaulting is more clearly expressed as:

            'id' : el.attrib.get('Id', 'unknown'), 

These str() calls are surprising:

            'missing_keys' : str(missing_keys), 
            'wrong_type_keys' : str(wrong_type_keys)

We can't keep them as lists? Ok, fine. Consider using repr() instead, so quoting issues are less worrisome.

plural

This is manifestly the wrong name:

def parse_into_dict(filename):

Clearly it returns a list of dicts.

type annotation

This is clear as written:

def run_a_posts_pipeline(project_id, filename, wtbq_dict):

It would be preferable to rename 3rd parameter:

def run_a_posts_pipeline(project_id: int, filename: Path, wtbq: dict) -> None:

performance

The review context suggests that speed matters.

Yet this submission's source code contains no description of

  • typical JSONL line size or complexity
  • observed app level throughput, in lines per second
  • automated tests
  • anticipated input line error rate
  • profiling

Minimally I anticipated that run_a_posts_pipeline would log a line count and elapsed time.

We don't know if most of the time was spent in this code or in the ElementTree library. It is possible that competing libraries would perform better on this task, but first we'd want some measurements.

Suppose this code runs in production for a couple of months, and then a maintenance engineer makes some edits. How would we know if a performance regression was introduced?


This code appears to achieve its design goals.

It is written in a reasonably clear style. Despite the fact that I don't know exactly what it does, I would be willing to delegate or accept maintenance tasks on it.

\$\endgroup\$
0

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.