<commit_before>/*
 * Copyright 2018 Google Inc. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.cloud.spark.bigquery.direct

import java.sql.{Date, Timestamp}
import java.util.UUID
import java.util.concurrent.{Callable, TimeUnit}

import com.google.api.gax.core.CredentialsProvider
import com.google.api.gax.rpc.FixedHeaderProvider
import com.google.auth.Credentials
import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions
import com.google.cloud.bigquery.storage.v1.{BigQueryReadClient, BigQueryReadSettings, CreateReadSessionRequest, DataFormat, ReadSession}
import com.google.cloud.bigquery.{BigQuery, BigQueryOptions, JobInfo, QueryJobConfiguration, Schema, StandardTableDefinition, TableDefinition, TableId, TableInfo}
import com.google.cloud.spark.bigquery.direct.DirectBigQueryRelation.{compileFilter, compileValue, quote}
import com.google.cloud.spark.bigquery.{BigQueryRelation, BigQueryUtil, BuildInfo, SchemaConverters, SparkBigQueryConnectorUserAgentProvider, SparkBigQueryOptions}
import com.google.common.cache.{Cache, CacheBuilder}
import org.apache.spark.Partition
import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.sources._
import org.apache.spark.sql.types.{StructField, StructType}

import scala.collection.JavaConverters._

private[bigquery] class DirectBigQueryRelation(
    options: SparkBigQueryOptions,
    table: TableInfo,
    getClient: SparkBigQueryOptions => BigQueryReadClient =
         DirectBigQueryRelation.createReadClient,
    bigQueryClient: SparkBigQueryOptions => BigQuery =
         DirectBigQueryRelation.createBigQueryClient)
    (@transient override val sqlContext: SQLContext)
    extends BigQueryRelation(options, table)(sqlContext)
        with TableScan with PrunedScan with PrunedFilteredScan {

  val tablePath: String =
    DirectBigQueryRelation.toTablePath(tableId)

  lazy val bigQuery = bigQueryClient(options)

  val topLevelFields = SchemaConverters
    .toSpark(table.getDefinition[TableDefinition].getSchema)
    .fields
    .map(field => (field.name, field))
    .toMap

  // used to cache the table instances in order to avoid redundant queries to
  // the BigQuery service
  case class DestinationTableBuilder(querySql: String) extends Callable[TableInfo] {
    override def call(): TableInfo = createTableFromQuery(querySql)
  }
  val destinationTableCache: Cache[String, TableInfo] =
    CacheBuilder.newBuilder()
      .expireAfterWrite(15, TimeUnit.MINUTES)
      .maximumSize(1000)
      .build()

  /**
   * Default parallelism to 1 reader per 400MB, which should be about the maximum allowed by the
   * BigQuery Storage API. The number of partitions returned may be significantly less depending
   * on a number of factors.
   */
  val DEFAULT_BYTES_PER_PARTITION = 400L * 1000 * 1000

  override val needConversion: Boolean = false
  override val sizeInBytes: Long = getNumBytes(defaultTableDefinition)
  // no added filters and with all column
  lazy val defaultTableDefinition: TableDefinition = table.getDefinition[TableDefinition]

  override def buildScan(): RDD[Row] = {
    buildScan(schema.fieldNames)
  }

  override def buildScan(requiredColumns: Array[String]): RDD[Row] = {
    buildScan(requiredColumns, Array())
  }

  override def buildScan(requiredColumns: Array[String], filters: Array[Filter]): RDD[Row] = {
    logInfo(
      s"""
         |Querying table $tableName, parameters sent from Spark:
         |requiredColumns=[${requiredColumns.mkString(",")}],
         |filters=[${filters.map(_.toString).mkString(",")}]"""
        .stripMargin.replace('\n', ' ').trim)
    val filter = getCompiledFilter(filters)
    val actualTable = getActualTable(requiredColumns, filter)

    logInfo(
      s"""
         |Going to read from ${BigQueryUtil.friendlyTableName(actualTable.getTableId)}
         |columns=[${requiredColumns.mkString(", ")}],
         |filter='$filter'"""
        .stripMargin.replace('\n', ' ').trim)

    if (options.optimizedEmptyProjection && requiredColumns.isEmpty) {
      generateEmptyRowRDD(actualTable, filter)
    } else {
      if (requiredColumns.isEmpty) {
        logDebug(s"Not using optimized empty projection")
      }
      val actualTableDefinition = actualTable.getDefinition[StandardTableDefinition]
      val actualTablePath = DirectBigQueryRelation.toTablePath(actualTable.getTableId)
      val readOptions = TableReadOptions.newBuilder()
        .addAllSelectedFields(requiredColumns.toList.asJava)
        .setRowRestriction(filter)
        .build()
      val requiredColumnSet = requiredColumns.toSet
      val prunedSchema = Schema.of(
        actualTableDefinition.getSchema.getFields.asScala
          .filter(f => requiredColumnSet.contains(f.getName)).asJava)

      val client = getClient(options)

      val maxNumPartitionsRequested = getMaxNumPartitionsRequested(actualTableDefinition)

      val readDataFormat = options.readDataFormat

      try {
        val session = client.createReadSession(
          CreateReadSessionRequest.newBuilder()
            .setParent(s"projects/${options.parentProject}")
            .setReadSession(ReadSession.newBuilder()
              .setDataFormat(readDataFormat)
              .setReadOptions(readOptions)
              .setTable(actualTablePath
            ))
            .setMaxStreamCount(maxNumPartitionsRequested)
            .build())
        val partitions = session.getStreamsList.asScala.map(_.getName)
          .zipWithIndex.map { case (name, i) => BigQueryPartition(name, i) }
          .toArray

        logInfo(s"Created read session for table '$tableName': ${session.getName}")

        // This is spammy, but it will make it clear to users the number of partitions they got and
        // why.
        if (!maxNumPartitionsRequested.equals(partitions.length)) {
          logInfo(
            s"""Requested $maxNumPartitionsRequested max partitions, but only
               |received ${partitions.length} from the BigQuery Storage API for
               |session ${session.getName}. Notice that the number of streams in
               |actual may be lower than the requested number, depending on the
               |amount parallelism that is reasonable for the table and the
               |maximum amount of parallelism allowed by the system."""
              .stripMargin.replace('\n', ' '))
        }

        BigQueryRDD.scanTable(
          sqlContext,
          partitions.asInstanceOf[Array[Partition]],
          session,
          prunedSchema,
          requiredColumns,
          options,
          getClient,
          bigQueryClient).asInstanceOf[RDD[Row]]

      } finally {
        // scanTable returns immediately not after the actual data is read.
        client.close()
      }
    }
  }

  def  generateEmptyRowRDD(tableInfo: TableInfo, filter: String) : RDD[Row] = {
    DirectBigQueryRelation.emptyRowRDDsCreated += 1
    val numberOfRows: Long = if (filter.length == 0) {
      // no need to query the table
      tableInfo.getNumRows.longValue
    } else {
      // run a query
      val table = DirectBigQueryRelation.toSqlTableReference(tableInfo.getTableId)
      val sql = s"SELECT COUNT(*) from `$table` WHERE $filter"
      val result = bigQuery.query(QueryJobConfiguration.of(sql))
      result.iterateAll.iterator.next.get(0).getLongValue
    }
    logDebug(s"Creating a DataFrame of empty roes of size $numberOfRows")
    sqlContext.sparkContext.range(0, numberOfRows)
      .map(_ => InternalRow.empty)
      .asInstanceOf[RDD[Row]]
  }

  def getActualTable(
      requiredColumns: Array[String],<commit_msg>Fixing typo<commit_after>/*
 * Copyright 2018 Google Inc. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.cloud.spark.bigquery.direct

import java.sql.{Date, Timestamp}
import java.util.UUID
import java.util.concurrent.{Callable, TimeUnit}

import com.google.api.gax.core.CredentialsProvider
import com.google.api.gax.rpc.FixedHeaderProvider
import com.google.auth.Credentials
import com.google.cloud.bigquery.storage.v1.ReadSession.TableReadOptions
import com.google.cloud.bigquery.storage.v1.{BigQueryReadClient, BigQueryReadSettings, CreateReadSessionRequest, DataFormat, ReadSession}
import com.google.cloud.bigquery.{BigQuery, BigQueryOptions, JobInfo, QueryJobConfiguration, Schema, StandardTableDefinition, TableDefinition, TableId, TableInfo}
import com.google.cloud.spark.bigquery.direct.DirectBigQueryRelation.{compileFilter, compileValue, quote}
import com.google.cloud.spark.bigquery.{BigQueryRelation, BigQueryUtil, BuildInfo, SchemaConverters, SparkBigQueryConnectorUserAgentProvider, SparkBigQueryOptions}
import com.google.common.cache.{Cache, CacheBuilder}
import org.apache.spark.Partition
import org.apache.spark.rdd.RDD
import org.apache.spark.sql._
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.sources._
import org.apache.spark.sql.types.{StructField, StructType}

import scala.collection.JavaConverters._

private[bigquery] class DirectBigQueryRelation(
    options: SparkBigQueryOptions,
    table: TableInfo,
    getClient: SparkBigQueryOptions => BigQueryReadClient =
         DirectBigQueryRelation.createReadClient,
    bigQueryClient: SparkBigQueryOptions => BigQuery =
         DirectBigQueryRelation.createBigQueryClient)
    (@transient override val sqlContext: SQLContext)
    extends BigQueryRelation(options, table)(sqlContext)
        with TableScan with PrunedScan with PrunedFilteredScan {

  val tablePath: String =
    DirectBigQueryRelation.toTablePath(tableId)

  lazy val bigQuery = bigQueryClient(options)

  val topLevelFields = SchemaConverters
    .toSpark(table.getDefinition[TableDefinition].getSchema)
    .fields
    .map(field => (field.name, field))
    .toMap

  // used to cache the table instances in order to avoid redundant queries to
  // the BigQuery service
  case class DestinationTableBuilder(querySql: String) extends Callable[TableInfo] {
    override def call(): TableInfo = createTableFromQuery(querySql)
  }
  val destinationTableCache: Cache[String, TableInfo] =
    CacheBuilder.newBuilder()
      .expireAfterWrite(15, TimeUnit.MINUTES)
      .maximumSize(1000)
      .build()

  /**
   * Default parallelism to 1 reader per 400MB, which should be about the maximum allowed by the
   * BigQuery Storage API. The number of partitions returned may be significantly less depending
   * on a number of factors.
   */
  val DEFAULT_BYTES_PER_PARTITION = 400L * 1000 * 1000

  override val needConversion: Boolean = false
  override val sizeInBytes: Long = getNumBytes(defaultTableDefinition)
  // no added filters and with all column
  lazy val defaultTableDefinition: TableDefinition = table.getDefinition[TableDefinition]

  override def buildScan(): RDD[Row] = {
    buildScan(schema.fieldNames)
  }

  override def buildScan(requiredColumns: Array[String]): RDD[Row] = {
    buildScan(requiredColumns, Array())
  }

  override def buildScan(requiredColumns: Array[String], filters: Array[Filter]): RDD[Row] = {
    logInfo(
      s"""
         |Querying table $tableName, parameters sent from Spark:
         |requiredColumns=[${requiredColumns.mkString(",")}],
         |filters=[${filters.map(_.toString).mkString(",")}]"""
        .stripMargin.replace('\n', ' ').trim)
    val filter = getCompiledFilter(filters)
    val actualTable = getActualTable(requiredColumns, filter)

    logInfo(
      s"""
         |Going to read from ${BigQueryUtil.friendlyTableName(actualTable.getTableId)}
         |columns=[${requiredColumns.mkString(", ")}],
         |filter='$filter'"""
        .stripMargin.replace('\n', ' ').trim)

    if (options.optimizedEmptyProjection && requiredColumns.isEmpty) {
      generateEmptyRowRDD(actualTable, filter)
    } else {
      if (requiredColumns.isEmpty) {
        logDebug(s"Not using optimized empty projection")
      }
      val actualTableDefinition = actualTable.getDefinition[StandardTableDefinition]
      val actualTablePath = DirectBigQueryRelation.toTablePath(actualTable.getTableId)
      val readOptions = TableReadOptions.newBuilder()
        .addAllSelectedFields(requiredColumns.toList.asJava)
        .setRowRestriction(filter)
        .build()
      val requiredColumnSet = requiredColumns.toSet
      val prunedSchema = Schema.of(
        actualTableDefinition.getSchema.getFields.asScala
          .filter(f => requiredColumnSet.contains(f.getName)).asJava)

      val client = getClient(options)

      val maxNumPartitionsRequested = getMaxNumPartitionsRequested(actualTableDefinition)

      val readDataFormat = options.readDataFormat

      try {
        val session = client.createReadSession(
          CreateReadSessionRequest.newBuilder()
            .setParent(s"projects/${options.parentProject}")
            .setReadSession(ReadSession.newBuilder()
              .setDataFormat(readDataFormat)
              .setReadOptions(readOptions)
              .setTable(actualTablePath
            ))
            .setMaxStreamCount(maxNumPartitionsRequested)
            .build())
        val partitions = session.getStreamsList.asScala.map(_.getName)
          .zipWithIndex.map { case (name, i) => BigQueryPartition(name, i) }
          .toArray

        logInfo(s"Created read session for table '$tableName': ${session.getName}")

        // This is spammy, but it will make it clear to users the number of partitions they got and
        // why.
        if (!maxNumPartitionsRequested.equals(partitions.length)) {
          logInfo(
            s"""Requested $maxNumPartitionsRequested max partitions, but only
               |received ${partitions.length} from the BigQuery Storage API for
               |session ${session.getName}. Notice that the number of streams in
               |actual may be lower than the requested number, depending on the
               |amount parallelism that is reasonable for the table and the
               |maximum amount of parallelism allowed by the system."""
              .stripMargin.replace('\n', ' '))
        }

        BigQueryRDD.scanTable(
          sqlContext,
          partitions.asInstanceOf[Array[Partition]],
          session,
          prunedSchema,
          requiredColumns,
          options,
          getClient,
          bigQueryClient).asInstanceOf[RDD[Row]]

      } finally {
        // scanTable returns immediately not after the actual data is read.
        client.close()
      }
    }
  }

  def  generateEmptyRowRDD(tableInfo: TableInfo, filter: String) : RDD[Row] = {
    DirectBigQueryRelation.emptyRowRDDsCreated += 1
    val numberOfRows: Long = if (filter.length == 0) {
      // no need to query the table
      tableInfo.getNumRows.longValue
    } else {
      // run a query
      val table = DirectBigQueryRelation.toSqlTableReference(tableInfo.getTableId)
      val sql = s"SELECT COUNT(*) from `$table` WHERE $filter"
      val result = bigQuery.query(QueryJobConfiguration.of(sql))
      result.iterateAll.iterator.next.get(0).getLongValue
    }
    logDebug(s"Creating a DataFrame of empty rows with size $numberOfRows")
    sqlContext.sparkContext.range(0, numberOfRows)
      .map(_ => InternalRow.empty)
      .asInstanceOf[RDD[Row]]
  }

  def getActualTable(
      requiredColumns: Array[String],
{'fraction_non_alphanumeric': 0.09142996464708034, 'fraction_numerical': 0.004266731683530416, 'mean_word_length': 3.4535830618892507, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '2502482', 'n_tokens_mistral': 4573, 'n_tokens_neox': 4436, 'n_words': 1136}