<commit_before>package io.eels.component.hive

import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.{ArrayBlockingQueue, CountDownLatch, Executors, TimeUnit}

import com.sksamuel.scalax.collection.BlockingQueueConcurrentIterator
import com.typesafe.config.ConfigFactory
import com.typesafe.scalalogging.slf4j.StrictLogging
import io.eels.{FrameSchema, InternalRow, Sink, Writer}
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.metastore.HiveMetaStoreClient

import scala.collection.JavaConverters._
import scala.collection.mutable

case class HiveSink(dbName: String,
                    tableName: String,
                    ioThreads: Int = 4,
                    dynamicPartitioning: Boolean = true,
                    createTable: Boolean = false,
                    overwriteTable: Boolean = false,
                    format: HiveFormat = HiveFormat.Text,
                    bufferSize: Int = 1000)
                   (implicit fs: FileSystem, hiveConf: HiveConf) extends Sink with StrictLogging {
  logger.info(s"Created HiveSink; createTable=$createTable, overwriteTable=$overwriteTable; format=$format")

  val config = ConfigFactory.load()
  val includePartitionsInData = config.getBoolean("hive.includePartitionsInData")

  def withIOThreads(ioThreads: Int): HiveSink = copy(ioThreads = ioThreads)
  def withDynamicPartitioning(dynamicPartitioning: Boolean): HiveSink = copy(dynamicPartitioning = dynamicPartitioning)

  /**
    * Returns the schema for the hive destination as an Eel format schema object.
    */
  def schema(implicit client: HiveMetaStoreClient): FrameSchema = {
    val schema = client.getSchema(dbName, tableName)
    FrameSchemaFn(schema.asScala)
  }

  override def writer: Writer = {
    logger.debug(s"HiveSinkWriter created")

    implicit val client = new HiveMetaStoreClient(hiveConf)

    val dialect = {
      val format = HiveOps.tableFormat(dbName, tableName)
      logger.debug(s"Table format is $format")
      HiveDialect(format)
    }

    val tablePath = HiveOps.tablePath(dbName, tableName)

    val partitionKeyNames = HiveOps.partitionKeyNames(dbName, tableName)
    logger.debug("Dynamic partitioning enabled: " + partitionKeyNames.mkString(","))

    val targetSchema = HiveSink.this.schema

    // we need an output stream per partition. Since the data can come in unordered, we need to
    // keep open a stream per partition path. This shouldn't be shared amongst threads until its made thread safe.
    val writers = mutable.Map.empty[Path, HiveWriter]

    def getOrCreateHiveWriter(row: InternalRow, sourceSchema: FrameSchema): HiveWriter = {
      val parts = RowPartitionParts(row, partitionKeyNames, sourceSchema)
      val partPath = HiveOps.partitionPath(dbName, tableName, parts, tablePath)
      writers.synchronized {
        writers.getOrElseUpdate(partPath, {
          val filePath = new Path(partPath, "eel_" + System.nanoTime)
          logger.debug(s"Creating hive writer for $filePath")
          if (dynamicPartitioning) {
            if (parts.nonEmpty)
              HiveOps.createPartitionIfNotExists(dbName, tableName, parts)
          } else if (!HiveOps.partitionExists(dbName, tableName, parts)) {
            sys.error(s"Partition $partPath does not exist and dynamicPartitioning=false")
          }
          dialect.writer(sourceSchema, targetSchema, filePath)
        })
      }
    }

    import com.sksamuel.scalax.concurrent.ThreadImplicits.toRunnable

    val queue = new ArrayBlockingQueue[InternalRow](bufferSize)
    val latch = new CountDownLatch(ioThreads)
    val executor = Executors.newFixedThreadPool(ioThreads)
    val count = new AtomicLong(0)
    var schema: FrameSchema = null


    for ( k <- 1 to ioThreads ) {
      executor.submit {
        logger.info(s"HiveSink thread $k started")
        try {
          BlockingQueueConcurrentIterator(queue, InternalRow.PoisonPill).foreach { row =>
            val writer = getOrCreateHiveWriter(row, schema)
            writer.synchronized {
              // need to strip out any partition information from the written data
              // keeping this as a list as I want it ordered and no need to waste cycles on an ordered map
              if (includePartitionsInData) {<commit_msg>Hive  no longer uses partition keys when writing schema #95<commit_after>package io.eels.component.hive

import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.{ArrayBlockingQueue, CountDownLatch, Executors, TimeUnit}

import com.sksamuel.scalax.collection.BlockingQueueConcurrentIterator
import com.typesafe.config.ConfigFactory
import com.typesafe.scalalogging.slf4j.StrictLogging
import io.eels.{FrameSchema, InternalRow, Sink, Writer}
import org.apache.hadoop.fs.{FileSystem, Path}
import org.apache.hadoop.hive.conf.HiveConf
import org.apache.hadoop.hive.metastore.HiveMetaStoreClient

import scala.collection.JavaConverters._
import scala.collection.mutable

case class HiveSink(dbName: String,
                    tableName: String,
                    ioThreads: Int = 4,
                    dynamicPartitioning: Boolean = true,
                    createTable: Boolean = false,
                    overwriteTable: Boolean = false,
                    format: HiveFormat = HiveFormat.Text,
                    bufferSize: Int = 1000)
                   (implicit fs: FileSystem, hiveConf: HiveConf) extends Sink with StrictLogging {
  logger.info(s"Created HiveSink; createTable=$createTable, overwriteTable=$overwriteTable; format=$format")

  val config = ConfigFactory.load()
  val includePartitionsInData = config.getBoolean("hive.includePartitionsInData")

  def withIOThreads(ioThreads: Int): HiveSink = copy(ioThreads = ioThreads)
  def withDynamicPartitioning(dynamicPartitioning: Boolean): HiveSink = copy(dynamicPartitioning = dynamicPartitioning)

  /**
    * Returns the schema for the hive destination as an Eel format schema object.
    */
  def schema(implicit client: HiveMetaStoreClient): FrameSchema = {
    val schema = client.getSchema(dbName, tableName)
    FrameSchemaFn(schema.asScala)
  }

  override def writer: Writer = {
    logger.debug(s"HiveSinkWriter created")

    implicit val client = new HiveMetaStoreClient(hiveConf)

    val dialect = {
      val format = HiveOps.tableFormat(dbName, tableName)
      logger.debug(s"Table format is $format")
      HiveDialect(format)
    }

    val tablePath = HiveOps.tablePath(dbName, tableName)

    val partitionKeyNames = HiveOps.partitionKeyNames(dbName, tableName)
    logger.debug("Dynamic partitioning enabled: " + partitionKeyNames.mkString(","))

    // we need an output stream per partition. Since the data can come in unordered, we need to
    // keep open a stream per partition path. This shouldn't be shared amongst threads until its made thread safe.
    val writers = mutable.Map.empty[Path, HiveWriter]

    def getOrCreateHiveWriter(row: InternalRow, sourceSchema: FrameSchema): HiveWriter = {
      val parts = RowPartitionParts(row, partitionKeyNames, sourceSchema)
      val partPath = HiveOps.partitionPath(dbName, tableName, parts, tablePath)
      writers.synchronized {
        writers.getOrElseUpdate(partPath, {
          val filePath = new Path(partPath, "eel_" + System.nanoTime)
          logger.debug(s"Creating hive writer for $filePath")
          if (dynamicPartitioning) {
            if (parts.nonEmpty)
              HiveOps.createPartitionIfNotExists(dbName, tableName, parts)
          } else if (!HiveOps.partitionExists(dbName, tableName, parts)) {
            sys.error(s"Partition $partPath does not exist and dynamicPartitioning=false")
          }
          val targetSchema = if (includePartitionsInData) {
            HiveSink.this.schema
          } else {
            partitionKeyNames.foldLeft(HiveSink.this.schema)((schema, name) => schema.removeColumn(name))
          }
          dialect.writer(sourceSchema, targetSchema, filePath)
        })
      }
    }

    import com.sksamuel.scalax.concurrent.ThreadImplicits.toRunnable

    val queue = new ArrayBlockingQueue[InternalRow](bufferSize)
    val latch = new CountDownLatch(ioThreads)
    val executor = Executors.newFixedThreadPool(ioThreads)
    val count = new AtomicLong(0)
    var schema: FrameSchema = null


    for ( k <- 1 to ioThreads ) {
      executor.submit {
        logger.info(s"HiveSink thread $k started")
        try {
          BlockingQueueConcurrentIterator(queue, InternalRow.PoisonPill).foreach { row =>
            val writer = getOrCreateHiveWriter(row, schema)
            writer.synchronized {
              // need to strip out any partition information from the written data
              // keeping this as a list as I want it ordered and no need to waste cycles on an ordered map
              if (includePartitionsInData) {
{'fraction_non_alphanumeric': 0.08205012925705293, 'fraction_numerical': 0.0020231538720917163, 'mean_word_length': 3.2070921985815604, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1691972', 'n_tokens_mistral': 2539, 'n_tokens_neox': 2484, 'n_words': 637}