Apache Spark Kafka SQL Connector

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.streaming.StreamingQuery
import org.apache.spark.sql.functions._

// The connector is registered automatically as "kafka" data source
// No direct imports of connector classes are needed

import org.apache.kafka.common.TopicPartition
import org.apache.spark.sql.kafka010.PartitionOffsetMap

import org.apache.spark.sql.SparkSession

val spark = SparkSession
  .builder()
  .appName("KafkaStreaming")
  .getOrCreate()

// Read from Kafka using structured streaming
val kafkaStream = spark
  .readStream
  .format("kafka")
  .option("kafka.bootstrap.servers", "localhost:9092")
  .option("subscribe", "my-topic")
  .option("startingOffsets", "latest")
  .load()

// Process the stream
val processedStream = kafkaStream
  .selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)", "topic", "partition", "offset", "timestamp")
  .writeStream
  .outputMode("append")
  .format("console")
  .start()

val kafkaBatch = spark
  .read
  .format("kafka")
  .option("kafka.bootstrap.servers", "localhost:9092")
  .option("subscribe", "my-topic")
  .option("startingOffsets", "earliest")
  .option("endingOffsets", "latest")
  .load()

val dataFrame = spark.createDataFrame(Seq(
  ("key1", "value1"),
  ("key2", "value2")
)).toDF("key", "value")

dataFrame
  .selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")
  .write
  .format("kafka")
  .option("kafka.bootstrap.servers", "localhost:9092")
  .option("topic", "output-topic")
  .save()

// Automatically registered - no direct usage
class KafkaSourceProvider extends DataSourceRegister 
  with StreamSourceProvider 
  with StreamSinkProvider 
  with RelationProvider 
  with CreatableRelationProvider 
  with SimpleTableProvider

// Consumer strategies are configured via options:
// .option("subscribe", "topic1,topic2,topic3")
// .option("subscribePattern", "prefix-.*")
// .option("assign", """{"topic1":[0,1,2],"topic2":[0,1]}""")

sealed trait ConsumerStrategy {
  def createConsumer(kafkaParams: ju.Map[String, Object]): Consumer[Array[Byte], Array[Byte]]
  def createAdmin(kafkaParams: ju.Map[String, Object]): Admin
  def assignedTopicPartitions(admin: Admin): Set[TopicPartition]
}

// Offset range limits for controlling read boundaries
sealed trait KafkaOffsetRangeLimit

case object EarliestOffsetRangeLimit extends KafkaOffsetRangeLimit
case object LatestOffsetRangeLimit extends KafkaOffsetRangeLimit
case class SpecificOffsetRangeLimit(partitionOffsets: Map[TopicPartition, Long]) extends KafkaOffsetRangeLimit
case class SpecificTimestampRangeLimit(topicTimestamps: Map[TopicPartition, Long], strategy: StrategyOnNoMatchStartingOffset.Value) extends KafkaOffsetRangeLimit
case class GlobalTimestampRangeLimit(timestamp: Long, strategy: StrategyOnNoMatchStartingOffset.Value) extends KafkaOffsetRangeLimit

// Schema with headers disabled (default)
val schemaWithoutHeaders = StructType(Array(
  StructField("key", BinaryType),
  StructField("value", BinaryType),
  StructField("topic", StringType),
  StructField("partition", IntegerType),
  StructField("offset", LongType),
  StructField("timestamp", TimestampType),
  StructField("timestampType", IntegerType)
))

// Schema with headers enabled (.option("includeHeaders", "true"))
val schemaWithHeaders = schemaWithoutHeaders.add(
  StructField("headers", ArrayType(StructType(Array(
    StructField("key", StringType),
    StructField("value", BinaryType)
  ))))
)

def kafkaSchema(includeHeaders: Boolean): StructType

// Micro-batch streaming
class KafkaMicroBatchStream extends MicroBatchStream 
  with SupportsTriggerAvailableNow 
  with ReportsSourceMetrics {
  
  def initialOffset(): Offset
  def latestOffset(): Offset
  def latestOffset(startOffset: Offset, readLimit: ReadLimit): Offset
  def planInputPartitions(start: Offset, end: Offset): Array[InputPartition]
  def createReaderFactory(): PartitionReaderFactory
  def commit(end: Offset): Unit
  def stop(): Unit
}

// Continuous streaming  
class KafkaContinuousStream extends ContinuousStream {
  def mergeOffsets(offsets: Array[PartitionOffset]): Offset
  def initialOffset(): Offset  
  def deserializeOffset(json: String): Offset
  def commit(end: Offset): Unit
  def stop(): Unit
}

class KafkaBatch extends Batch {
  def planInputPartitions(): Array[InputPartition]
  def createReaderFactory(): PartitionReaderFactory
}

class KafkaBatchPartitionReader extends PartitionReader[InternalRow] {
  def next(): Boolean
  def get(): UnsafeRow
  def close(): Unit
  def currentMetricsValues(): Array[CustomTaskMetric]
}

// Core writer functionality
object KafkaWriter {
  val TOPIC_ATTRIBUTE_NAME: String = "topic"
  val KEY_ATTRIBUTE_NAME: String = "key"
  val VALUE_ATTRIBUTE_NAME: String = "value"
  val HEADERS_ATTRIBUTE_NAME: String = "headers"
  val PARTITION_ATTRIBUTE_NAME: String = "partition"
  
  def write(sparkSession: SparkSession, queryExecution: QueryExecution, 
            kafkaParams: ju.Map[String, Object], topic: Option[String]): Unit
}

// V2 DataSource write implementation
case class KafkaWrite(topic: Option[String], producerParams: ju.Map[String, Object], schema: StructType) extends Write {
  def description(): String
  def toBatch: BatchWrite
  def toStreaming: StreamingWrite
}

object KafkaExceptions {
  def mismatchedTopicPartitionsBetweenEndOffsetAndPrefetched(
      tpsForPrefetched: Set[TopicPartition],
      tpsForEndOffset: Set[TopicPartition]): SparkException

  def endOffsetHasGreaterOffsetForTopicPartitionThanPrefetched(
      prefetchedOffset: Map[TopicPartition, Long],
      endOffset: Map[TopicPartition, Long]): SparkException

  def lostTopicPartitionsInEndOffsetWithTriggerAvailableNow(
      tpsForLatestOffset: Set[TopicPartition],
      tpsForEndOffset: Set[TopicPartition]): SparkException

  def endOffsetHasGreaterOffsetForTopicPartitionThanLatestWithTriggerAvailableNow(
      latestOffset: Map[TopicPartition, Long],
      endOffset: Map[TopicPartition, Long]): SparkException
}

## Custom Metrics

The connector exposes custom metrics for monitoring:

- `offsetOutOfRange`: Number of offsets that were out of range
- `dataLoss`: Number of data loss events detected

These metrics integrate with Spark's metrics system and can be monitored through Spark UI and external monitoring systems.