com.azure.cosmos.spark.ChangeFeedScanPartitionReaderFactory.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of azure-cosmos-spark_3-5_2-12 Show documentation

OLTP Spark 3.5 Connector for Azure Cosmos DB SQL API

There is a newer version: 4.34.0

Show newest version

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

package com.azure.cosmos.spark

import com.azure.cosmos.spark.diagnostics.{DiagnosticsContext, LoggerHelper}
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.connector.read.{InputPartition, PartitionReader, PartitionReaderFactory}
import org.apache.spark.sql.types.StructType

private case class ChangeFeedScanPartitionReaderFactory
(
  config: Map[String, String],
  readSchema: StructType,
  diagnosticsOperationContext: DiagnosticsContext,
  cosmosClientStateHandles: Broadcast[CosmosClientMetadataCachesSnapshots],
  diagnosticsConfig: DiagnosticsConfig,
  sparkEnvironmentInfo: String
) extends PartitionReaderFactory {

  @transient private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass)

  log.logTrace(s"Instantiated ${this.getClass.getSimpleName}")

  override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
    val changeFeedPartition = partition.asInstanceOf[CosmosInputPartition]
    ChangeFeedPartitionReader(
      changeFeedPartition,
      config,
      readSchema,
      diagnosticsOperationContext,
      cosmosClientStateHandles,
      diagnosticsConfig,
      sparkEnvironmentInfo)
  }
}