Skip to content

Commit

Permalink
[gcp] Fix BQ BIGNUMERIC type support
Browse files Browse the repository at this point in the history
  • Loading branch information
RustedBones committed Jul 11, 2024
1 parent b52036a commit e09ffa0
Show file tree
Hide file tree
Showing 10 changed files with 193 additions and 105 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import com.google.api.client.googleapis.json.GoogleJsonResponseException
import com.google.api.services.bigquery.model.{Dataset, DatasetReference}
import com.google.protobuf.ByteString
import com.spotify.scio.bigquery.client.BigQuery
import com.spotify.scio.bigquery.types.{Geography, Json}
import com.spotify.scio.bigquery.types.{BigNumeric, Geography, Json}
import org.joda.time._
import org.joda.time.format.DateTimeFormat
import org.slf4j.LoggerFactory
Expand Down Expand Up @@ -53,7 +53,8 @@ object PopulateTestData {
time: LocalTime,
datetime: LocalDateTime,
geography: Geography,
json: Json
json: Json,
bigNumeric: BigNumeric
)

@BigQueryType.toTable
Expand All @@ -69,7 +70,8 @@ object PopulateTestData {
time: Option[LocalTime],
datetime: Option[LocalDateTime],
geography: Option[Geography],
json: Option[Json]
json: Option[Json],
bigNumeric: Option[BigNumeric]
)

@BigQueryType.toTable
Expand All @@ -85,7 +87,8 @@ object PopulateTestData {
time: List[LocalTime],
datetime: List[LocalDateTime],
geography: List[Geography],
json: List[Json]
json: List[Json],
bigNumeric: List[BigNumeric]
)

case class Record(int: Long, string: String)
Expand Down Expand Up @@ -209,7 +212,8 @@ object PopulateTestData {
dt.toLocalTime.plusMillis(i),
dt.toLocalDateTime.plusMillis(i),
Geography(s"POINT($i $i)"),
Json(s"""{"value": $i}""")
Json(s"""{"value": $i}"""),
BigNumeric(BigDecimal(i))
)
}

Expand All @@ -228,7 +232,8 @@ object PopulateTestData {
Some(dt.toLocalTime.plusMillis(i)),
Some(dt.toLocalDateTime.plusMillis(i)),
Some(Geography(s"POINT($i $i)")),
Some(Json(s"""{"value": $i}"""))
Some(Json(s"""{"value": $i}""")),
Some(BigNumeric(BigDecimal(i)))
)
}

Expand All @@ -247,7 +252,8 @@ object PopulateTestData {
List(dt.toLocalTime.plusMillis(i)),
List(dt.toLocalDateTime.plusMillis(i)),
List(Geography(s"POINT($i $i)")),
List(Json(s"""{"value": $i}"""))
List(Json(s"""{"value": $i}""")),
List(BigNumeric(BigDecimal(i)))
)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,9 @@ class BigQueryStorageIT extends AnyFlatSpec with Matchers {
dt.toLocalDate.plusDays(i),
dt.toLocalTime.plusMillis(i),
dt.toLocalDateTime.plusMillis(i),
s"POINT($i $i)", // geography is not an avro logical type
s"""{"value":$i}""" // json is not an avro logical type
Geography(s"POINT($i $i)"),
Json(s"""{"value":$i}"""),
BigNumeric(BigDecimal(i))
)
}.asJava
val (sc, _) = ContextAndArgs(
Expand All @@ -77,8 +78,9 @@ class BigQueryStorageIT extends AnyFlatSpec with Matchers {
Some(dt.toLocalDate.plusDays(i)),
Some(dt.toLocalTime.plusMillis(i)),
Some(dt.toLocalDateTime.plusMillis(i)),
Some(s"POINT($i $i)"), // geography is not an avro logical type
Some(s"""{"value":$i}""") // json is not an avro logical type
Some(Geography(s"POINT($i $i)")),
Some(Json(s"""{"value":$i}""")),
Some(BigNumeric(BigDecimal(i)))
)
}.asJava
val (sc, _) = ContextAndArgs(
Expand All @@ -104,8 +106,9 @@ class BigQueryStorageIT extends AnyFlatSpec with Matchers {
List(dt.toLocalDate.plusDays(i)),
List(dt.toLocalTime.plusMillis(i)),
List(dt.toLocalDateTime.plusMillis(i)),
List(s"POINT($i $i)"), // geography is not an avro logical type
List(s"""{"value":$i}""") // json is not an avro logical type
List(Geography(s"POINT($i $i)")),
List(Json(s"""{"value":$i}""")),
List(BigNumeric(BigDecimal(i)))
)
}.asJava
val (sc, _) = ContextAndArgs(
Expand Down Expand Up @@ -192,7 +195,8 @@ class BigQueryStorageIT extends AnyFlatSpec with Matchers {
Some(dt.toLocalTime.plusMillis(i)),
Some(dt.toLocalDateTime.plusMillis(i)),
Some(Geography(s"POINT($i $i)")),
Some(Json(s"""{"value":$i}"""))
Some(Json(s"""{"value":$i}""")),
Some(BigNumeric(BigDecimal(i)))
)
}.asJava
val (sc, _) = ContextAndArgs(
Expand Down Expand Up @@ -221,7 +225,8 @@ class BigQueryStorageIT extends AnyFlatSpec with Matchers {
dt.toLocalTime.plusMillis(i),
dt.toLocalDateTime.plusMillis(i),
Geography(s"POINT($i $i)"),
Json(s"""{"value":$i}""")
Json(s"""{"value":$i}"""),
BigNumeric(BigDecimal(i))
)
}.asJava
val (sc, _) = ContextAndArgs(
Expand Down Expand Up @@ -260,7 +265,8 @@ class BigQueryStorageIT extends AnyFlatSpec with Matchers {
Some(dt.toLocalTime.plusMillis(i)),
Some(dt.toLocalDateTime.plusMillis(i)),
Some(Geography(s"POINT($i $i)")),
Some(Json(s"""{"value":$i}"""))
Some(Json(s"""{"value":$i}""")),
Some(BigNumeric(BigDecimal(i)))
)
}.asJava
val (sc, _) = ContextAndArgs(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -300,13 +300,12 @@ object Numeric {
val MaxNumericPrecision = 38
val MaxNumericScale = 9

private[this] val DecimalConverter = new DecimalConversion
private[this] val DecimalLogicalType = LogicalTypes.decimal(MaxNumericPrecision, MaxNumericScale)
private val DecimalConverter = new DecimalConversion
private val DecimalLogicalType = LogicalTypes.decimal(MaxNumericPrecision, MaxNumericScale)

def apply(value: String): BigDecimal = apply(BigDecimal(value))

def apply(value: BigDecimal): BigDecimal = {
// NUMERIC's max scale is 9, precision is 38
val scaled = if (value.scale > MaxNumericScale) {
value.setScale(MaxNumericScale, scala.math.BigDecimal.RoundingMode.HALF_UP)
} else {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -99,17 +99,13 @@ object StorageUtil {
}
case Type.STRING =>
// FIXME: schema.getLogicalType == null in this case, BigQuery service side bug?
if (schema.getProp("logicalType") == "datetime") {
"DATETIME"
} else {
schema.getLogicalType match {
case null => "STRING"
case t if t.getName == "datetime" => "DATETIME"
case t if t.getName == "geography" => "GEOGRAPHY"
case t if t.getName == "json" => "JSON"
case t =>
throw new IllegalStateException(s"Unsupported logical type: $t")
}
val logicalType = schema.getProp("logicalType")
val sqlType = schema.getProp("sqlType")
(logicalType, sqlType) match {
case ("datetime", _) => "DATETIME"
case (_, "GEOGRAPHY") => "GEOGRAPHY"
case (_, "JSON") => "JSON"
case _ => "STRING"
}
case Type.RECORD =>
tableField.setFields(getFieldSchemas(schema).asJava)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,6 @@

package com.spotify.scio.bigquery.types

class BigQueryTag extends scala.annotation.StaticAnnotation with Serializable
import scala.annotation.StaticAnnotation

class BigQueryTag extends StaticAnnotation with Serializable
Original file line number Diff line number Diff line change
Expand Up @@ -74,22 +74,26 @@ private[types] object ConverterProvider {
tpe match {
case t if provider.shouldOverrideType(c)(t) =>
provider.createInstance(c)(t, q"$tree")
case t if t =:= typeOf[Boolean] => q"$tree.asInstanceOf[Boolean]"
case t if t =:= typeOf[Int] => q"$tree.asInstanceOf[Long].toInt"
case t if t =:= typeOf[Long] => q"$tree.asInstanceOf[Long]"
case t if t =:= typeOf[Float] => q"$tree.asInstanceOf[Double].toFloat"
case t if t =:= typeOf[Double] => q"$tree.asInstanceOf[Double]"
case t if t =:= typeOf[String] => q"$tree.toString"
case t if t =:= typeOf[Boolean] =>
q"$tree.asInstanceOf[Boolean]"
case t if t =:= typeOf[Int] =>
q"$tree.asInstanceOf[Long].toInt"
case t if t =:= typeOf[Long] =>
q"$tree.asInstanceOf[Long]"
case t if t =:= typeOf[Float] =>
q"$tree.asInstanceOf[Double].toFloat"
case t if t =:= typeOf[Double] =>
q"$tree.asInstanceOf[Double]"
case t if t =:= typeOf[String] =>
q"$tree.toString"
case t if t =:= typeOf[BigDecimal] =>
q"_root_.com.spotify.scio.bigquery.Numeric.parse($tree)"

case t if t =:= typeOf[ByteString] =>
val b = q"$tree.asInstanceOf[_root_.java.nio.ByteBuffer]"
q"_root_.com.google.protobuf.ByteString.copyFrom($b.asReadOnlyBuffer())"
case t if t =:= typeOf[Array[Byte]] =>
val b = q"$tree.asInstanceOf[_root_.java.nio.ByteBuffer]"
q"_root_.java.util.Arrays.copyOfRange($b.array(), $b.position(), $b.limit())"

case t if t =:= typeOf[Instant] =>
q"_root_.com.spotify.scio.bigquery.Timestamp.parse($tree)"
case t if t =:= typeOf[LocalDate] =>
Expand All @@ -102,12 +106,13 @@ private[types] object ConverterProvider {
q"_root_.com.spotify.scio.bigquery.types.Geography($tree.toString)"
case t if t =:= typeOf[Json] =>
q"_root_.com.spotify.scio.bigquery.types.Json($tree.toString)"

case t if t =:= typeOf[BigNumeric] =>
q"_root_.com.spotify.scio.bigquery.types.BigNumeric.parse($tree)"
case t if isCaseClass(c)(t) =>
val fn = TermName("r" + t.typeSymbol.name)
val nestedRecord = TermName("r" + t.typeSymbol.name)
q"""{
val $fn = $tree.asInstanceOf[_root_.org.apache.avro.generic.GenericRecord]
${constructor(t, fn)}
val $nestedRecord = $tree.asInstanceOf[_root_.org.apache.avro.generic.GenericRecord]
${constructor(t, Ident(nestedRecord))}
}
"""
case _ => c.abort(c.enclosingPosition, s"Unsupported type: $tpe")
Expand All @@ -122,11 +127,11 @@ private[types] object ConverterProvider {
q"asScala($tree.asInstanceOf[$jl]).iterator.map(x => ${cast(q"x", tpe)}).toList"
}

def field(symbol: Symbol, fn: TermName): Tree = {
def field(symbol: Symbol, record: Tree): Tree = {
val name = symbol.name.toString
val tpe = symbol.asMethod.returnType

val tree = q"$fn.get($name)"
val tree = q"$record.get($name)"
if (tpe.erasure =:= typeOf[Option[_]].erasure) {
option(tree, tpe.typeArgs.head)
} else if (tpe.erasure =:= typeOf[List[_]].erasure) {
Expand All @@ -136,10 +141,10 @@ private[types] object ConverterProvider {
}
}

def constructor(tpe: Type, fn: TermName): Tree = {
def constructor(tpe: Type, record: Tree): Tree = {
val companion = tpe.typeSymbol.companion
val gets = tpe.erasure match {
case t if isCaseClass(c)(t) => getFields(c)(t).map(s => field(s, fn))
case t if isCaseClass(c)(t) => getFields(c)(t).map(s => field(s, record))
case _ => c.abort(c.enclosingPosition, s"Unsupported type: $tpe")
}
q"$companion(..$gets)"
Expand All @@ -148,11 +153,9 @@ private[types] object ConverterProvider {
// =======================================================================
// Entry point
// =======================================================================

val tn = TermName("r")
q"""(r: _root_.org.apache.avro.generic.GenericRecord) => {
import _root_.scala.jdk.javaapi.CollectionConverters._
${constructor(tpe, tn)}
${constructor(tpe, q"r")}
}
"""
}
Expand All @@ -178,8 +181,10 @@ private[types] object ConverterProvider {

case t if t =:= typeOf[BigDecimal] =>
q"_root_.com.spotify.scio.bigquery.Numeric($tree).toString"
case t if t =:= typeOf[ByteString] => q"_root_.java.nio.ByteBuffer.wrap($tree.toByteArray)"
case t if t =:= typeOf[Array[Byte]] => q"_root_.java.nio.ByteBuffer.wrap($tree)"
case t if t =:= typeOf[ByteString] =>
q"_root_.java.nio.ByteBuffer.wrap($tree.toByteArray)"
case t if t =:= typeOf[Array[Byte]] =>
q"_root_.java.nio.ByteBuffer.wrap($tree)"

case t if t =:= typeOf[Instant] => q"$tree.getMillis * 1000"
case t if t =:= typeOf[LocalDate] =>
Expand All @@ -194,8 +199,11 @@ private[types] object ConverterProvider {
q"$tree.wkt"
case t if t =:= typeOf[Json] =>
q"$tree.wkt"
case t if t =:= typeOf[BigNumeric] =>
q"_root_.com.spotify.scio.bigquery.types.BigNumeric($tree.wkt).toString"

case t if isCaseClass(c)(t) => // nested records
// nested records
case t if isCaseClass(c)(t) =>
val fn = TermName("r" + t.typeSymbol.name)
q"""{
val $fn = $tree
Expand Down Expand Up @@ -298,6 +306,8 @@ private[types] object ConverterProvider {
q"_root_.com.spotify.scio.bigquery.types.Geography($s)"
case t if t =:= typeOf[Json] =>
q"_root_.com.spotify.scio.bigquery.types.Json($s)"
case t if t =:= typeOf[BigNumeric] =>
q"_root_.com.spotify.scio.bigquery.types.BigNumeric($s)"

case t if isCaseClass(c)(t) => // nested records
val fn = TermName("r" + t.typeSymbol.name)
Expand Down Expand Up @@ -403,6 +413,8 @@ private[types] object ConverterProvider {
q"$tree.wkt"
case t if t =:= typeOf[Json] =>
q"$tree.wkt"
case t if t =:= typeOf[BigNumeric] =>
q"$tree.wkt"

case t if isCaseClass(c)(t) => // nested records
val fn = TermName("r" + t.typeSymbol.name)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,7 @@ private[types] object SchemaProvider {
case t if t =:= typeOf[Double] => ("FLOAT", Iterable.empty)
case t if t =:= typeOf[String] => ("STRING", Iterable.empty)
case t if t =:= typeOf[BigDecimal] => ("NUMERIC", Iterable.empty)
case t if t =:= typeOf[BigNumeric] => ("BIGNUMERIC", Iterable.empty)

case t if t =:= typeOf[ByteString] => ("BYTES", Iterable.empty)
case t if t =:= typeOf[Array[Byte]] => ("BYTES", Iterable.empty)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -271,14 +271,14 @@ private[types] object TypeProvider {
case "FLOAT" | "FLOAT64" => (tq"_root_.scala.Double", Nil)
case "STRING" => (tq"_root_.java.lang.String", Nil)
case "NUMERIC" => (tq"_root_.scala.BigDecimal", Nil)
case "BIGNUMERIC" => (tq"_root_.scala.BigDecimal", Nil)
case "BYTES" => (tq"_root_.com.google.protobuf.ByteString", Nil)
case "TIMESTAMP" => (tq"_root_.org.joda.time.Instant", Nil)
case "DATE" => (tq"_root_.org.joda.time.LocalDate", Nil)
case "TIME" => (tq"_root_.org.joda.time.LocalTime", Nil)
case "DATETIME" => (tq"_root_.org.joda.time.LocalDateTime", Nil)
case "GEOGRAPHY" => (tq"_root_.com.spotify.scio.bigquery.types.Geography", Nil)
case "JSON" => (tq"_root_.com.spotify.scio.bigquery.types.Json", Nil)
case "BIGNUMERIC" => (tq"_root_.com.spotify.scio.bigquery.types.BigNumeric", Nil)
case "RECORD" | "STRUCT" =>
val name = NameProvider.getUniqueName(tfs.getName)
val (fields, records) = toFields(tfs.getFields)
Expand Down
Loading

0 comments on commit e09ffa0

Please sign in to comment.