From c7cdd47d35f1d44bf447dd2752ac2ec791b6964f Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Tue, 6 Jun 2023 17:23:24 +0530 Subject: [PATCH 01/12] added check for pub sub --- weather_mv/loader_pipeline/pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/weather_mv/loader_pipeline/pipeline.py b/weather_mv/loader_pipeline/pipeline.py index e8d02cc..0488d26 100644 --- a/weather_mv/loader_pipeline/pipeline.py +++ b/weather_mv/loader_pipeline/pipeline.py @@ -47,7 +47,7 @@ def pattern_to_uris(match_pattern: str, is_zarr: bool = False) -> t.Iterable[str def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None: all_uris = list(pattern_to_uris(known_args.uris, known_args.zarr)) - if not all_uris: + if not all_uris and (known_args.topic or known_args.subscription): raise FileNotFoundError(f"File pattern '{known_args.uris}' matched no objects") # First URI is useful to get an example data shard. It also can be a Zarr path. From 9aa71aa72a326e4ffcfcfb367f892d601b0b5441 Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Tue, 6 Jun 2023 18:14:12 +0530 Subject: [PATCH 02/12] added create table stage to bq pipeline --- weather_mv/loader_pipeline/bq.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index 10f5aeb..364d002 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -148,8 +148,10 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis region=pipeline_options_dict.get('region')) logger.info('Region validation completed successfully.') - def __post_init__(self): - """Initializes Sink by creating a BigQuery table based on user input.""" + def create_bq_table(self, uri: str) -> str: + if self.table: + return uri + if self.zarr: self.xarray_open_dataset_kwargs = self.zarr_kwargs with open_dataset(self.first_uri, self.xarray_open_dataset_kwargs, @@ -176,6 +178,7 @@ def __post_init__(self): try: table = bigquery.Table(self.output_table, schema=table_schema) self.table = bigquery.Client().create_table(table, exists_ok=True) + return uri except Exception as e: logger.error(f'Unable to create table in BigQuery: {e}') raise @@ -243,6 +246,7 @@ def expand(self, paths): """Extract rows of variables from data paths into a BigQuery table.""" extracted_rows = ( paths + | 'CreateTable' >> beam.Map(self.create_bq_table) | 'PrepareCoordinates' >> beam.FlatMap(self.prepare_coordinates) | beam.Reshuffle() | 'ExtractRows' >> beam.FlatMapTuple(self.extract_rows) From 5ce00028a16591cb604ea6dd967724e63bee030c Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Tue, 6 Jun 2023 18:16:28 +0530 Subject: [PATCH 03/12] return uri from dry run --- weather_mv/loader_pipeline/bq.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index 364d002..2dab829 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -151,7 +151,7 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis def create_bq_table(self, uri: str) -> str: if self.table: return uri - + if self.zarr: self.xarray_open_dataset_kwargs = self.zarr_kwargs with open_dataset(self.first_uri, self.xarray_open_dataset_kwargs, @@ -172,7 +172,7 @@ def create_bq_table(self, uri: str) -> str: if self.dry_run: logger.debug('Created the BigQuery table with schema...') logger.debug(f'\n{pformat(table_schema)}') - return + return uri # Create the table in BigQuery try: From 36237d44a1c64cd879782e06843c353b5a024bd7 Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Wed, 7 Jun 2023 15:36:39 +0530 Subject: [PATCH 04/12] init write to big query from output table var --- weather_mv/loader_pipeline/bq.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index 2dab829..c58adae 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -17,6 +17,7 @@ import json import logging import os +import re import typing as t from pprint import pformat @@ -130,6 +131,12 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis pipeline_options = PipelineOptions(pipeline_args) pipeline_options_dict = pipeline_options.get_all_options() + if known_args.output_table: + # checking if the output table is in format (..). + pattern = r'^[\w-]+\.[\w-]+\.[\w-]+$' + if not bool(re.match(pattern, known_args.output_table)): + raise RuntimeError("Output table is not in correct format.") + if known_args.area: assert len(known_args.area) == 4, 'Must specify exactly 4 lat/long values for area: N, W, S, E boundaries.' @@ -148,6 +155,13 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis region=pipeline_options_dict.get('region')) logger.info('Region validation completed successfully.') + def __post_init__(self): + self.table = None + project, dataset, table_name = self.output_table.split('.') + self.project = project + self.dataset = dataset + self.table_name = table_name + def create_bq_table(self, uri: str) -> str: if self.table: return uri @@ -256,9 +270,11 @@ def expand(self, paths): ( extracted_rows | 'WriteToBigQuery' >> WriteToBigQuery( - project=self.table.project, - dataset=self.table.dataset_id, - table=self.table.table_id, + project=self.project, + dataset=self.dataset, + table=self.table_name, + method="STREAMING_INSERTS", + batch_size=1000, write_disposition=BigQueryDisposition.WRITE_APPEND, create_disposition=BigQueryDisposition.CREATE_NEVER) ) From bc753108a64ea65c1d57ce3d7de1c7c6787ae85b Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Wed, 7 Jun 2023 15:37:24 +0530 Subject: [PATCH 05/12] remove method and batch_size --- weather_mv/loader_pipeline/bq.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index c58adae..3843921 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -273,8 +273,6 @@ def expand(self, paths): project=self.project, dataset=self.dataset, table=self.table_name, - method="STREAMING_INSERTS", - batch_size=1000, write_disposition=BigQueryDisposition.WRITE_APPEND, create_disposition=BigQueryDisposition.CREATE_NEVER) ) From 17a88b73ec89463135e0063d4f49bc7d6025b2fc Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Thu, 8 Jun 2023 11:10:59 +0530 Subject: [PATCH 06/12] removed first uri from bq --- weather_mv/loader_pipeline/bq.py | 3 ++- weather_mv/loader_pipeline/pipeline.py | 8 ++++++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index 3843921..5dd7518 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -168,7 +168,7 @@ def create_bq_table(self, uri: str) -> str: if self.zarr: self.xarray_open_dataset_kwargs = self.zarr_kwargs - with open_dataset(self.first_uri, self.xarray_open_dataset_kwargs, + with open_dataset(uri, self.xarray_open_dataset_kwargs, self.disable_grib_schema_normalization, self.tif_metadata_for_datetime, is_zarr=self.zarr) as open_ds: # Define table from user input @@ -258,6 +258,7 @@ def extract_rows(self, uri: str, coordinates: t.List[t.Dict]) -> t.Iterator[t.Di def expand(self, paths): """Extract rows of variables from data paths into a BigQuery table.""" + print(f"paths {paths}") extracted_rows = ( paths | 'CreateTable' >> beam.Map(self.create_bq_table) diff --git a/weather_mv/loader_pipeline/pipeline.py b/weather_mv/loader_pipeline/pipeline.py index 0488d26..dac1a8c 100644 --- a/weather_mv/loader_pipeline/pipeline.py +++ b/weather_mv/loader_pipeline/pipeline.py @@ -47,11 +47,15 @@ def pattern_to_uris(match_pattern: str, is_zarr: bool = False) -> t.Iterable[str def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None: all_uris = list(pattern_to_uris(known_args.uris, known_args.zarr)) - if not all_uris and (known_args.topic or known_args.subscription): + if not all_uris and "topic" not in known_args: raise FileNotFoundError(f"File pattern '{known_args.uris}' matched no objects") # First URI is useful to get an example data shard. It also can be a Zarr path. - known_args.first_uri = next(iter(all_uris)) + if all_uris: + known_args.first_uri = next(iter(all_uris)) + else: + # If it's a streaming pipeline, it will allow first_uri to be empty. + known_args.first_uri = "" with beam.Pipeline(argv=pipeline_args) as p: if known_args.topic or known_args.subscription: From 5684b9ac3eb8b42657c3e7c4e131009c13cddcce Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Thu, 8 Jun 2023 11:27:02 +0530 Subject: [PATCH 07/12] copy zarr kwargs to open dataset kwargs in bq post_init --- weather_mv/loader_pipeline/bq.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index 5dd7518..ec528a5 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -161,13 +161,13 @@ def __post_init__(self): self.project = project self.dataset = dataset self.table_name = table_name + if self.zarr: + self.xarray_open_dataset_kwargs = self.zarr_kwargs def create_bq_table(self, uri: str) -> str: if self.table: return uri - if self.zarr: - self.xarray_open_dataset_kwargs = self.zarr_kwargs with open_dataset(uri, self.xarray_open_dataset_kwargs, self.disable_grib_schema_normalization, self.tif_metadata_for_datetime, is_zarr=self.zarr) as open_ds: From 63a25aa1bea56e817b81f0696562ab5d39218688 Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Thu, 8 Jun 2023 12:26:17 +0530 Subject: [PATCH 08/12] first_uri to None is case all_uris is empty --- weather_mv/loader_pipeline/pipeline.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/weather_mv/loader_pipeline/pipeline.py b/weather_mv/loader_pipeline/pipeline.py index dac1a8c..4f34f79 100644 --- a/weather_mv/loader_pipeline/pipeline.py +++ b/weather_mv/loader_pipeline/pipeline.py @@ -55,7 +55,7 @@ def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None known_args.first_uri = next(iter(all_uris)) else: # If it's a streaming pipeline, it will allow first_uri to be empty. - known_args.first_uri = "" + known_args.first_uri = None with beam.Pipeline(argv=pipeline_args) as p: if known_args.topic or known_args.subscription: From cb6b4737abc40e1599efdc31b9d71a8ff6649dea Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Thu, 8 Jun 2023 12:34:15 +0530 Subject: [PATCH 09/12] minor fixes --- weather_mv/loader_pipeline/bq.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index ec528a5..bc5fec5 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -133,9 +133,9 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis if known_args.output_table: # checking if the output table is in format (..
). - pattern = r'^[\w-]+\.[\w-]+\.[\w-]+$' - if not bool(re.match(pattern, known_args.output_table)): - raise RuntimeError("Output table is not in correct format.") + output_table_pattern = r'^[\w-]+\.[\w-]+\.[\w-]+$' + if not bool(re.match(output_table_pattern, known_args.output_table)): + raise RuntimeError("output_table is not in correct format.") if known_args.area: assert len(known_args.area) == 4, 'Must specify exactly 4 lat/long values for area: N, W, S, E boundaries.' @@ -156,16 +156,20 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis logger.info('Region validation completed successfully.') def __post_init__(self): - self.table = None - project, dataset, table_name = self.output_table.split('.') + """Initializes BigQuery table based on user input.""" + project, dataset_id, table_id = self.output_table.split('.') self.project = project - self.dataset = dataset - self.table_name = table_name + self.dataset_id = dataset_id + self.table_id = table_id + self.table = None + if self.zarr: self.xarray_open_dataset_kwargs = self.zarr_kwargs def create_bq_table(self, uri: str) -> str: + """Create a big query table for the first uri. After table is created, subsequent uris are returned.""" if self.table: + # Skip table creation. return uri with open_dataset(uri, self.xarray_open_dataset_kwargs, @@ -258,7 +262,6 @@ def extract_rows(self, uri: str, coordinates: t.List[t.Dict]) -> t.Iterator[t.Di def expand(self, paths): """Extract rows of variables from data paths into a BigQuery table.""" - print(f"paths {paths}") extracted_rows = ( paths | 'CreateTable' >> beam.Map(self.create_bq_table) @@ -272,8 +275,8 @@ def expand(self, paths): extracted_rows | 'WriteToBigQuery' >> WriteToBigQuery( project=self.project, - dataset=self.dataset, - table=self.table_name, + dataset=self.dataset_id, + table=self.table_id, write_disposition=BigQueryDisposition.WRITE_APPEND, create_disposition=BigQueryDisposition.CREATE_NEVER) ) From e3ca3b1dc634e46e12f61e0db6ace37d28f111e9 Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Thu, 8 Jun 2023 14:28:30 +0530 Subject: [PATCH 10/12] minor nits --- weather_mv/loader_pipeline/bq.py | 4 ++-- weather_mv/loader_pipeline/pipeline.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index bc5fec5..f5d338d 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -135,7 +135,7 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis # checking if the output table is in format (..
). output_table_pattern = r'^[\w-]+\.[\w-]+\.[\w-]+$' if not bool(re.match(output_table_pattern, known_args.output_table)): - raise RuntimeError("output_table is not in correct format.") + raise RuntimeError("output_table is not in correct format (..). ") if known_args.area: assert len(known_args.area) == 4, 'Must specify exactly 4 lat/long values for area: N, W, S, E boundaries.' @@ -168,8 +168,8 @@ def __post_init__(self): def create_bq_table(self, uri: str) -> str: """Create a big query table for the first uri. After table is created, subsequent uris are returned.""" + # Skip table creation. if self.table: - # Skip table creation. return uri with open_dataset(uri, self.xarray_open_dataset_kwargs, diff --git a/weather_mv/loader_pipeline/pipeline.py b/weather_mv/loader_pipeline/pipeline.py index 4f34f79..85d9230 100644 --- a/weather_mv/loader_pipeline/pipeline.py +++ b/weather_mv/loader_pipeline/pipeline.py @@ -47,7 +47,7 @@ def pattern_to_uris(match_pattern: str, is_zarr: bool = False) -> t.Iterable[str def pipeline(known_args: argparse.Namespace, pipeline_args: t.List[str]) -> None: all_uris = list(pattern_to_uris(known_args.uris, known_args.zarr)) - if not all_uris and "topic" not in known_args: + if not all_uris and not known_args.topic: raise FileNotFoundError(f"File pattern '{known_args.uris}' matched no objects") # First URI is useful to get an example data shard. It also can be a Zarr path. From 324b5389262ce52490b9466c01b9240c50a0ab5d Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Fri, 9 Jun 2023 12:51:28 +0530 Subject: [PATCH 11/12] minor changes --- weather_mv/loader_pipeline/bq.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index f5d338d..e53a992 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -157,10 +157,7 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis def __post_init__(self): """Initializes BigQuery table based on user input.""" - project, dataset_id, table_id = self.output_table.split('.') - self.project = project - self.dataset_id = dataset_id - self.table_id = table_id + self.project, self.dataset_id, self.table_id = self.output_table.split('.') self.table = None if self.zarr: From 820c34c1e4828a09fd1d8ba676a8e51caeadd488 Mon Sep 17 00:00:00 2001 From: Aniket Singh Rawat Date: Tue, 13 Jun 2023 17:45:08 +0530 Subject: [PATCH 12/12] lint issue --- weather_mv/loader_pipeline/bq.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/weather_mv/loader_pipeline/bq.py b/weather_mv/loader_pipeline/bq.py index e53a992..6328a15 100644 --- a/weather_mv/loader_pipeline/bq.py +++ b/weather_mv/loader_pipeline/bq.py @@ -158,7 +158,7 @@ def validate_arguments(cls, known_args: argparse.Namespace, pipeline_args: t.Lis def __post_init__(self): """Initializes BigQuery table based on user input.""" self.project, self.dataset_id, self.table_id = self.output_table.split('.') - self.table = None + self.table = None if self.zarr: self.xarray_open_dataset_kwargs = self.zarr_kwargs