text: """KafkaSource entity."""
  
  from pyspark.sql.dataframe import DataFrame, StructType
  from pyspark.sql.functions import col, struct
  
  from butterfree.clients import SparkClient
  from butterfree.configs import environment
  from butterfree.extract.pre_processing import explode_json_column
  from butterfree.extract.readers.reader import Reader
  
  
  class KafkaReader(Reader):
      """Responsible for get data from a Kafka topic.
  
      Attributes:
          id: unique string id for register the reader as a view on the metastore
          value_schema: expected schema of the default column named "value" from Kafka.
          topic: string with the Kafka topic name to subscribe.
          connection_string: string with hosts and ports to connect.
              The string need to be in the format: host1:port,host2:port,...,hostN:portN.
              The argument is not necessary if is passed as a environment variable
              named KAFKA_CONSUMER_CONNECTION_STRING.
          topic_options: additional options for consuming from topic. See docs:
              https://spark.apache.org/docs/latest/structured-streaming-kafka-integration.html.
          stream: flag to indicate the reading mode: stream or batch
  
      The default df schema coming from Kafka reader of Spark is the following:
          key:string
          value:string
          topic:string
          partition:integer
          offset:long
          timestamp:timestamp
          timestampType:integer
  
      But using this reader and passing the desired schema under value_schema we would
      have the following result:
  
      With value_schema declared as:
  
      >>> value_schema = StructType(
      ...     [
      ...         StructField("ts", LongType(), nullable=True),
      ...         StructField("id", LongType(), nullable=True),
      ...         StructField("type", StringType(), nullable=True),
      ...     ]
      ... )
  
      The output df schema would be:
          ts:long
          id:long
          type:string
          kafka_metadata:struct
              key:string
              topic:string
              value:string
              partition:integer
              offset:long
              timestamp:timestamp
              timestampType:integer
  
      Instantiation example:
  
      >>> from butterfree.extract.readers import KafkaReader
      >>> from butterfree.clients import SparkClient
      >>> from pyspark.sql.types import StructType, StructField, StringType, LongType
      >>> spark_client = SparkClient()
      >>> value_schema = StructType(
      ...     [
      ...         StructField("ts", LongType(), nullable=True),
      ...         StructField("id", LongType(), nullable=True),
      ...         StructField("type", StringType(), nullable=True),
      ...     ]
      ... )
      >>> kafka_reader = KafkaReader(
      ...                 id="kafka_reader_id",
      ...                 topic="topic",
      ...                 value_schema=value_schema
      ...                 connection_string="host1:port,host2:port",
      ...                )
      >>> df = kafka_reader.consume(spark_client)
  
      This last method will use the Spark Client, as default, to read
      the desired topic, loading data into a dataframe, according to
      KafkaReader class arguments.
  
      In this case, however, a temp view will be created, containing
      the transformed data.
  
      """
  
      KAFKA_COLUMNS = [
          "key",
          "topic",
          "value",
          "partition",
          "offset",
          "timestamp",
          "timestampType",
      ]
  
      __name__ = "Kafka Reader"
  
      def __init__(
          self,
          id: str,
          topic: str,
          value_schema: StructType,
          connection_string: str = None,
          topic_options: dict = None,
          stream: bool = True,
      ):
          super().__init__(id)
          if not isinstance(topic, str):
              raise ValueError("topic must be a string with the topic name")
          if not isinstance(value_schema, StructType):
              raise ValueError(
                  "value_schema must be a StructType with the schema "
                  'of the JSON presented in "value" Kafka column'
              )
          self.topic = topic
          self.value_schema = value_schema
          self.connection_string = connection_string or environment.get_variable(
              "KAFKA_CONSUMER_CONNECTION_STRING"
          )
          self.options = dict(
              {
                  "kafka.bootstrap.servers": self.connection_string,
                  "subscribe": self.topic,
              },
              **topic_options if topic_options else {},
          )
          self.stream = stream
  
      def _struct_df(self, df: DataFrame) -> DataFrame:
          """Struct the output dataframe generated by the reader.
  
          Under the default "value" column coming from Kafka there are the custom
          fields created by some producer. This function will struct the dataframe as
          to get all desired fields from "value" and insert all Kafka default columns,
          including "value", under "kafka_metadata" column. It is important to notice
          that the declared value_schema suffer from the same effects described in
          explode_json_column method in pre_processing module.
  
          Args:
              df: direct dataframe output from from KafkaReader.
  
          Returns:
              Structured dataframe with kafka value fields as columns.
                  All other default fields from Kafka will be stored under
                  "kafka_metadata" column.
  
          """
          df = df.withColumn("kafka_metadata", struct(*self.KAFKA_COLUMNS))
          df = explode_json_column(df, column="value", json_schema=self.value_schema)
          return df.select(
              [field.name for field in self.value_schema] + ["kafka_metadata"]
          )
  
      def consume(self, client: SparkClient) -> DataFrame:
          """Extract data from a kafka topic.
  
          When stream mode it will get all the new data arriving at the topic in a
          streaming dataframe. When not in stream mode it will get all data
          available in the kafka topic.
  
          Args:
              client: client responsible for connecting to Spark session.
  
          Returns:
              Dataframe with data from topic.
  
          """
          # read using client and cast key and value columns from binary to string
          raw_df = (
              client.read(format="kafka", options=self.options, stream=self.stream)
              .withColumn("key", col("key").cast("string"))
              .withColumn("value", col("value").cast("string"))
          )
  
          # apply schema defined in self.value_schema
          return self._struct_df(raw_df)
  