-
Notifications
You must be signed in to change notification settings - Fork 513
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Provide module specific kryo coder registrar (#4753)
* Provide module specific kryo coder registrar * Add back avro dep in core
- Loading branch information
1 parent
924282d
commit 15dba83
Showing
10 changed files
with
285 additions
and
103 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
31 changes: 31 additions & 0 deletions
31
scio-avro/src/main/scala/com/spotify/scio/coders/AvroKryoRegistrar.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,31 @@ | ||
| /* | ||
| * Copyright 2023 Spotify AB | ||
| * | ||
| * 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.spotify.scio.coders | ||
|
|
||
| import com.spotify.scio.coders.instances.kryo.{GenericAvroSerializer, SpecificAvroSerializer} | ||
| import com.twitter.chill._ | ||
| import org.apache.avro.generic.GenericRecord | ||
| import org.apache.avro.specific.SpecificRecord | ||
|
|
||
| @KryoRegistrar | ||
| class AvroKryoRegistrar extends IKryoRegistrar { | ||
| override def apply(k: Kryo): Unit = { | ||
| k.forSubclass[SpecificRecord](new SpecificAvroSerializer) | ||
| k.forSubclass[GenericRecord](new GenericAvroSerializer) | ||
| } | ||
|
|
||
| } |
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
62 changes: 62 additions & 0 deletions
62
scio-core/src/main/scala/com/spotify/scio/coders/instances/kryo/GrpcSerializer.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| /* | ||
| * Copyright 2020 Spotify AB. | ||
| * | ||
| * 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.spotify.scio.coders.instances.kryo | ||
|
|
||
| import com.esotericsoftware.kryo.Kryo | ||
| import com.esotericsoftware.kryo.io.{Input, Output} | ||
| import com.twitter.chill.KSerializer | ||
| import io.grpc.{Metadata, Status, StatusRuntimeException} | ||
|
|
||
| private[coders] class StatusSerializer extends KSerializer[Status] { | ||
| override def write(kryo: Kryo, output: Output, status: Status): Unit = { | ||
| output.writeInt(status.getCode().value()) | ||
| output.writeString(status.getDescription) | ||
| kryo.writeClassAndObject(output, status.getCause) | ||
| } | ||
|
|
||
| override def read(kryo: Kryo, input: Input, `type`: Class[Status]): Status = { | ||
| val code = input.readInt() | ||
| val description = input.readString() | ||
| val cause = kryo.readClassAndObject(input).asInstanceOf[Throwable] | ||
|
|
||
| Status | ||
| .fromCodeValue(code) | ||
| .withDescription(description) | ||
| .withCause(cause) | ||
| } | ||
| } | ||
|
|
||
| private[coders] class StatusRuntimeExceptionSerializer extends KSerializer[StatusRuntimeException] { | ||
| lazy val statusSer = new StatusSerializer() | ||
|
|
||
| override def write(kryo: Kryo, output: Output, e: StatusRuntimeException): Unit = { | ||
| kryo.writeObject(output, e.getStatus, statusSer) | ||
| kryo.writeObjectOrNull(output, e.getTrailers, classOf[Metadata]) | ||
| } | ||
|
|
||
| override def read( | ||
| kryo: Kryo, | ||
| input: Input, | ||
| `type`: Class[StatusRuntimeException] | ||
| ): StatusRuntimeException = { | ||
| val status = kryo.readObject(input, classOf[Status], statusSer) | ||
| val trailers = kryo.readObjectOrNull(input, classOf[Metadata]) | ||
|
|
||
| new StatusRuntimeException(status, trailers) | ||
| } | ||
| } |
65 changes: 0 additions & 65 deletions
65
scio-core/src/main/scala/com/spotify/scio/coders/instances/kryo/GrpcSerializers.scala
This file was deleted.
Oops, something went wrong.
32 changes: 32 additions & 0 deletions
32
scio-google-cloud-platform/src/main/scala/com/spotify/scio/coders/GcpKryoRegistrar.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,32 @@ | ||
| /* | ||
| * Copyright 2023 Spotify AB | ||
| * | ||
| * 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.spotify.scio.coders | ||
|
|
||
| import com.google.cloud.bigtable.grpc.scanner.BigtableRetriesExhaustedException | ||
| import com.spotify.scio.bigquery.TableRow | ||
| import com.spotify.scio.coders.instances.kryo.CoderSerializer | ||
| import com.spotify.scio.coders.instances.kryo.BigtableRetriesExhaustedExceptionSerializer | ||
| import com.twitter.chill._ | ||
| import org.apache.beam.sdk.io.gcp.bigquery.TableRowJsonCoder | ||
|
|
||
| @KryoRegistrar | ||
| class GcpKryoRegistrar extends IKryoRegistrar { | ||
| override def apply(k: Kryo): Unit = { | ||
| k.forClass[TableRow](new CoderSerializer(TableRowJsonCoder.of())) | ||
| k.forClass[BigtableRetriesExhaustedException](new BigtableRetriesExhaustedExceptionSerializer) | ||
| } | ||
| } |
43 changes: 43 additions & 0 deletions
43
...-cloud-platform/src/main/scala/com/spotify/scio/coders/instances/kryo/GcpSerializer.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,43 @@ | ||
| /* | ||
| * Copyright 2023 Spotify AB | ||
| * | ||
| * 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.spotify.scio.coders.instances.kryo | ||
|
|
||
| import com.esotericsoftware.kryo.serializers.DefaultSerializers.StringSerializer | ||
| import com.google.cloud.bigtable.grpc.scanner.BigtableRetriesExhaustedException | ||
| import com.twitter.chill._ | ||
|
|
||
| private[coders] class BigtableRetriesExhaustedExceptionSerializer | ||
| extends KSerializer[BigtableRetriesExhaustedException] { | ||
|
|
||
| private lazy val stringSerializer = new StringSerializer() | ||
| private lazy val statusExceptionSerializer = new StatusRuntimeExceptionSerializer() | ||
|
|
||
| override def write(kryo: Kryo, output: Output, e: BigtableRetriesExhaustedException): Unit = { | ||
| kryo.writeObject(output, e.getMessage, stringSerializer) | ||
| kryo.writeObject(output, e.getCause, statusExceptionSerializer) | ||
| } | ||
|
|
||
| override def read( | ||
| kryo: Kryo, | ||
| input: Input, | ||
| `type`: Class[BigtableRetriesExhaustedException] | ||
| ): BigtableRetriesExhaustedException = { | ||
| val message = kryo.readObject(input, classOf[String], stringSerializer) | ||
| val cause = kryo.readObject(input, classOf[Throwable], statusExceptionSerializer) | ||
| new BigtableRetriesExhaustedException(message, cause) | ||
| } | ||
| } |
68 changes: 68 additions & 0 deletions
68
...oud-platform/src/test/scala/com/spotify/scio/coders/instance/kryo/GcpSerializerTest.scala
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,68 @@ | ||
| /* | ||
| * Copyright 2023 Spotify AB | ||
| * | ||
| * 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.spotify.scio.coders.instance.kryo | ||
|
|
||
| import com.google.cloud.bigtable.grpc.scanner.BigtableRetriesExhaustedException | ||
| import com.spotify.scio.coders.instances.kryo.GrpcSerializerTest.eqStatusRuntimeException | ||
| import com.spotify.scio.coders.{Coder, CoderMaterializer} | ||
| import io.grpc.{Metadata, Status, StatusRuntimeException} | ||
| import org.apache.beam.sdk.util.CoderUtils | ||
| import org.scalactic.Equality | ||
| import org.scalatest.flatspec.AnyFlatSpec | ||
| import org.scalatest.matchers.should.Matchers | ||
|
|
||
| object GcpSerializerTest { | ||
|
|
||
| implicit val eqBigtableRetriesExhaustedException: Equality[BigtableRetriesExhaustedException] = { | ||
| case (a: BigtableRetriesExhaustedException, b: BigtableRetriesExhaustedException) => | ||
| a.getMessage == b.getMessage && | ||
| ((Option(a.getCause), Option(b.getCause)) match { | ||
| case (None, None) => true | ||
| case (Some(ac: StatusRuntimeException), Some(bc: StatusRuntimeException)) => | ||
| eqStatusRuntimeException.areEqual(ac, bc) | ||
| case _ => | ||
| false | ||
| }) | ||
| case _ => false | ||
| } | ||
|
|
||
| } | ||
|
|
||
| class GcpSerializerTest extends AnyFlatSpec with Matchers { | ||
|
|
||
| import GcpSerializerTest._ | ||
|
|
||
| "GcpSerializer" should "roundtrip" in { | ||
| val metadata = new Metadata() | ||
| metadata.put(Metadata.Key.of[String]("k", Metadata.ASCII_STRING_MARSHALLER), "v") | ||
| val cause = new StatusRuntimeException( | ||
| Status.OK.withCause(new RuntimeException("bar")).withDescription("bar"), | ||
| metadata | ||
| ) | ||
| roundtrip(new BigtableRetriesExhaustedException("Error", cause)) | ||
| } | ||
|
|
||
| private def roundtrip(t: BigtableRetriesExhaustedException): Unit = { | ||
| val kryoBCoder = CoderMaterializer.beamWithDefault(Coder[BigtableRetriesExhaustedException]) | ||
|
|
||
| val bytes = CoderUtils.encodeToByteArray(kryoBCoder, t) | ||
| val copy = CoderUtils.decodeFromByteArray(kryoBCoder, bytes) | ||
|
|
||
| t shouldEqual copy | ||
| } | ||
|
|
||
| } |
Oops, something went wrong.