diff --git a/.travis.yml b/.travis.yml index 9723317..6248ec9 100644 --- a/.travis.yml +++ b/.travis.yml @@ -1,8 +1,8 @@ language: scala scala: - 2.10.6 - - 2.11.8 -jdk: oraclejdk7 + - 2.11.12 +jdk: oraclejdk8 cache: directories: diff --git a/README.md b/README.md index 4ecea51..ffb36f8 100644 --- a/README.md +++ b/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. diff --git a/build.sbt b/build.sbt index dc68e23..a676b80 100644 --- a/build.sbt +++ b/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 := { diff --git a/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala b/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala index 880e08b..5ee8246 100644 --- a/src/main/scala/com/spotify/spark/bigquery/BigQueryClient.scala +++ b/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,28 +67,52 @@ 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 } @@ -97,19 +120,21 @@ private[bigquery] class BigQueryClient(conf: Configuration) { 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) { diff --git a/src/main/scala/com/spotify/spark/bigquery/BigQueryDataFrame.scala b/src/main/scala/com/spotify/spark/bigquery/BigQueryDataFrame.scala index 89e2324..70d500b 100644 --- a/src/main/scala/com/spotify/spark/bigquery/BigQueryDataFrame.scala +++ b/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) diff --git a/src/main/scala/com/spotify/spark/bigquery/BigQuerySQLContext.scala b/src/main/scala/com/spotify/spark/bigquery/BigQuerySQLContext.scala index 3dbd1ce..9429fe5 100644 --- a/src/main/scala/com/spotify/spark/bigquery/BigQuerySQLContext.scala +++ b/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)