def testConstruct(self):
big_query_to_elwc_example_gen = component.BigQueryToElwcExampleGen(
query='query',
elwc_config=elwc_config_pb2.ElwcConfig(
context_feature_fields=['query_id', 'query_content']))
self.assertEqual(
standard_artifacts.Examples.TYPE_NAME,
big_query_to_elwc_example_gen.outputs['examples'].type_name)
def testConstructWithInputConfig(self):
big_query_to_elwc_example_gen = component.BigQueryToElwcExampleGen(
elwc_config=elwc_config_pb2.ElwcConfig(
context_feature_fields=['query_id', 'query_content']),
input_config=example_gen_pb2.Input(splits=[
example_gen_pb2.Input.Split(name='train', pattern='query1'),
example_gen_pb2.Input.Split(name='eval', pattern='query2'),
example_gen_pb2.Input.Split(name='test', pattern='query3')
]))
self.assertEqual(
standard_artifacts.Examples.TYPE_NAME,
big_query_to_elwc_example_gen.outputs['examples'].type_name)
def testConstructWithOutputConfig(self):
big_query_to_elwc_example_gen = component.BigQueryToElwcExampleGen(
query='query',
elwc_config=elwc_config_pb2.ElwcConfig(
context_feature_fields=['query_id', 'query_content']),
output_config=example_gen_pb2.
Output(split_config=example_gen_pb2.SplitConfig(splits=[
example_gen_pb2.SplitConfig.Split(name='train',
hash_buckets=2),
example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1),
example_gen_pb2.SplitConfig.Split(name='test', hash_buckets=1)
])))
self.assertEqual(
standard_artifacts.Examples.TYPE_NAME,
big_query_to_elwc_example_gen.outputs['examples'].type_name)