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

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

OLTP Spark 3.2 Connector for Azure Cosmos DB SQL API

There is a newer version: 4.33.1

Show newest version

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

package com.azure.cosmos.spark

import com.azure.cosmos.models.CosmosParameterizedQuery
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 ItemsScanPartitionReaderFactory
(
  config: Map[String, String],
  readSchema: StructType,
  cosmosQuery: CosmosParameterizedQuery,
  diagnosticsOperationContext: DiagnosticsContext,
  cosmosClientStateHandles: Broadcast[CosmosClientMetadataCachesSnapshots],
  diagnosticsConfig: DiagnosticsConfig
) extends PartitionReaderFactory {

  @transient private lazy val log = LoggerHelper.getLogger(diagnosticsConfig, this.getClass)
  log.logInfo(s"Instantiated ${this.getClass.getSimpleName}")

  override def createReader(partition: InputPartition): PartitionReader[InternalRow] = {
    val feedRange = partition.asInstanceOf[CosmosInputPartition].feedRange
    log.logInfo(s"Creating an ItemsPartitionReader to read from feed-range [$feedRange]")

    ItemsPartitionReader(config,
      feedRange,
      readSchema,
      cosmosQuery,
      diagnosticsOperationContext,
      cosmosClientStateHandles,
      diagnosticsConfig
    )
  }
}