feat: Delta file support - GCS to BQ template enhancement

python/dataproc_templates/gcs/README.md

@@ -11,7 +11,11 @@ It uses the [Spark BigQuery connector](https://cloud.google.com/dataproc-serverl
 * `gcs.bigquery.input.format`: Input file format (one of: avro,parquet,csv,json,delta)
 * `gcs.bigquery.temp.bucket.name`: Temporary bucket for the Spark BigQuery connector
 * `gcs.bigquery.output.mode`: Output write mode (one of: append,overwrite,ignore,errorifexists)(Defaults to append)
+
 #### Optional Arguments
+* `gcs.bigquery.direct.write.method`: pass method as 'direct' to use [BigQuery Storage Write API](https://docs.cloud.google.com/bigquery/docs/write-api)
+* `gcs.to.bigquery.temp.view.name`: Temp view name for creating a spark sql view on source data. This name has to match with the table name that will be used in the SQLquery
+* `gcs.to.bigquery.sql.query`: SQL query for data transformation. This must use the temp view name as the table to query from.
 * `gcs.bigquery.input.chartoescapequoteescaping`: Sets a single character used for escaping the escape for the quote character. The default value is escape character when escape and quote characters are different, \0 otherwise
 * `gcs.bigquery.input.columnnameofcorruptrecord`: Allows renaming the new field having malformed string created by PERMISSIVE mode
 * `gcs.bigquery.input.comment`: Sets a single character used for skipping lines beginning with this character. By default it is disabled
@@ -51,6 +55,9 @@ usage: main.py [-h]
                --gcs.bigquery.output.dataset GCS.BIGQUERY.OUTPUT.DATASET
                --gcs.bigquery.output.table GCS.BIGQUERY.OUTPUT.TABLE
                --gcs.bigquery.input.format {avro,parquet,csv,json,delta}
+               [--gcs.bigquery.direct.write.method GCS.BIGQUERY.DIRECT.WRITE.METHOD]
+               [--gcs.to.bigquery.temp.view.name GCS.TO.BIGQUERY.TEMP.VIEW.NAME]
+               [--gcs.to.bigquery.sql.query GCS.TO.BIGQUERY.SQL.QUERY]
                [--gcs.bigquery.input.chartoescapequoteescaping GCS.BIGQUERY.INPUT.CHARTOESCAPEQUOTEESCAPING]
                [--gcs.bigquery.input.columnnameofcorruptrecord GCS.BIGQUERY.INPUT.COLUMNNAMEOFCORRUPTRECORD]
                [--gcs.bigquery.input.comment GCS.BIGQUERY.INPUT.COMMENT]
@@ -92,6 +99,10 @@ options:
                         BigQuery output table name
   --gcs.bigquery.input.format {avro,parquet,csv,json,delta}
                         Input file format (one of: avro,parquet,csv,json,delta)
+  --gcs.to.bigquery.temp.view.name GCS.TO.BIGQUERY.TEMP.VIEW.NAME
+                        Temp view name for creating a spark sql view on source data. This name has to match with the table name that will be used in the SQL query
+  --gcs.to.bigquery.sql.query GCS.TO.BIGQUERY.SQL.QUERY
+                        SQL query for data transformation. This must use the temp view name as the table to query from.
   --gcs.bigquery.input.chartoescapequoteescaping GCS.BIGQUERY.INPUT.CHARTOESCAPEQUOTEESCAPING
                         Sets a single character used for escaping the escape for the quote character. The default value is escape character when escape and quote characters are
                         different, \0 otherwise

python/dataproc_templates/gcs/gcs_to_bigquery.py

@@ -15,8 +15,8 @@
 from typing import Dict, Sequence, Optional, Any
 from logging import Logger
 import argparse
+import sys
 import pprint
-
 from pyspark.sql import SparkSession
 
 from dataproc_templates import BaseTemplate
@@ -92,10 +92,35 @@ def parse_args(args: Optional[Sequence[str]] = None) -> Dict[str, Any]:
                 constants.OUTPUT_MODE_ERRORIFEXISTS
             ]
         )
+        parser.add_argument(
+            f'--{constants.GCS_TO_BQ_TEMP_VIEW_NAME}',
+            dest=constants.GCS_TO_BQ_TEMP_VIEW_NAME,
+            required=False,
+            default="",
+            help='Temp view name for creating a spark sql view on source data. This name has to match with the table name that will be used in the SQL query'
+        )
+        parser.add_argument(
+            f'--{constants.GCS_TO_BQ_SQL_QUERY}',
+            dest=constants.GCS_TO_BQ_SQL_QUERY,
+            required=False,
+            default="",
+            help='SQL query for data transformation. This must use the temp view name as the table to query from.'
+        )
+
+        parser.add_argument(
+            f'--{constants.GCS_TO_BQ_DIRECT_WRITE_METHOD}',
+            dest=constants.GCS_TO_BQ_DIRECT_WRITE_METHOD,
+            required=False,
+            default="",
+            help='Pass Direct write method to use Storage Write API'
+        )
 
         known_args: argparse.Namespace
         known_args, _ = parser.parse_known_args(args)
 
+        if getattr(known_args, constants.GCS_TO_BQ_SQL_QUERY) and not getattr(known_args, constants.GCS_TO_BQ_TEMP_VIEW_NAME):
+            sys.exit('ArgumentParser Error: Temp view name cannot be null if you want to do data transformations with query')
+
         return vars(known_args)
 
     def run(self, spark: SparkSession, args: Dict[str, Any]) -> None:
@@ -109,6 +134,9 @@ def run(self, spark: SparkSession, args: Dict[str, Any]) -> None:
         big_query_table: str = args[constants.GCS_BQ_OUTPUT_TABLE]
         bq_temp_bucket: str = args[constants.GCS_BQ_LD_TEMP_BUCKET_NAME]
         output_mode: str = args[constants.GCS_BQ_OUTPUT_MODE]
+        bq_temp_view: str = args[constants.GCS_TO_BQ_TEMP_VIEW_NAME]
+        sql_query: str = args[constants.GCS_TO_BQ_SQL_QUERY]
+        write_method: str = args[constants.GCS_TO_BQ_DIRECT_WRITE_METHOD]
 
         logger.info(
             "Starting Cloud Storage to BigQuery Spark job with parameters:\n"
@@ -120,10 +148,76 @@ def run(self, spark: SparkSession, args: Dict[str, Any]) -> None:
             spark, args, input_location, input_format, "gcs.bigquery.input."
         )
 
-        # Write
-        input_data.write \
+        if sql_query:
+            print("------- RUNNING OVERRIDE SQL QUERY -------")
+            print("sql_query : ", sql_query)
+            input_data.createOrReplaceTempView(bq_temp_view)
+            # Execute SQL
+            transformed_data = spark.sql(sql_query)
+        else:
+            transformed_data = input_data
+
+
+        def datatype_mapping(source_schema_map, df):
+            """
+            Generates a SQL SELECT clause with column and datatype transformations
+            based on the source schema map. This is used to build a SQL query.
+            """  
+            transformed_columns = []
+            # Loop through each (column_name, datatype) pair from the source DataFrame
+            for column_name, source_type in source_schema_map.items():
+                source_type = source_type.lower() # 
+                # --- Apply Datatype Mappings based on user logic ---
+                if "map" in source_type:
+                    transformed_column_expression = f"to_json({column_name}) AS {column_name}"
+                # Handle ARRAY: Convert to a JSON string
+                elif "array" in source_type:
+                    # transformed_column_expression = f"to_json({column_name}) AS {column_name}"
+                    transformed_column_expression = f"COALESCE({column_name}, ARRAY()) AS {column_name}"
+                else:
+                    # Default case: Keep the column as is (e.g., string, int, boolean)
+                    transformed_column_expression = column_name
+
+                transformed_columns.append(transformed_column_expression)
+
+            # Build the final SQL query: "SELECT col1, CAST(col2...), to_json(col3...)"
+            query = "SELECT " + ",\n       ".join(transformed_columns) + "\n  FROM TEMP_DF"
+            print("------- GENERATED SQL QUERY -------")
+            print(query)
+
+            # Register the original DataFrame as a temporary table to query it
+            df.createOrReplaceTempView("TEMP_DF")
+            # Run the SQL query to get the new, transformed DataFrame
+            override_df = spark.sql(query)
+            return override_df
+
+
+        # --- Call the function ---
+        # Get the schema from the source DataFrame as a dictionary
+        source_schema_map = dict(transformed_data.dtypes)
+
+        # Run the mapping function to get the transformed DataFrame
+        if input_format == "delta":
+            output_data = datatype_mapping(source_schema_map, transformed_data)
+        else:
+            output_data = transformed_data
+
+        # output_data.printSchema()
+        # output_data.show(3, False)
+
+        if write_method == 'direct':
+            # BQ Write direct method
+            output_data.write \
             .format(constants.FORMAT_BIGQUERY) \
             .option(constants.TABLE, big_query_dataset + "." + big_query_table) \
-            .option(constants.GCS_BQ_TEMP_BUCKET, bq_temp_bucket) \
+            .option(constants.BQ_WRITE_METHOD, constants.BQ_DIRECT_WRITE_METHOD) \
             .mode(output_mode) \
             .save()
+        else:
+            # BQ Write Indirect method
+            output_data.write \
+                .format(constants.FORMAT_BIGQUERY) \
+                .option(constants.TABLE, big_query_dataset + "." + big_query_table) \
+                .option(constants.GCS_BQ_TEMP_BUCKET, bq_temp_bucket) \
+                .mode(output_mode) \
+                .save()

python/dataproc_templates/util/template_constants.py

@@ -174,6 +174,8 @@
 BQ_ALLOW_FIELD_RELAXATION="allowFieldRelaxation"
 BQ_BIGNUMERIC_DEFAULT_PRECISION="bigNumericDefaultPrecision"
 BQ_BIGNUMERIC_DEFAULT_SCALE="bigNumericDefaultScale"
+BQ_WRITE_METHOD="writeMethod"
+BQ_DIRECT_WRITE_METHOD="direct"
 
 OPTION_DEFAULT = "default"
 OPTION_HELP = "help"
@@ -653,6 +655,10 @@ def get_es_spark_connector_input_options(prefix):
 GCS_BQ_OUTPUT_MODE = "gcs.bigquery.output.mode"
 GCS_BQ_TEMP_BUCKET = "temporaryGcsBucket"
 GCS_BQ_LD_TEMP_BUCKET_NAME = "gcs.bigquery.temp.bucket.name"
+GCS_TO_BQ_TEMP_VIEW_NAME = "gcs.bigquery.temp.view.name"
+GCS_TO_BQ_SQL_QUERY = "gcs.bigquery.sql.query"
+GCS_TO_BQ_DIRECT_WRITE_METHOD = "gcs.bigquery.direct.write.method"
+
 
 # GCS to JDBC
 GCS_JDBC_INPUT_LOCATION = "gcs.jdbc.input.location"