text: # Copyright 2020 The Merlin Authors
  #
  # 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.
  
  from abc import ABC, abstractmethod
  from typing import Iterable, MutableMapping
  
  import yaml
  from google.protobuf import json_format
  from pyspark.sql.types import ArrayType, DataType, DoubleType, FloatType, \
      IntegerType, LongType, StringType
  
  from merlinpyspark.spec.prediction_job_pb2 import BigQuerySink, BigQuerySource, Model, \
      PredictionJob, ResultType, ModelType, SaveMode
  
  
  def printJobConfig(spec_path):
      with open(spec_path, "r") as f:
          print(f"""
  ===============================================================
  Prediction job configuration:
  {f.read()}
  ===============================================================
          """)
  
  
  def load(spec_path: str) -> "JobConfig":
      """
      Read a job spec file (yaml file) and return nested dictionary
      containing the job config
      :param spec_path: path to yaml config file
      :return: JobConfig object
      """
      printJobConfig(spec_path)
      prediction_job = PredictionJob()
      with open(spec_path, "r") as f:
          return JobConfig(
              json_format.ParseDict(yaml.safe_load(f), prediction_job))
  
  
  class JobConfig:
      BIGQUERY_SOURCE = "bigquery_source"
      BIGQUERY_SINK = "bigquery_sink"
      GCS_SOURCE = "gcs_source"
      GCS_SINK = "gcs_sink"
  
      def __init__(self, config: PredictionJob):
          self._config = config
  
      def source(self) -> "SourceConfig":
          src = self._config.WhichOneof("source")
          if src is None:
              raise ValueError("missing mandatory source config")
  
          if src == self.BIGQUERY_SOURCE:
              return BigQuerySourceConfig(self._config.bigquery_source)
          else:
              raise ValueError(f"{src} is not implemented")
  
      def model(self) -> "ModelConfig":
          return ModelConfig(self._config.model)
  
      def sink(self) -> "SinkConfig":
          sink = self._config.WhichOneof("sink")
          if sink is None:
              raise ValueError("missing mandatory sink config")
  
          if sink == self.BIGQUERY_SINK:
              return BigQuerySinkConfig(self._config.bigquery_sink)
          else:
              raise ValueError(f"{sink} is not implemented")
  
  
  class SourceConfig(ABC):
      @abstractmethod
      def source_type(self) -> str:
          pass
  
      @abstractmethod
      def options(self) -> MutableMapping[str, str]:
          pass
  
      @abstractmethod
      def features(self) -> Iterable[str]:
          pass
  
  
  class BigQuerySourceConfig(SourceConfig):
      TYPE = "bigquery"
  
      def __init__(self, bq_src_proto: BigQuerySource):
          self._proto = bq_src_proto
  
      def source_type(self) -> str:
          return self.TYPE
  
      def options(self) -> MutableMapping[str, str]:
          return self._proto.options
  
      def features(self) -> Iterable[str]:
          return self._proto.features
  
      def table(self) -> str:
          return self._proto.table
  
  
  class ModelConfig:
      PRIMITIVE_TYPE_MAP = {
          ResultType.DOUBLE: DoubleType(),
          ResultType.FLOAT: FloatType(),
          ResultType.INTEGER: IntegerType(),
          ResultType.LONG: LongType(),
          ResultType.STRING: StringType(),
      }
  
      def __init__(self, model_proto: Model):
          self._proto = model_proto
  
      def model_type(self) -> str:
          return ModelType.Name(self._proto.type)
  
      def model_uri(self) -> str:
          return self._proto.uri
  
      def result_type(self) -> DataType:
          rt = self._proto.result.type
          if rt == ResultType.ARRAY:
              it = self._proto.result.item_type
              if it in self.PRIMITIVE_TYPE_MAP:
                  return ArrayType(self.PRIMITIVE_TYPE_MAP.get(it))
              raise ValueError(f"unknown item type for array: {it}")
  
          if rt in self.PRIMITIVE_TYPE_MAP:
              return self.PRIMITIVE_TYPE_MAP.get(rt)
  
          raise ValueError(f"unknown result type: {rt}")
  
  
  class SinkConfig(ABC):
      @abstractmethod
      def sink_type(self) -> str:
          pass
  
      @abstractmethod
      def result_column(self) -> str:
          pass
  
      @abstractmethod
      def options(self) -> MutableMapping[str, str]:
          pass
  
      @abstractmethod
      def save_mode(self) -> str:
          pass
  
  
  class BigQuerySinkConfig(SinkConfig):
      TYPE = "bigquery"
      DEFAULT_RESULT_COLUMN = "prediction"
      
      def __init__(self, bq_sink_proto: BigQuerySink):
          self._proto = bq_sink_proto
  
      def sink_type(self) -> str:
          return self.TYPE
  
      def result_column(self) -> str:
          col = self._proto.result_column
          if col == "":
              return self.DEFAULT_RESULT_COLUMN
          return col
  
      def options(self) -> MutableMapping[str, str]:
          return self._proto.options
  
      def save_mode(self) -> str:
          return SaveMode.Name(self._proto.save_mode).lower()
  
      def table(self) -> str:
          return self._proto.table
  
      def staging_bucket(self) -> str:
          return self._proto.staging_bucket
  