[Feature/Addition] Adding Avro support to Google BQ/GCS Connectors

[matthewkrausse]

Wanted to start a discussion here...

According to Big Query documentation, Avro is the preferred format for loading into Big Query from GCS. It seems like we could see improvements to load times using this format instead of .csv files.

A toavro() method exists in the PETL library so we can just expose that as well.

Looking at the code, it seems like we could probably add a flag in the bq.copy() method as to the preferred file format as to keep this from being a breaking change. Want to hear from folks if there are any thoughts and if it is a feasible project?

[matthewkrausse]

Just putting some stuff in here we would need to look at for making this change.

def copy(
        self,
        tbl: Table,
        table_name: str,
        if_exists: str = "fail",
        max_errors: int = 0,
        tmp_gcs_bucket: Optional[str] = None,
        gcs_client: Optional[GoogleCloudStorage] = None,
        job_config: Optional[LoadJobConfig] = None,
        template_table: Optional[str] = None,
        ignoreheader: int = 1,
        nullas: Optional[str] = None,
        allow_quoted_newlines: bool = True,
        allow_jagged_rows: bool = True,
        quote: Optional[str] = None,
        schema: Optional[List[dict]] = None,
        max_timeout: int = 21600,
        convert_dict_list_columns_to_json: bool = True,
        **load_kwargs,
    ):
        """
        Copy a :ref:`parsons-table` into Google BigQuery via Google Cloud Storage.

        `Args:`
            tbl: obj
                The Parsons Table to copy into BigQuery.
            table_name: str
                The table name to load the data into. Will be used to generate load schema
                if no custom schema or template table are supplied and if_exists is
                set to "truncate" or "append".
            if_exists: str
                If the table already exists, either ``fail``, ``append``, ``drop``
                or ``truncate`` the table.
            max_errors: int
                The maximum number of rows that can error and be skipped before
                the job fails.
            tmp_gcs_bucket: str
                The name of the Google Cloud Storage bucket to use to stage the data to load
                into BigQuery. Required if `GCS_TEMP_BUCKET` is not specified or set on
                the class instance.
            gcs_client: object
                The GoogleCloudStorage Connector to use for loading data into Google Cloud Storage.
            job_config: object
                A LoadJobConfig object to provide to the underlying call to load_table_from_uri
                on the BigQuery client. The function will create its own if not provided.
            template_table: str
                Table name to be used as the load schema. Load operation wil use the same
                columns and data types as the template table.
            max_timeout: int
                The maximum number of seconds to wait for a request before the job fails.
            convert_dict_list_columns_to_json: bool
                If set to True, will convert any dict or list columns (which cannot by default be successfully loaded to BigQuery to JSON strings)
            **load_kwargs: kwargs
                Arguments to pass to the underlying load_table_from_uri call on the BigQuery
                client.
        """
        data_type = "csv"
        tmp_gcs_bucket = (
            tmp_gcs_bucket
            or self.tmp_gcs_bucket
            or check_env.check("GCS_TEMP_BUCKET", tmp_gcs_bucket)
        )
        if not tmp_gcs_bucket:
            raise ValueError(
                "Must set GCS_TEMP_BUCKET environment variable or pass in tmp_gcs_bucket parameter"
            )

        self._validate_copy_inputs(if_exists=if_exists, data_type=data_type)

        # If our source table is loaded from CSV with no transformations
        # The original source file will be directly loaded to GCS
        # We may need to pass along a custom delimiter to BigQuery
        # Otherwise we use the default comma
        if isinstance(tbl.table, petl.io.csv_py3.CSVView):
            csv_delimiter = tbl.table.csvargs.get("delimiter", ",")
        else:
            csv_delimiter = ","

        if convert_dict_list_columns_to_json:
            # Convert dict columns to JSON strings
            for field in tbl.get_columns_type_stats():
                if "dict" in field["type"] or "list" in field["type"]:
                    new_petl = tbl.table.addfield(
                        field["name"] + "_replace",
                        lambda row: json.dumps(row[field["name"]]),
                    )
                    new_tbl = Table(new_petl)
                    new_tbl.remove_column(field["name"])
                    new_tbl.rename_column(field["name"] + "_replace", field["name"])
                    new_tbl.materialize()
                    tbl = new_tbl

        job_config = self._process_job_config(
            job_config=job_config,
            destination_table_name=table_name,
            if_exists=if_exists,
            max_errors=max_errors,
            data_type=data_type,
            template_table=template_table,
            parsons_table=tbl,
            ignoreheader=ignoreheader,
            nullas=nullas,
            allow_quoted_newlines=allow_quoted_newlines,
            allow_jagged_rows=allow_jagged_rows,
            quote=quote,
            custom_schema=schema,
            csv_delimiter=csv_delimiter,
        )

        # Reorder schema to match table to ensure compatibility
        schema = []
        for column in tbl.columns:
            try:
                schema_row = [i for i in job_config.schema if i.name.lower() == column.lower()][0]
            except IndexError:
                raise IndexError(f"Column found in Table that was not found in schema: {column}")
            schema.append(schema_row)
        job_config.schema = schema

        gcs_client = gcs_client or GoogleCloudStorage(app_creds=self.app_creds)
        temp_blob_name = f"{uuid.uuid4()}.{data_type}"
        temp_blob_uri = gcs_client.upload_table(tbl, tmp_gcs_bucket, temp_blob_name)

        # load CSV from Cloud Storage into BigQuery
        try:
            self._load_table_from_uri(
                source_uris=temp_blob_uri,
                destination=self.get_table_ref(table_name=table_name),
                job_config=job_config,
                max_timeout=max_timeout,
                **load_kwargs,
            )
        finally:
            gcs_client.delete_blob(tmp_gcs_bucket, temp_blob_name)

   def _validate_copy_inputs(self, if_exists: str, data_type: str):
        if if_exists not in ["fail", "truncate", "append", "drop"]:
            raise ValueError(
                f"Unexpected value for if_exists: {if_exists}, must be one of "
                '"append", "drop", "truncate", or "fail"'
            )
        if data_type not in ["csv", "json"]:
            raise ValueError(f"Only supports csv or json files [data_type = {data_type}]")

def _process_job_config(
        self,
        destination_table_name: str,
        if_exists: str,
        max_errors: int,
        data_type: str,
        csv_delimiter: Optional[str] = ",",
        ignoreheader: Optional[int] = 1,
        nullas: Optional[str] = None,
        allow_quoted_newlines: Optional[bool] = None,
        allow_jagged_rows: Optional[bool] = None,
        quote: Optional[str] = None,
        job_config: Optional[LoadJobConfig] = None,
        custom_schema: Optional[list] = None,
        template_table: Optional[str] = None,
        parsons_table: Optional[Table] = None,
    ) -> LoadJobConfig:
        """
        Internal function to neatly process a user-supplied job configuration object.
        As a convention, if both the job_config and keyword arguments specify a value,
        we defer to the job_config.

        `Args`:
            job_config: `LoadJobConfig`
                Optionally supplied GCS `LoadJobConfig` object

        `Returns`:
            A `LoadJobConfig` object
        """

        if not job_config:
            job_config = bigquery.LoadJobConfig()

        job_config.schema = self._get_job_config_schema(
            job_config=job_config,
            destination_table_name=destination_table_name,
            if_exists=if_exists,
            parsons_table=parsons_table,
            custom_schema=custom_schema,
            template_table=template_table,
        )
        if not job_config.schema:
            job_config.autodetect = True

        if not job_config.create_disposition:
            job_config.create_disposition = bigquery.CreateDisposition.CREATE_IF_NEEDED

        if not job_config.max_bad_records:
            job_config.max_bad_records = max_errors

        if not job_config.skip_leading_rows and data_type == "csv":
            job_config.skip_leading_rows = ignoreheader

        if not job_config.source_format:
            job_config.source_format = (
                bigquery.SourceFormat.CSV
                if data_type == "csv"
                else bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            )

        if not job_config.field_delimiter:
            if data_type == "csv":
                job_config.field_delimiter = csv_delimiter
            if nullas:
                job_config.null_marker = nullas

        destination_table_exists = self.table_exists(destination_table_name)
        if not job_config.write_disposition:
            if if_exists == "append":
                job_config.write_disposition = bigquery.WriteDisposition.WRITE_APPEND
            elif if_exists == "truncate":
                job_config.write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
            elif destination_table_exists and if_exists == "fail":
                raise Exception("Table already exists.")
            elif if_exists == "drop" and destination_table_exists:
                self.delete_table(destination_table_name)
                job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY
            else:
                job_config.write_disposition = bigquery.WriteDisposition.WRITE_EMPTY

        if not job_config.allow_quoted_newlines and allow_quoted_newlines is not None:
            job_config.allow_quoted_newlines = allow_quoted_newlines

        if data_type == "csv" and allow_jagged_rows is not None:
            job_config.allow_jagged_rows = allow_jagged_rows

        if not job_config.quote_character and quote is not None:
            job_config.quote_character = quote

        return job_config