Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[gcp] Fix BQ BIGNUMERIC type support #5428

Merged
merged 1 commit into from
Jul 11, 2024
Merged

[gcp] Fix BQ BIGNUMERIC type support #5428

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 13 additions & 7 deletions integration/src/main/scala/com/spotify/scio/bigquery/PopulateTestData.scala
View file
Open in desktop
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)))
)
}
}
24 changes: 15 additions & 9 deletions integration/src/test/scala/com/spotify/scio/bigquery/types/BigQueryStorageIT.scala
View file
Open in desktop
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
Comment on lines -53 to -54
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I found out that BQ adds custom sqlType property to the schema. We can reconstruct the original 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
5 changes: 2 additions & 3 deletions scio-google-cloud-platform/src/main/scala/com/spotify/scio/bigquery/BigQueryTypes.scala
View file
Open in desktop
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
18 changes: 7 additions & 11 deletions scio-google-cloud-platform/src/main/scala/com/spotify/scio/bigquery/StorageUtil.scala
View file
Open in desktop
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
4 changes: 3 additions & 1 deletion scio-google-cloud-platform/src/main/scala/com/spotify/scio/bigquery/types/BigQueryTag.scala
View file
Open in desktop
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
56 changes: 34 additions & 22 deletions ...gle-cloud-platform/src/main/scala/com/spotify/scio/bigquery/types/ConverterProvider.scala
View file
Open in desktop
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)}
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is technically an Ident(TermName("r")), adapt macro

${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
1 change: 1 addition & 0 deletions ...google-cloud-platform/src/main/scala/com/spotify/scio/bigquery/types/SchemaProvider.scala
View file
Open in desktop
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
2 changes: 1 addition & 1 deletion scio-google-cloud-platform/src/main/scala/com/spotify/scio/bigquery/types/TypeProvider.scala
View file
Open in desktop
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