188 lines
8.4 KiB
YAML
188 lines
8.4 KiB
YAML
name: BigQueryExampleGen
|
|
description: |-
|
|
Official TFX BigQueryExampleGen component.
|
|
|
|
The BigQuery examplegen component takes a query, and generates train
|
|
and eval examples for downsteam components.
|
|
|
|
|
|
Args:
|
|
query: BigQuery sql string, query result will be treated as a single
|
|
split, can be overwritten by input_config.
|
|
input_config: An example_gen_pb2.Input instance with Split.pattern as
|
|
BigQuery sql string. If set, it overwrites the 'query' arg, and allows
|
|
different queries per split.
|
|
output_config: An example_gen_pb2.Output instance, providing output
|
|
configuration. If unset, default splits will be 'train' and 'eval' with
|
|
size 2:1.
|
|
Returns:
|
|
examples: Optional channel of 'ExamplesPath' for output train and
|
|
eval examples.
|
|
|
|
Raises:
|
|
RuntimeError: Only one of query and input_config should be set.
|
|
inputs:
|
|
- {name: query, type: String, optional: true}
|
|
- name: input_config
|
|
type:
|
|
JsonObject: {data_type: 'proto:tfx.components.example_gen.Input'}
|
|
optional: true
|
|
- name: output_config
|
|
type:
|
|
JsonObject: {data_type: 'proto:tfx.components.example_gen.Output'}
|
|
optional: true
|
|
- name: custom_config
|
|
type:
|
|
JsonObject: {data_type: 'proto:tfx.components.example_gen.CustomConfig'}
|
|
optional: true
|
|
outputs:
|
|
- {name: examples, type: Examples}
|
|
implementation:
|
|
container:
|
|
image: tensorflow/tfx:0.21.4
|
|
command:
|
|
- python3
|
|
- -u
|
|
- -c
|
|
- |
|
|
def _make_parent_dirs_and_return_path(file_path: str):
|
|
import os
|
|
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
|
return file_path
|
|
|
|
def BigQueryExampleGen(
|
|
examples_path ,
|
|
|
|
query = None,
|
|
input_config = None,
|
|
output_config = None,
|
|
custom_config = None,
|
|
):
|
|
"""
|
|
Official TFX BigQueryExampleGen component.
|
|
|
|
The BigQuery examplegen component takes a query, and generates train
|
|
and eval examples for downsteam components.
|
|
|
|
Args:
|
|
query: BigQuery sql string, query result will be treated as a single
|
|
split, can be overwritten by input_config.
|
|
input_config: An example_gen_pb2.Input instance with Split.pattern as
|
|
BigQuery sql string. If set, it overwrites the 'query' arg, and allows
|
|
different queries per split.
|
|
output_config: An example_gen_pb2.Output instance, providing output
|
|
configuration. If unset, default splits will be 'train' and 'eval' with
|
|
size 2:1.
|
|
Returns:
|
|
examples: Optional channel of 'ExamplesPath' for output train and
|
|
eval examples.
|
|
|
|
Raises:
|
|
RuntimeError: Only one of query and input_config should be set.
|
|
"""
|
|
from tfx.components.example_gen.csv_example_gen.component import BigQueryExampleGen as component_class
|
|
|
|
#Generated code
|
|
import json
|
|
import os
|
|
import tensorflow
|
|
from google.protobuf import json_format, message
|
|
from tfx.types import Artifact, channel_utils, artifact_utils
|
|
|
|
arguments = locals().copy()
|
|
|
|
component_class_args = {}
|
|
|
|
for name, execution_parameter in component_class.SPEC_CLASS.PARAMETERS.items():
|
|
argument_value_obj = argument_value = arguments.get(name, None)
|
|
if argument_value is None:
|
|
continue
|
|
parameter_type = execution_parameter.type
|
|
if isinstance(parameter_type, type) and issubclass(parameter_type, message.Message): # Maybe FIX: execution_parameter.type can also be a tuple
|
|
argument_value_obj = parameter_type()
|
|
json_format.Parse(argument_value, argument_value_obj)
|
|
component_class_args[name] = argument_value_obj
|
|
|
|
for name, channel_parameter in component_class.SPEC_CLASS.INPUTS.items():
|
|
artifact_path = arguments[name + '_path']
|
|
if artifact_path:
|
|
artifact = channel_parameter.type()
|
|
artifact.uri = artifact_path + '/' # ?
|
|
if channel_parameter.type.PROPERTIES and 'split_names' in channel_parameter.type.PROPERTIES:
|
|
# Recovering splits
|
|
subdirs = tensorflow.io.gfile.listdir(artifact_path)
|
|
artifact.split_names = artifact_utils.encode_split_names(sorted(subdirs))
|
|
component_class_args[name] = channel_utils.as_channel([artifact])
|
|
|
|
component_class_instance = component_class(**component_class_args)
|
|
|
|
input_dict = {name: channel.get() for name, channel in component_class_instance.inputs.get_all().items()}
|
|
output_dict = {name: channel.get() for name, channel in component_class_instance.outputs.get_all().items()}
|
|
exec_properties = component_class_instance.exec_properties
|
|
|
|
# Generating paths for output artifacts
|
|
for name, artifacts in output_dict.items():
|
|
base_artifact_path = arguments[name + '_path']
|
|
# Are there still cases where output channel has multiple artifacts?
|
|
for idx, artifact in enumerate(artifacts):
|
|
subdir = str(idx + 1) if idx > 0 else ''
|
|
artifact.uri = os.path.join(base_artifact_path, subdir) # Ends with '/'
|
|
|
|
print('component instance: ' + str(component_class_instance))
|
|
|
|
#executor = component_class.EXECUTOR_SPEC.executor_class() # Same
|
|
executor = component_class_instance.executor_spec.executor_class()
|
|
executor.Do(
|
|
input_dict=input_dict,
|
|
output_dict=output_dict,
|
|
exec_properties=exec_properties,
|
|
)
|
|
|
|
import argparse
|
|
_parser = argparse.ArgumentParser(prog='BigQueryExampleGen', description="Official TFX BigQueryExampleGen component.\n\n The BigQuery examplegen component takes a query, and generates train\n and eval examples for downsteam components.\n\n\n Args:\n query: BigQuery sql string, query result will be treated as a single\n split, can be overwritten by input_config.\n input_config: An example_gen_pb2.Input instance with Split.pattern as\n BigQuery sql string. If set, it overwrites the 'query' arg, and allows\n different queries per split.\n output_config: An example_gen_pb2.Output instance, providing output\n configuration. If unset, default splits will be 'train' and 'eval' with\n size 2:1.\n Returns:\n examples: Optional channel of 'ExamplesPath' for output train and\n eval examples.\n\n Raises:\n RuntimeError: Only one of query and input_config should be set.")
|
|
_parser.add_argument("--query", dest="query", type=str, required=False, default=argparse.SUPPRESS)
|
|
_parser.add_argument("--input-config", dest="input_config", type=str, required=False, default=argparse.SUPPRESS)
|
|
_parser.add_argument("--output-config", dest="output_config", type=str, required=False, default=argparse.SUPPRESS)
|
|
_parser.add_argument("--custom-config", dest="custom_config", type=str, required=False, default=argparse.SUPPRESS)
|
|
_parser.add_argument("--examples", dest="examples_path", type=_make_parent_dirs_and_return_path, required=True, default=argparse.SUPPRESS)
|
|
_parsed_args = vars(_parser.parse_args())
|
|
_output_files = _parsed_args.pop("_output_paths", [])
|
|
|
|
_outputs = BigQueryExampleGen(**_parsed_args)
|
|
|
|
_output_serializers = [
|
|
|
|
]
|
|
|
|
import os
|
|
for idx, output_file in enumerate(_output_files):
|
|
try:
|
|
os.makedirs(os.path.dirname(output_file))
|
|
except OSError:
|
|
pass
|
|
with open(output_file, 'w') as f:
|
|
f.write(_output_serializers[idx](_outputs[idx]))
|
|
args:
|
|
- if:
|
|
cond: {isPresent: query}
|
|
then:
|
|
- --query
|
|
- {inputValue: query}
|
|
- if:
|
|
cond: {isPresent: input_config}
|
|
then:
|
|
- --input-config
|
|
- {inputValue: input_config}
|
|
- if:
|
|
cond: {isPresent: output_config}
|
|
then:
|
|
- --output-config
|
|
- {inputValue: output_config}
|
|
- if:
|
|
cond: {isPresent: custom_config}
|
|
then:
|
|
- --custom-config
|
|
- {inputValue: custom_config}
|
|
- --examples
|
|
- {outputPath: examples}
|