Bug fixes on 0.2.2 + adding standard sql support + documentation

.travis.yml

@@ -1,8 +1,8 @@
 language: scala
 scala:
   - 2.10.6
-  - 2.11.8
-jdk: oraclejdk7
+  - 2.11.12
+jdk: oraclejdk8
 
 cache:
   directories:

README.md

@@ -17,11 +17,11 @@ Google BigQuery support for Spark, SQL, and DataFrames.
 | 0.2.x | 2.x.y | Active development |
 | 0.1.x | 1.x.y | Development halted |
 
-To use the package in a Google [Cloud Dataproc](https://cloud.google.com/dataproc/) cluster:
+Building:
 
-install `org.apache.avro_avro-ipc-1.7.7.jar` to `~/.ivy2/jars`
+```sbt clean assembly```
 
-`spark-shell --packages com.spotify:spark-bigquery_2.10:0.2.2`
+Assembly doesn't like the latest version of Java (currently 11) so set JAVA_HOME to point to Java 8.
 
 To use it in a local SBT console:
 
@@ -45,26 +45,18 @@ Usage:
 // Load everything from a table
 val table = sqlContext.bigQueryTable("bigquery-public-data:samples.shakespeare")
 
-// Load results from a SQL query
-// Only legacy SQL dialect is supported for now
+// Load results from a legacy SQL query
 val df = sqlContext.bigQuerySelect(
   "SELECT word, word_count FROM [bigquery-public-data:samples.shakespeare]")
+
+// Load results from a standard SQL query
+val df = sqlContext.bigQuerySelect(
+  "SELECT word, word_count FROM `bigquery-public-data.samples.shakespeare`", useStandardSql = true)
 
-// Save data to a table
+  // Save data to a table
 df.saveAsBigQueryTable("my-project:my_dataset.my_table")
 ```
 
-If you'd like to write nested records to BigQuery, be sure to specify an Avro Namespace.
-BigQuery is unable to load Avro Namespaces with a leading dot (`.nestedColumn`) on nested records.
-
-```scala
-// BigQuery is able to load fields with namespace 'myNamespace.nestedColumn'
-df.saveAsBigQueryTable("my-project:my_dataset.my_table", tmpWriteOptions = Map("recordNamespace" -> "myNamespace"))
-```
-See also
-[Loading Avro Data from Google Cloud Storage](https://cloud.google.com/bigquery/docs/loading-data-cloud-storage-avro)
-for data type mappings and limitations. For example loading arrays of arrays is not supported.
-
 # License
 
 Copyright 2016 Spotify AB.

build.sbt

@@ -17,22 +17,22 @@
 
 name := "spark-bigquery"
 organization := "com.spotify"
-scalaVersion := "2.12.8"
-crossScalaVersions := Seq("2.10.6", "2.11.11", "2.12.8")
+scalaVersion := "2.11.12"
+crossScalaVersions := Seq("2.10.6", "2.11.12")
 
 spName := "spotify/spark-bigquery"
 sparkVersion := "2.4.0"
-sparkComponents := Seq("core", "sql", "avro")
+sparkComponents := Seq("core", "sql")
 spAppendScalaVersion := true
 spIncludeMaven := true
 
 libraryDependencies ++= Seq(
   "com.databricks" %% "spark-avro" % "4.0.0",
-  "com.google.cloud.bigdataoss" % "bigquery-connector" % "hadoop2-0.13.13"
+  "com.google.cloud.bigdataoss" % "bigquery-connector" % "0.10.2-hadoop2"
     exclude ("com.google.guava", "guava-jdk5"),
   "org.slf4j" % "slf4j-simple" % "1.7.21",
   "joda-time" % "joda-time" % "2.9.3",
-  "org.scalatest" %% "scalatest" % "3.0.7" % "test"
+  "org.scalatest" %% "scalatest" % "2.2.1" % "test"
 )
 
 assemblyMergeStrategy in assembly := {

src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala

@@ -17,6 +17,7 @@
 
 package com.spotify.spark.bigquery
 
+import java.io.FileInputStream
 import java.util.UUID
 import java.util.concurrent.TimeUnit
 
@@ -39,16 +40,14 @@ import scala.util.Random
 import scala.util.control.NonFatal
 
 private[bigquery] object BigQueryClient {
-  val STAGING_DATASET_ID = "bq.staging_dataset.id"
+  // the following is specific to this connector...
   val STAGING_DATASET_PREFIX = "bq.staging_dataset.prefix"
   val STAGING_DATASET_PREFIX_DEFAULT = "spark_bigquery_staging_"
   val STAGING_DATASET_LOCATION = "bq.staging_dataset.location"
+
   val STAGING_DATASET_LOCATION_DEFAULT = "US"
   val STAGING_DATASET_TABLE_EXPIRATION_MS = 86400000L
   val STAGING_DATASET_DESCRIPTION = "Spark BigQuery staging dataset"
-  val QUERY_JOB_PRIORITY = "bq.query_job.priority"
-  val JOB_PRIORITY_INTERACTIVE = "INTERACTIVE"
-  val JOB_PRIORITY_BATCH = "BATCH"
 
   private var instance: BigQueryClient = null
 
@@ -68,48 +67,74 @@ private[bigquery] class BigQueryClient(conf: Configuration) {
 
   private val SCOPES = List(BigqueryScopes.BIGQUERY).asJava
 
-  private val bigquery: Bigquery = {
-    val credential = GoogleCredential.getApplicationDefault.createScoped(SCOPES)
-    new Bigquery.Builder(new NetHttpTransport, new JacksonFactory, credential)
+  private val googleCredential: GoogleCredential = {
+    // TODO get this working for pk12 credentials
+
+    val jsonKeyFileLocation: String = conf.get("fs.gs.auth.service.account.json.keyfile", "")
+    //"fs.gs.auth.service.account.json.keyfile"
+    if (jsonKeyFileLocation.nonEmpty) {
+      var optionFileInputStream: Option[FileInputStream] = None
+
+      try {
+        optionFileInputStream = Some(new FileInputStream(jsonKeyFileLocation))
+
+        GoogleCredential.fromStream(optionFileInputStream.get).createScoped(SCOPES)
+      } finally {
+        if (optionFileInputStream.isDefined) {
+          optionFileInputStream.get.close()
+        }
+      }
+    } else {
+      GoogleCredential.getApplicationDefault.createScoped(SCOPES)
+    }
+  }
+
+  private val bigquery: Bigquery =
+    new Bigquery.Builder(new NetHttpTransport, new JacksonFactory, googleCredential)
       .setApplicationName("spark-bigquery")
       .build()
-  }
 
   private def projectId: String = conf.get(BigQueryConfiguration.PROJECT_ID_KEY)
 
-  private val queryCache: LoadingCache[String, TableReference] =
+  private case class CacheKey(sqlQuery: String, useStandardSql: Boolean)
+
+  private val queryCache: LoadingCache[CacheKey, TableReference] =
     CacheBuilder.newBuilder()
       .expireAfterWrite(STAGING_DATASET_TABLE_EXPIRATION_MS, TimeUnit.MILLISECONDS)
-      .build[String, TableReference](new CacheLoader[String, TableReference] {
-      override def load(key: String): TableReference = {
-        val sqlQuery = key
+      .build[CacheKey, TableReference](new CacheLoader[CacheKey, TableReference] {
+      override def load(key: CacheKey): TableReference = {
+        val sqlQuery = key.sqlQuery
         logger.info(s"Executing query $sqlQuery")
 
         val location = conf.get(STAGING_DATASET_LOCATION, STAGING_DATASET_LOCATION_DEFAULT)
         val destinationTable = temporaryTable(location)
+        val tableName = BigQueryStrings.toString(destinationTable)
         logger.info(s"Destination table: $destinationTable")
 
-        val job = createQueryJob(sqlQuery, destinationTable, dryRun = false)
+        val job = createQueryJob(sqlQuery, destinationTable,
+          dryRun = false, useLegacySql = !key.useStandardSql)
         waitForJob(job)
         destinationTable
       }
     })
 
   private def inConsole = Thread.currentThread().getStackTrace.exists(
     _.getClassName.startsWith("scala.tools.nsc.interpreter."))
-  private val DEFAULT_PRIORITY = if (inConsole) JOB_PRIORITY_INTERACTIVE else JOB_PRIORITY_BATCH
+
+  private val PRIORITY = if (inConsole) "INTERACTIVE" else "BATCH"
   private val TABLE_ID_PREFIX = "spark_bigquery"
   private val JOB_ID_PREFIX = "spark_bigquery"
   private val TIME_FORMATTER = DateTimeFormat.forPattern("yyyyMMddHHmmss")
 
   /**
-   * Perform a BigQuery SELECT query and save results to a temporary table.
-   */
-  def query(sqlQuery: String): TableReference = queryCache.get(sqlQuery)
+    * Perform a BigQuery SELECT query and save results to a temporary table.
+    */
+  def query(sqlQuery: String, useStandardSql: Boolean = false): TableReference =
+    queryCache.get(CacheKey(sqlQuery, useStandardSql))
 
   /**
-   * Load an Avro data set on GCS to a BigQuery table.
-   */
+    * Load an Avro data set on GCS to a BigQuery table.
+    */
   def load(gcsPath: String, destinationTable: TableReference,
            writeDisposition: WriteDisposition.Value = null,
            createDisposition: CreateDisposition.Value = null): Unit = {
@@ -125,10 +150,6 @@ private[bigquery] class BigQueryClient(conf: Configuration) {
     if (createDisposition != null) {
       loadConfig = loadConfig.setCreateDisposition(createDisposition.toString)
     }
-    if (destinationTable.getTableId.contains("$")) {
-      val timePartitioning = new TimePartitioning().setType("DAY")
-      loadConfig = loadConfig.setTimePartitioning(timePartitioning)
-    }
 
     val jobConfig = new JobConfiguration().setLoad(loadConfig)
     val jobReference = createJobReference(projectId, JOB_ID_PREFIX)
@@ -146,8 +167,7 @@ private[bigquery] class BigQueryClient(conf: Configuration) {
   private def stagingDataset(location: String): DatasetReference = {
     // Create staging dataset if it does not already exist
     val prefix = conf.get(STAGING_DATASET_PREFIX, STAGING_DATASET_PREFIX_DEFAULT)
-    conf.setIfUnset(STAGING_DATASET_ID, prefix + location.toLowerCase)
-    val datasetId = conf.get(STAGING_DATASET_ID)
+    val datasetId = prefix + location.toLowerCase
     try {
       bigquery.datasets().get(projectId, datasetId).execute()
       logger.info(s"Staging dataset $projectId:$datasetId already exists")
@@ -181,12 +201,10 @@ private[bigquery] class BigQueryClient(conf: Configuration) {
 
   private def createQueryJob(sqlQuery: String,
                              destinationTable: TableReference,
-                             dryRun: Boolean): Job = {
-
-    val priority = conf.get(QUERY_JOB_PRIORITY, DEFAULT_PRIORITY)
-    var queryConfig = new JobConfigurationQuery()
+                             dryRun: Boolean, useLegacySql: Boolean): Job = {
+    var queryConfig = new JobConfigurationQuery().setUseLegacySql(useLegacySql)
       .setQuery(sqlQuery)
-      .setPriority(priority)
+      .setPriority(PRIORITY)
       .setCreateDisposition("CREATE_IF_NEEDED")
       .setWriteDisposition("WRITE_EMPTY")
     if (destinationTable != null) {

src/main/scala/com/spotify/spark/bigquery/BigQueryDataFrame.scala

@@ -47,16 +47,11 @@ class BigQueryDataFrame(df: DataFrame) {
     */
   def saveAsBigQueryTable(tableRef: TableReference,
                           writeDisposition: WriteDisposition.Value,
-                          createDisposition: CreateDisposition.Value,
-                          tmpWriteOptions: Map[String, String]): Unit = {
+                          createDisposition: CreateDisposition.Value): Unit = {
     val bucket = conf.get(BigQueryConfiguration.GCS_BUCKET_KEY)
     val temp = s"spark-bigquery-${System.currentTimeMillis()}=${Random.nextInt(Int.MaxValue)}"
     val gcsPath = s"gs://$bucket/hadoop/tmp/spark-bigquery/$temp"
-
-    tmpWriteOptions match {
-      case null => df.write.format("avro").save(gcsPath)
-      case _ => df.write.options(tmpWriteOptions).format("avro").save(gcsPath)
-    }
+    df.write.avro(gcsPath)
 
     val fdf = bq.load(gcsPath, tableRef, writeDisposition, createDisposition)
     delete(new Path(gcsPath))
@@ -68,13 +63,11 @@ class BigQueryDataFrame(df: DataFrame) {
     */
   def saveAsBigQueryTable(tableSpec: String,
                           writeDisposition: WriteDisposition.Value = null,
-                          createDisposition: CreateDisposition.Value = null,
-                          tmpWriteOptions: Map[String,String] = null): Unit =
+                          createDisposition: CreateDisposition.Value = null): Unit =
     saveAsBigQueryTable(
       BigQueryStrings.parseTableReference(tableSpec),
       writeDisposition,
-      createDisposition,
-      tmpWriteOptions)
+      createDisposition)
 
   private def delete(path: Path): Unit = {
     val fs = FileSystem.get(path.toUri, conf)

src/main/scala/com/spotify/spark/bigquery/BigQuerySQLContext.scala

@@ -32,12 +32,13 @@ import org.apache.spark.sql.{DataFrame, Row, SQLContext}
 
 /**
   * Easier class to expose to python and java
+  *
   * @param sqlContext SQLContext
   */
 class BigQuerySQLContext(sqlContext: SQLContext) {
 
   val sc: SparkContext = sqlContext.sparkContext
-  val conf: Configuration = new Configuration(sc.hadoopConfiguration)
+  val conf: Configuration = sc.hadoopConfiguration
   lazy val bq: BigQueryClient = BigQueryClient.getInstance(conf)
 
   // Register GCS implementation
@@ -69,18 +70,6 @@ class BigQuerySQLContext(sqlContext: SQLContext) {
   def setBigQueryDatasetLocation(location: String): Unit =
     conf.set(BigQueryClient.STAGING_DATASET_LOCATION, location)
 
-  /**
-    * Set BigQuery query job priority.
-    * @param interactive Whether to execute query jobs as interactive
-    *                    or batch queries.
-    * @see https://cloud.google.com/bigquery/docs/query-overview
-    */
-  def setBigQueryJobPriority(interactive: Boolean): Unit = {
-    import BigQueryClient._
-    val priority = if(interactive) JOB_PRIORITY_INTERACTIVE else JOB_PRIORITY_BATCH
-    conf.set(QUERY_JOB_PRIORITY, priority)
-  }
-
   /**
     * Set GCP JSON key file.
     */
@@ -100,9 +89,12 @@ class BigQuerySQLContext(sqlContext: SQLContext) {
 
   /**
     * Perform a BigQuery SELECT query and load results as a [[DataFrame]].
-    * @param sqlQuery SQL query in SQL-2011 dialect.
+    *
+    * @param sqlQuery       SQL query in SQL-2011 dialect.
+    * @param useStandardSql Set for Standard SQL dialect.
     */
-  def bigQuerySelect(sqlQuery: String): DataFrame = bigQueryTable(bq.query(sqlQuery))
+  def bigQuerySelect(sqlQuery: String, useStandardSql: Boolean = false): DataFrame =
+    bigQueryTable(bq.query(sqlQuery, useStandardSql = useStandardSql))
 
   /**
     * Load a BigQuery table as a [[DataFrame]].
@@ -115,9 +107,11 @@ class BigQuerySQLContext(sqlContext: SQLContext) {
     BigQueryConfiguration.configureBigQueryInput(
       conf, tableRef.getProjectId, tableRef.getDatasetId, tableRef.getTableId)
 
+    val fClass = classOf[AvroBigQueryInputFormat]
+    val kClass = classOf[LongWritable]
+    val vClass = classOf[GenericData.Record]
     val rdd = sc
-      .newAPIHadoopRDD(conf, classOf[AvroBigQueryInputFormat], classOf[LongWritable],
-        classOf[GenericData.Record])
+      .newAPIHadoopRDD(conf, fClass, kClass, vClass)
       .map(_._2)
     val schemaString = rdd.map(_.getSchema.toString).first()
     val schema = new Schema.Parser().parse(schemaString)