marketing-analytics/predicting/future-customer-value-segments/fcvs_pipeline_bq.py

# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import logging
import operator

import apache_beam as beam
from apache_beam import io
from apache_beam import pvalue
from apache_beam.options import pipeline_options
from apache_beam.options import value_provider
from apache_beam.transforms import util

import bigquery_mod as bq_mod
import common as c


class RuntimeOptions(pipeline_options.PipelineOptions):
    """Specifies runtime options for the pipeline.

    Class defining the arguments that can be passed to the pipeline to
    customize the execution.
    """
    @classmethod
    def _add_argparse_args(cls, parser):
        parser.add_value_provider_argument(f'--{c._OPTION_INPUT_BQ_QUERY}')
        parser.add_value_provider_argument(f'--{c._OPTION_INPUT_BQ_PROJECT}')
        parser.add_value_provider_argument(f'--{c._OPTION_TEMP_GCS_LOCATION}')
        parser.add_value_provider_argument(f'--{c._OPTION_OUTPUT_FOLDER}')
        parser.add_value_provider_argument(f'--{c._OPTION_OUTPUT_BQ_PROJECT}')
        parser.add_value_provider_argument(f'--{c._OPTION_OUTPUT_BQ_DATASET}')
        parser.add_value_provider_argument(
            f'--{c._OPTION_CUSTOMER_ID_COLUMN_NAME}')
        parser.add_value_provider_argument(
            f'--{c._OPTION_TRANSACTION_DATE_COLUMN_NAME}')
        parser.add_value_provider_argument(
            f'--{c._OPTION_SALES_COLUMN_NAME}')
        parser.add_value_provider_argument(
            f'--{c._OPTION_EXTRA_DIMENSION_COLUMN_NAME}')
        parser.add_value_provider_argument(f'--{c._OPTION_DATE_PARSING_PATTERN}')
        parser.add_value_provider_argument(
            f'--{c._OPTION_MODEL_TIME_GRANULARITY}',
            default=c.TimeGranularityParams.GRANULARITY_WEEKLY)
        parser.add_value_provider_argument(
            f'--{c._OPTION_FREQUENCY_MODEL_TYPE}', default=c._MODEL_TYPE_MBGNBD)
        parser.add_value_provider_argument(
            f'--{c._OPTION_CALIBRATION_START_DATE}')
        parser.add_value_provider_argument(f'--{c._OPTION_CALIBRATION_END_DATE}')
        parser.add_value_provider_argument(f'--{c._OPTION_COHORT_START_DATE}')
        parser.add_value_provider_argument(f'--{c._OPTION_COHORT_END_DATE}')
        parser.add_value_provider_argument(f'--{c._OPTION_HOLDOUT_END_DATE}')
        parser.add_value_provider_argument(
            f'--{c._OPTION_PREDICTION_PERIOD}', default=52, type=int)
        parser.add_value_provider_argument(
            f'--{c._OPTION_OUTPUT_SEGMENTS}', default=5, type=int)
        parser.add_value_provider_argument(
            f'--{c._OPTION_TRANSACTION_FREQUENCY_THRESHOLD}', default=15,
            type=int)
        parser.add_value_provider_argument(
            f'--{c._OPTION_PENALIZER_COEF}', default=0.0, type=float)
        parser.add_value_provider_argument(
            f'--{c._OPTION_ROUND_NUMBERS}', default="False")


def run(argv=None):
    """Main function.

    Main function containing the Apache Beam pipeline describing how to process
    the input CSV file to generate the LTV predictions.
    """
    parser = argparse.ArgumentParser()
    _, pipeline_args = parser.parse_known_args(argv)
    options = pipeline_options.PipelineOptions(pipeline_args)
    runtime_options = options.view_as(RuntimeOptions)

    with beam.Pipeline(options=options) as pipeline:
        options = (
            pipeline
            | 'Create single element Stream containing options dict' >>
                beam.Create([options.get_all_options()])
            | beam.Map(lambda x: {
                  k: v.get() if isinstance(v, value_provider.ValueProvider)
                             else v
                  for (k, v) in x.items()
              })
            | beam.Map(c.set_extra_options)
        )

        full_elog = (
            pipeline
            | bq_mod.ReadFromBigQuery(
                project=getattr(runtime_options, c._OPTION_INPUT_BQ_PROJECT),
                query=getattr(runtime_options, c._OPTION_INPUT_BQ_QUERY),
                gcs_location=getattr(runtime_options, c._OPTION_TEMP_GCS_LOCATION),
                use_standard_sql=True,
                priority='BATCH'
            )
            | beam.FlatMap(
                c.bq_row_to_list,
                pvalue.AsSingleton(options))  # (customer_id, date_str, date,
                                              #  sales, extra_dimension?)
        )

        full_elog_merged = (
            full_elog
            | beam.Filter(lambda x: x[3] > 0)  # sales > 0
            | beam.Map(lambda x: ((x[0], x[1]), x))  # key: (customer_id, date)
            | 'Group full elog by customer and date' >> beam.GroupByKey()
            | beam.Map(c.merge_full_elog_by_customer_and_date)  # (customer_id,
                                                                #  date_str, date,
                                                                #  sales)
        )

        min_max_dates = (
            full_elog_merged
            | beam.Map(lambda x: x[2])  # date
            | beam.CombineGlobally(c.MinMaxDatesFn())
            | beam.Map(c.min_max_dates_dict)
        )

        limits_dates = (
            min_max_dates
            | beam.FlatMap(c.limit_dates_boundaries, pvalue.AsSingleton(options))
        )

        cohort = (
            full_elog_merged
            | beam.FlatMap(c.filter_customers_in_cohort,
                           pvalue.AsSingleton(limits_dates))
            | 'Distinct Customer IDs in Cohort' >> util.Distinct()
        )

        cohort_count = (
            cohort
            | 'Count cohort entries' >> beam.combiners.Count.Globally()
        )

        cohort_set = (
            cohort
            | beam.Map(lambda x: (x, 1))
        )

        all_customer_ids = (
            full_elog_merged
            | beam.Map(lambda x: x[0])  # key: customer_id
            | 'Distinct all Customer IDs' >> util.Distinct()
        )

        all_customer_ids_count = (
            all_customer_ids
            | 'Count all customers' >> beam.combiners.Count.Globally()
        )

        num_customers = (
            pipeline
            | 'Create single elem Stream I' >> beam.Create([1])
            | beam.FlatMap(c.count_customers,
                           pvalue.AsSingleton(cohort_count),
                           pvalue.AsSingleton(all_customer_ids_count),
                           pvalue.AsSingleton(options))
        )

        cal_hol_elog = (
            full_elog_merged
            | beam.FlatMap(c.filter_cohort_records_in_cal_hol,
                           pvalue.AsDict(cohort_set),
                           pvalue.AsSingleton(limits_dates))
        )

        cal_hol_elog_count = (
            cal_hol_elog
            | 'Count cal hol elog entries' >> beam.combiners.Count.Globally()
        )

        calibration = (
            cal_hol_elog
            | beam.FlatMap(c.filter_records_in_calibration,
                           pvalue.AsSingleton(limits_dates))
        )

        num_txns_total = (
            full_elog_merged
            | beam.FlatMap(c.filter_records_in_cal_hol,
                           pvalue.AsSingleton(limits_dates))
            | 'Count num txns total' >> beam.combiners.Count.Globally()
        )

        num_txns = (
            pipeline
            | 'Create single elem Stream II' >> beam.Create([1])
            | beam.FlatMap(c.count_txns,
                           pvalue.AsSingleton(cal_hol_elog_count),
                           pvalue.AsSingleton(num_txns_total),
                           pvalue.AsSingleton(options))
        )

        calcbs = (
            calibration
            | beam.Map(lambda x: (x[0], x))
            | 'Group calibration elog by customer id' >> beam.GroupByKey()
            | beam.FlatMap(
                c.create_cal_cbs,
                pvalue.AsSingleton(options),
                pvalue.AsSingleton(limits_dates)
            )  # (customer_id, number_of_transactions, average_order_value,
               #  frequency, recency, total_time_observed)
        )

        first_transaction_dates_by_customer = (
            cal_hol_elog
            | beam.Map(lambda x: (x[0], x))  # customer_id
            | 'Group cal hol elog by customer id' >> beam.GroupByKey()
            | beam.Map(lambda x: (x[0], min(map(operator.itemgetter(2), x[1])))
                       )  # item 2 -> date
        )

        cal_hol_elog_repeat = (
            cal_hol_elog
            | beam.FlatMap(c.filter_first_transaction_date_records,
                           pvalue.AsDict(first_transaction_dates_by_customer))
            | beam.FlatMap(
                c.calculate_time_unit_numbers,  # (customer_id, date,
                                                #  time_unit_number)
                pvalue.AsSingleton(options),
                pvalue.AsSingleton(limits_dates))
            | beam.Map(lambda x: (x[2], 1))  # key: time_unit_number
            | 'Group cal hol elog repeat by time unit number' >>
            beam.GroupByKey()
            | beam.Map(lambda x: (x[0], sum(x[1]))
                       )  # (time_unit_number, occurrences)
        )

        repeat_tx = (
            pipeline
            | 'Create single elem Stream III' >> beam.Create([1])
            | beam.FlatMap(c.calculate_cumulative_repeat_transactions,
                           pvalue.AsIter(cal_hol_elog_repeat)
                           )  # (time_unit_number, repeat_transactions,
                              #  repeat_transactions_cumulative)
        )

        model_validation = (
            pipeline
            | 'Create single elem Stream IV' >> beam.Create([1])
            | beam.FlatMap(c.calculate_model_fit_validation,
                           pvalue.AsSingleton(options),
                           pvalue.AsSingleton(limits_dates),
                           pvalue.AsIter(calcbs),
                           pvalue.AsIter(repeat_tx),
                           pvalue.AsSingleton(num_customers),
                           pvalue.AsSingleton(num_txns))
        )

        _ = (
            model_validation
            | beam.Map(c.raise_error_if_invalid_mape)
        )

        _ = (
            model_validation
            | beam.Map(lambda x: x[0])
            | 'Write to validation_params table' >>
                io.WriteToBigQuery(
                    table=c.TableValueProvider(
                        getattr(runtime_options, c._OPTION_OUTPUT_BQ_PROJECT),
                        getattr(runtime_options, c._OPTION_OUTPUT_BQ_DATASET),
                        'validation_params'
                    ),
                    custom_gcs_temp_location=getattr(runtime_options, c._OPTION_TEMP_GCS_LOCATION),
                    validate=False,
                    schema={
                        'fields': [
                            {'name': 'calibration_start_date', 'type': 'STRING'},
                            {'name': 'calibration_end_date', 'type': 'STRING'},
                            {'name': 'cohort_start_date', 'type': 'STRING'},
                            {'name': 'cohort_end_date', 'type': 'STRING'},
                            {'name': 'holdout_end_date', 'type': 'STRING'},
                            {'name': 'model_time_granularity', 'type': 'STRING'},
                            {'name': 'model', 'type': 'RECORD',
                                'fields': [
                                    {'name': 'frequency_model', 'type': 'STRING'},
                                    {'name': 'num_customers_cohort', 'type': 'INTEGER'},
                                    {'name': 'perc_customers_cohort', 'type': 'FLOAT'},
                                    {'name': 'num_transactions_validation', 'type': 'INTEGER'},
                                    {'name': 'perc_transactions_validation', 'type': 'FLOAT'},
                                    {'name': 'validation_mape', 'type': 'STRING'},
                                ]}
                        ]
                    },
                    write_disposition=io.BigQueryDisposition.WRITE_TRUNCATE,
                    create_disposition=io.BigQueryDisposition.CREATE_IF_NEEDED)
        )

        fullcbs_without_extra_dimension = (
            full_elog_merged
            | beam.Map(lambda x: (x[0], x))  # key: customer_id
            | 'Group full merged elog by customer id' >> beam.GroupByKey()
            | beam.FlatMap(
                c.create_fullcbs,
                pvalue.AsSingleton(options),
                pvalue.AsSingleton(min_max_dates)
            )  # (customer_id, number_of_transactions, historical_aov,
               #  frequency, recency, total_time_observed)
        )

        full_elog_if_extra_dimension = (
            full_elog
            | 'Discard records if no extra dimension' >> beam.FlatMap(
                c.discard_if_no_extra_dimension, pvalue.AsSingleton(options))
        )

        extra_dimensions_stats = (
            full_elog_if_extra_dimension
            | beam.Map(lambda x: ((x[0], x[4]), x)
                       )  # key: (customer_id, extra_dimension)
            | 'Group full elog by customer id and extra dimension' >>
                beam.GroupByKey()
            | beam.Map(
                c.create_extra_dimensions_stats
            )  # (customer_id, extra_dimension, dimension_count, tot_sales,
               #  max_dimension_date)
        )

        top_dimension_per_customer = (
            extra_dimensions_stats
            | beam.Map(lambda x: (x[0], x))  # customer_id
            | 'Group extra dimension stats by customer id' >> beam.GroupByKey()
            | beam.Map(
                c.extract_top_extra_dimension
            )  # (customer_id, extra_dimension, dimension_count, tot_sales,
               #  max_dimension_date)
        )

        customer_dimension_map = (
            top_dimension_per_customer
            | beam.Map(
                lambda x: (x[0], x[1]))  # (customer_id, extra_dimension)
        )

        prediction = (
            pipeline
            | 'Create single elem Stream V' >> beam.Create([1])
            | beam.FlatMap(
                c.calculate_prediction,
                pvalue.AsSingleton(options),
                pvalue.AsIter(fullcbs_without_extra_dimension),
                pvalue.AsSingleton(num_customers),
                pvalue.AsSingleton(num_txns)
            )  # [customer_id, p_alive, predicted_purchases, future_aov,
               #  historical_aov, expected_value, frequency, recency,
               #  total_time_observed], prediction_params
        )

        prediction_by_customer_no_segments_no_extra_dimension = (
            prediction
            | beam.FlatMap(lambda x: x[0])  # Extract predictions by customer
        )

        prediction_by_customer_no_segments = (
            prediction_by_customer_no_segments_no_extra_dimension
            | beam.FlatMap(
                c.add_top_extra_dimension_to_fullcbs,
                pvalue.AsSingleton(options),
                pvalue.AsDict(customer_dimension_map)
            )  # [customer_id, p_alive, predicted_purchases, future_aov
               #  historical_aov, expected_value, frequency, recency,
               #  total_time_observed, extra_dimension?]
        )

        _ = (
            prediction
            | beam.Map(lambda x: x[1])  # Extract prediction params
            | 'Write to prediction_params table' >>
                io.WriteToBigQuery(
                    table=c.TableValueProvider(
                        getattr(runtime_options, c._OPTION_OUTPUT_BQ_PROJECT),
                        getattr(runtime_options, c._OPTION_OUTPUT_BQ_DATASET),
                        'prediction_params'
                    ),
                    custom_gcs_temp_location=getattr(runtime_options, c._OPTION_TEMP_GCS_LOCATION),
                    validate=False,
                    schema={
                        'fields': [
                            {'name': 'prediction_period', 'type': 'INTEGER'},
                            {'name': 'prediction_period_unit', 'type': 'STRING'},
                            {'name': 'model_time_granularity', 'type': 'STRING'},
                            {'name': 'customers_modeled', 'type': 'INTEGER'},
                            {'name': 'transactions_observed', 'type': 'INTEGER'},
                            {'name': 'frequency_model', 'type': 'STRING'},
                            {'name': 'bgnbd_model_params', 'type': 'RECORD',
                                'fields': [
                                    {'name': 'a', 'type': 'FLOAT'},
                                    {'name': 'b', 'type': 'FLOAT'},
                                    {'name': 'r', 'type': 'FLOAT'},
                                    {'name': 'alpha', 'type': 'FLOAT'}
                                ]},
                            {'name': 'bgbb_model_params', 'type': 'RECORD',
                                'fields': [
                                    {'name': 'alpha', 'type': 'FLOAT'},
                                    {'name': 'beta', 'type': 'FLOAT'},
                                    {'name': 'gamma', 'type': 'FLOAT'},
                                    {'name': 'delta', 'type': 'FLOAT'}
                                ]},
                            {'name': 'paretonbd_model_params',
                                'type': 'RECORD',
                                'fields': [
                                    {'name': 'r', 'type': 'FLOAT'},
                                    {'name': 's', 'type': 'FLOAT'},
                                    {'name': 'alpha', 'type': 'FLOAT'},
                                    {'name': 'beta', 'type': 'FLOAT'}
                                ]},
                            {'name': 'gamma_gamma_params',
                                'type': 'RECORD',
                                'fields': [
                                    {'name': 'p', 'type': 'FLOAT'},
                                    {'name': 'q', 'type': 'FLOAT'},
                                    {'name': 'v', 'type': 'FLOAT'}
                                ]}
                        ]
                    },
                    write_disposition=io.BigQueryDisposition.WRITE_TRUNCATE,
                    create_disposition=io.BigQueryDisposition.CREATE_IF_NEEDED)
        )

        num_rows = (
            full_elog_merged
            | 'Count num rows in full elog merged' >>
                beam.combiners.Count.Globally()
        )

        segment_predictions_exact = (
            pipeline
            | 'Create single elem Stream VII' >> beam.Create([1])
            | beam.FlatMap(lambda _, rows_count: [
                rows_count <= c._SEGMENT_PREDICTION_THRESHOLD],
                           pvalue.AsSingleton(num_rows))
        )

        sharded_cust_predictions_no_segments_exact, \
            sharded_cust_predictions_no_segments_hash = (
                prediction_by_customer_no_segments
                | beam.FlatMap(
                    c.prediction_sharded,
                    pvalue.AsSingleton(options),
                    pvalue.AsSingleton(segment_predictions_exact)
                )  # [customer_id, p_alive, predicted_purchases, future_aov,
                   #  historical_aov, expected_value, frequency, recency,
                   #  total_time_observed, extra_dimension?]
                | beam.Partition(lambda x, _: 0 if x[1] else 1, 2)
            )

        # BEGIN of "exact" branch
        prediction_by_customer_exact = (
            pipeline
            | 'Create single elem Stream VIII' >> beam.Create([1])
            | beam.FlatMap(c.split_in_ntiles_exact,
                           pvalue.AsSingleton(options),
                           pvalue.AsIter(
                               sharded_cust_predictions_no_segments_exact)
                           )  # [customer_id, p_alive, predicted_purchases,
                              #  future_aov, historical_aov, expected_value,
                              #  frequency, recency, total_time_observed,
                              #  segment, extra_dimension?]
        )
        # END of "exact" branch

        # BEGIN of "hash" branch
        customer_count_by_expected_value = (
            sharded_cust_predictions_no_segments_hash
            | beam.Map(lambda x: (x[0][5], 1))  # (expected_value, 1)
            | 'Group customer predictions by expected value' >>
                beam.GroupByKey()
            | beam.Map(
                lambda x: (x[0], sum(x[1])))  # expected_value, customers_count
        )

        hash_segment_limits = (
            pipeline
            | 'Create single elem Stream IX' >> beam.Create([1])
            | beam.FlatMap(c.expected_values_segment_limits,
                           pvalue.AsSingleton(options),
                           pvalue.AsIter(customer_count_by_expected_value),
                           pvalue.AsSingleton(all_customer_ids_count))
        )

        prediction_by_customer_hash = (
            sharded_cust_predictions_no_segments_hash
            | beam.Map(lambda x: x[0])
            | beam.FlatMap(c.split_in_ntiles_hash,
                           pvalue.AsSingleton(hash_segment_limits)
                           )  # [customer_id, p_alive, predicted_purchases,
                              #  future_aov, historical_aov, expected_value,
                              #  frequency, recency, total_time_observed,
                              #  segment, extra_dimension?]
        )
        # END of "hash" branch

        prediction_by_customer = (
            # only one of these two streams will contains values
            (prediction_by_customer_exact, prediction_by_customer_hash)
            | beam.Flatten()
            | beam.Map(c.clean_nan_and_inf)
        )

        _ = (
            prediction_by_customer
            | beam.FlatMap(lambda x, opts: [x + ['']]
                           if not opts[c._OPTION_EXTRA_DIMENSION_EXISTS] else [x],
                           pvalue.AsSingleton(options))
            | 'prediction_by_customer to Dict' >> beam.Map(c.list_to_dict, [
                    'customer_id', 'p_alive', 'predicted_purchases',
                    'future_aov', 'historical_aov',
                    'expected_value', 'frequency', 'recency',
                    'total_time_observed', 'segment',
                    'extra_dimension'])
            | 'Write to prediction_by_customer table' >>
                io.WriteToBigQuery(
                    table=c.TableValueProvider(
                        getattr(runtime_options, c._OPTION_OUTPUT_BQ_PROJECT),
                        getattr(runtime_options, c._OPTION_OUTPUT_BQ_DATASET),
                        'prediction_by_customer'
                    ),
                    custom_gcs_temp_location=getattr(runtime_options, c._OPTION_TEMP_GCS_LOCATION),
                    validate=False,
                    schema='customer_id:STRING, p_alive:FLOAT64'
                           ', predicted_purchases:FLOAT64'
                           ', future_aov:FLOAT64, historical_aov:FLOAT64'
                           ', expected_value:FLOAT64, frequency:INT64'
                           ', recency:FLOAT64'
                           ', total_time_observed:FLOAT64, segment:INT64'
                           ', extra_dimension:STRING',
                    write_disposition=io.BigQueryDisposition.WRITE_TRUNCATE,
                    create_disposition=io.BigQueryDisposition.CREATE_IF_NEEDED)
        )

        prediction_summary_temp = (
            prediction_by_customer
            | beam.Map(lambda x: (x[9], x))  # key: segment
            | 'Group customer predictions by segment' >> beam.GroupByKey()
            | beam.FlatMap(c.generate_prediction_summary,
                           pvalue.AsSingleton(options)
                           )  # (segment, average_retention_probability,
                              #  average_predicted_customer_value,
                              #  average_predicted_order_value,
                              #  average_predicted_purchases, total_customer_value,
                              #  number_of_customers)
        )

        tot_equity = (
            prediction_summary_temp
            | beam.Map(lambda x: x[5])  # total_customer_value
            | beam.CombineGlobally(sum)
        )

        prediction_summary = (
            prediction_summary_temp
            | beam.FlatMap(
                c.calculate_perc_of_total_customer_value,
                pvalue.AsSingleton(tot_equity),
                pvalue.AsSingleton(options)
                )  # (segment, average_retention_probability,
                   #  average_predicted_customer_value,
                   #  average_predicted_order_value,
                   #  average_predicted_purchases,
                   #  total_customer_value, number_of_customers,
                   #  perc_of_total_customer_value)
        )

        _ = (
            prediction_summary
            | 'prediction_summary to Dict' >> beam.Map(c.list_to_dict, [
                'segment', 'average_retention_probability',
                'average_predicted_customer_value',
                'average_predicted_order_value', 'average_predicted_purchases',
                'total_customer_value', 'number_of_customers',
                'perc_of_total_customer_value'])
            | 'Write to prediction_summary table' >> io.WriteToBigQuery(
                table=c.TableValueProvider(
                    getattr(runtime_options, c._OPTION_OUTPUT_BQ_PROJECT),
                    getattr(runtime_options, c._OPTION_OUTPUT_BQ_DATASET),
                    'prediction_summary'
                ),
                custom_gcs_temp_location=getattr(runtime_options, c._OPTION_TEMP_GCS_LOCATION),
                validate=False,
                schema='segment:INT64 ,average_retention_probability:FLOAT64'
                       ', average_predicted_customer_value:FLOAT64'
                       ', average_predicted_order_value:FLOAT64'
                       ', average_predicted_purchases:FLOAT64'
                       ', total_customer_value:FLOAT64'
                       ', number_of_customers:FLOAT64'
                       ', perc_of_total_customer_value:FLOAT64',
                write_disposition=io.BigQueryDisposition.WRITE_TRUNCATE,
                create_disposition=io.BigQueryDisposition.CREATE_IF_NEEDED)
        )

        prediction_summary_extra_dimension = (
            prediction_by_customer
            | 'Discard prediction if there is not extra dimension' >>
                beam.FlatMap(c.discard_if_no_extra_dimension,
                             pvalue.AsSingleton(options))
            | beam.Map(lambda x: (x[10], x))  # extra dimension
            | 'Group customer predictions by extra dimension' >>
                beam.GroupByKey()
            | beam.FlatMap(c.generate_prediction_summary_extra_dimension,
                           pvalue.AsSingleton(tot_equity),
                           pvalue.AsSingleton(options))
        )

        _ = (
            prediction_summary_extra_dimension
            | 'prediction_summary_extra_dimension to Dict' >> beam.Map(c.list_to_dict, [
                'extra_dimension', 'average_retention_probability',
                'average_predicted_customer_value',
                'average_predicted_order_value',
                'average_predicted_purchases', 'total_customer_value',
                'number_of_customers', 'perc_of_total_customer_value'])
            | 'Write to prediction_summary_extra_dimension table' >> io.WriteToBigQuery(
                table=c.TableValueProvider(
                    getattr(runtime_options, c._OPTION_OUTPUT_BQ_PROJECT),
                    getattr(runtime_options, c._OPTION_OUTPUT_BQ_DATASET),
                    'prediction_summary_extra_dimension'
                ),
                custom_gcs_temp_location=getattr(runtime_options, c._OPTION_TEMP_GCS_LOCATION),
                validate=False,
                schema='extra_dimension:STRING'
                       ', average_retention_probability:FLOAT64'
                       ', average_predicted_customer_value:FLOAT64'
                       ', average_predicted_order_value:FLOAT64'
                       ', average_predicted_purchases:FLOAT64'
                       ', total_customer_value:FLOAT64'
                       ', number_of_customers:INT64'
                       ', perc_of_total_customer_value:FLOAT64',
                write_disposition=io.BigQueryDisposition.WRITE_TRUNCATE,
                create_disposition=io.BigQueryDisposition.CREATE_IF_NEEDED)
        )


if __name__ == '__main__':
    logging.getLogger().setLevel(logging.INFO)
    run()