From ff303c6815d00448e91e591b31de5f0cb3a29659 Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Tue, 4 Aug 2026 20:49:51 +0100 Subject: [PATCH 1/3] perf issues in google-cloud-bigquery-storage --- .../bigquery/storage/impl/ArrowSource.scala | 23 ++++++++++++++----- .../bigquery/storage/impl/AvroSource.scala | 7 ++++-- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala index 786bd69c7..50dcf91f9 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala @@ -16,7 +16,7 @@ package org.apache.pekko.stream.connectors.googlecloud.bigquery.storage.impl import org.apache.pekko import pekko.NotUsed import pekko.stream.connectors.googlecloud.bigquery.storage.BigQueryRecord -import pekko.stream.scaladsl.Source +import pekko.stream.scaladsl.{ Merge, Source } import com.google.cloud.bigquery.storage.v1.arrow.{ ArrowRecordBatch, ArrowSchema } import com.google.cloud.bigquery.storage.v1.storage.BigQueryReadClient import com.google.cloud.bigquery.storage.v1.stream.ReadSession @@ -34,11 +34,17 @@ object ArrowSource { def readRecordsMerged(client: BigQueryReadClient, readSession: ReadSession): Source[List[BigQueryRecord], NotUsed] = readMerged(client, readSession) - .map(a => new SimpleRowReader(readSession.schema.arrowSchema.get).read(a)) + .map { a => + val reader = new SimpleRowReader(readSession.schema.arrowSchema.get) + try reader.read(a) + finally reader.close() + } def readMerged(client: BigQueryReadClient, session: ReadSession): Source[ArrowRecordBatch, NotUsed] = - read(client, session) - .reduce((a, b) => a.merge(b)) + read(client, session) match { + case Seq(single) => single + case sources => Source.combine(sources.head, sources.tail.head, sources.tail.tail: _*)(Merge(_)) + } def readRecords(client: BigQueryReadClient, session: ReadSession): Seq[Source[BigQueryRecord, NotUsed]] = read(client, session) @@ -56,9 +62,14 @@ object ArrowSource { } -final class SimpleRowReader(val schema: ArrowSchema) extends AutoCloseable { +object SimpleRowReader { + val DefaultAllocationLimit: Long = 256L * 1024 * 1024 // 256 MB +} + +final class SimpleRowReader(val schema: ArrowSchema, allocationLimit: Long = SimpleRowReader.DefaultAllocationLimit) + extends AutoCloseable { - val allocator = new RootAllocator(Long.MaxValue) + val allocator = new RootAllocator(allocationLimit) val sd = MessageSerializer.deserializeSchema( new ReadChannel( diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala index 47398cd47..2c4b15a33 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala @@ -16,7 +16,7 @@ package org.apache.pekko.stream.connectors.googlecloud.bigquery.storage.impl import org.apache.pekko import pekko.NotUsed import pekko.stream.connectors.googlecloud.bigquery.storage.BigQueryRecord -import pekko.stream.scaladsl.Source +import pekko.stream.scaladsl.{ Merge, Source } import com.google.cloud.bigquery.storage.v1.avro.AvroRows import com.google.cloud.bigquery.storage.v1.storage.BigQueryReadClient import com.google.cloud.bigquery.storage.v1.stream.ReadSession @@ -30,7 +30,10 @@ object AvroSource { } def readMerged(client: BigQueryReadClient, session: ReadSession): Source[AvroRows, NotUsed] = - read(client, session).reduce((a, b) => a.merge(b)) + read(client, session) match { + case Seq(single) => single + case sources => Source.combine(sources.head, sources.tail.head, sources.tail.tail: _*)(Merge(_)) + } def readRecords(client: BigQueryReadClient, session: ReadSession): Seq[Source[BigQueryRecord, NotUsed]] = read(client, session) From d6e9f197fb8953cc97db7457d1f4a502148b3a2e Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Tue, 4 Aug 2026 21:11:28 +0100 Subject: [PATCH 2/3] refactor --- .../storage/BigQueryStorageSettings.scala | 31 ++++++++++++++++--- .../bigquery/storage/impl/ArrowSource.scala | 27 +++++++--------- .../bigquery/storage/impl/AvroSource.scala | 7 ++--- .../scaladsl/BigQueryArrowStorage.scala | 16 +++++----- .../scaladsl/GrpcBigQueryStorageReader.scala | 3 +- .../scaladsl/BigQueryArrowStorageSpec.scala | 5 +-- .../scaladsl/BigQueryStorageSpec.scala | 3 +- 7 files changed, 55 insertions(+), 37 deletions(-) diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/BigQueryStorageSettings.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/BigQueryStorageSettings.scala index d4a699221..2582e0a8a 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/BigQueryStorageSettings.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/BigQueryStorageSettings.scala @@ -19,7 +19,8 @@ import com.typesafe.config.Config final class BigQueryStorageSettings private ( val host: String, val port: Int, - val rootCa: Option[String] = None) { + val rootCa: Option[String] = None, + val arrowAllocatorBytes: Long = BigQueryStorageSettings.DefaultArrowAllocatorBytes) { /** * Endpoint hostname where the gRPC connection is made. @@ -38,19 +39,34 @@ final class BigQueryStorageSettings private ( def withRootCa(rootCa: String): BigQueryStorageSettings = copy(rootCa = Some(rootCa)) - private def copy(host: String = host, port: Int = port, rootCa: Option[String] = rootCa) = - new BigQueryStorageSettings(host, port, rootCa) + /** + * Maximum bytes the Arrow root allocator may reserve per batch. + * The allocator is created per batch and closed after reading, so this bounds + * native memory for a single batch rather than the lifetime of the stream. + */ + def withArrowAllocatorBytes(bytes: Long): BigQueryStorageSettings = + copy(arrowAllocatorBytes = bytes) + + private def copy( + host: String = host, + port: Int = port, + rootCa: Option[String] = rootCa, + arrowAllocatorBytes: Long = arrowAllocatorBytes) = + new BigQueryStorageSettings(host, port, rootCa, arrowAllocatorBytes) override def toString: String = "BigQueryStorageSettings(" + s"host=$host, " + s"port=$port, " + - s"rootCa=$rootCa" + + s"rootCa=$rootCa, " + + s"arrowAllocatorBytes=$arrowAllocatorBytes" + ")" } object BigQueryStorageSettings { + val DefaultArrowAllocatorBytes: Long = 512L * 1024 * 1024 // 512 MB + /** * Create settings for unsecure (no tls), unauthenticated (no root ca) * and unauthorized (no call credentials) endpoint. @@ -73,7 +89,12 @@ object BigQueryStorageSettings { case _ => bigQueryConfig } - Seq(setRootCa).foldLeft(bigQueryConfig) { + val setAllocatorBytes = (bigQueryConfig: BigQueryStorageSettings) => + if (config.hasPath("arrowAllocatorBytes")) + bigQueryConfig.withArrowAllocatorBytes(config.getBytes("arrowAllocatorBytes")) + else bigQueryConfig + + Seq(setRootCa, setAllocatorBytes).foldLeft(bigQueryConfig) { case (c, f) => f(c) } } diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala index 50dcf91f9..cb9e5986b 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala @@ -16,7 +16,7 @@ package org.apache.pekko.stream.connectors.googlecloud.bigquery.storage.impl import org.apache.pekko import pekko.NotUsed import pekko.stream.connectors.googlecloud.bigquery.storage.BigQueryRecord -import pekko.stream.scaladsl.{ Merge, Source } +import pekko.stream.scaladsl.Source import com.google.cloud.bigquery.storage.v1.arrow.{ ArrowRecordBatch, ArrowSchema } import com.google.cloud.bigquery.storage.v1.storage.BigQueryReadClient import com.google.cloud.bigquery.storage.v1.stream.ReadSession @@ -32,24 +32,24 @@ import scala.jdk.CollectionConverters._ object ArrowSource { - def readRecordsMerged(client: BigQueryReadClient, readSession: ReadSession): Source[List[BigQueryRecord], NotUsed] = + def readRecordsMerged(client: BigQueryReadClient, + readSession: ReadSession, + allocatorBytes: Long): Source[List[BigQueryRecord], NotUsed] = readMerged(client, readSession) .map { a => - val reader = new SimpleRowReader(readSession.schema.arrowSchema.get) + val reader = new SimpleRowReader(readSession.schema.arrowSchema.get, allocatorBytes) try reader.read(a) finally reader.close() } def readMerged(client: BigQueryReadClient, session: ReadSession): Source[ArrowRecordBatch, NotUsed] = - read(client, session) match { - case Seq(single) => single - case sources => Source.combine(sources.head, sources.tail.head, sources.tail.tail: _*)(Merge(_)) - } + read(client, session).reduce((a, b) => a.merge(b)) - def readRecords(client: BigQueryReadClient, session: ReadSession): Seq[Source[BigQueryRecord, NotUsed]] = + def readRecords(client: BigQueryReadClient, session: ReadSession, + allocatorBytes: Long): Seq[Source[BigQueryRecord, NotUsed]] = read(client, session) .map { a => - a.map(new SimpleRowReader(session.schema.arrowSchema.get).read(_)) + a.map(new SimpleRowReader(session.schema.arrowSchema.get, allocatorBytes).read(_)) .mapConcat(c => c) } @@ -62,14 +62,9 @@ object ArrowSource { } -object SimpleRowReader { - val DefaultAllocationLimit: Long = 256L * 1024 * 1024 // 256 MB -} - -final class SimpleRowReader(val schema: ArrowSchema, allocationLimit: Long = SimpleRowReader.DefaultAllocationLimit) - extends AutoCloseable { +final class SimpleRowReader(val schema: ArrowSchema, allocatorBytes: Long) extends AutoCloseable { - val allocator = new RootAllocator(allocationLimit) + val allocator = new RootAllocator(allocatorBytes) val sd = MessageSerializer.deserializeSchema( new ReadChannel( diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala index 2c4b15a33..47398cd47 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/AvroSource.scala @@ -16,7 +16,7 @@ package org.apache.pekko.stream.connectors.googlecloud.bigquery.storage.impl import org.apache.pekko import pekko.NotUsed import pekko.stream.connectors.googlecloud.bigquery.storage.BigQueryRecord -import pekko.stream.scaladsl.{ Merge, Source } +import pekko.stream.scaladsl.Source import com.google.cloud.bigquery.storage.v1.avro.AvroRows import com.google.cloud.bigquery.storage.v1.storage.BigQueryReadClient import com.google.cloud.bigquery.storage.v1.stream.ReadSession @@ -30,10 +30,7 @@ object AvroSource { } def readMerged(client: BigQueryReadClient, session: ReadSession): Source[AvroRows, NotUsed] = - read(client, session) match { - case Seq(single) => single - case sources => Source.combine(sources.head, sources.tail.head, sources.tail.tail: _*)(Merge(_)) - } + read(client, session).reduce((a, b) => a.merge(b)) def readRecords(client: BigQueryReadClient, session: ReadSession): Seq[Source[BigQueryRecord, NotUsed]] = read(client, session) diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorage.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorage.scala index f1cd727fe..9553c4b2d 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorage.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorage.scala @@ -42,7 +42,7 @@ object BigQueryArrowStorage { tableId, readOptions, maxNumStreams, - (_, client, session) => ArrowSource.readRecordsMerged(client, session)) + (_, client, session, allocatorBytes) => ArrowSource.readRecordsMerged(client, session, allocatorBytes)) .flatMapConcat(a => a) def readRecords(projectId: String, @@ -55,7 +55,7 @@ object BigQueryArrowStorage { tableId, readOptions, maxNumStreams, - (_, client, session) => ArrowSource.readRecords(client, session)) + (_, client, session, allocatorBytes) => ArrowSource.readRecords(client, session, allocatorBytes)) def readMerged(projectId: String, datasetId: String, @@ -67,7 +67,7 @@ object BigQueryArrowStorage { tableId, readOptions, maxNumStreams, - (schema, client, session) => (schema, ArrowSource.readMerged(client, session))) + (schema, client, session, _) => (schema, ArrowSource.readMerged(client, session))) def read(projectId: String, datasetId: String, @@ -79,20 +79,22 @@ object BigQueryArrowStorage { tableId, readOptions, maxNumStreams, - (schema, client, session) => (schema, ArrowSource.read(client, session))) + (schema, client, session, _) => (schema, ArrowSource.read(client, session))) private def readAndMapTo[T](projectId: String, datasetId: String, tableId: String, readOptions: Option[TableReadOptions], maxNumStreams: Int, - fx: (ArrowSchema, BigQueryReadClient, ReadSession) => T): Source[T, Future[NotUsed]] = + fx: (ArrowSchema, BigQueryReadClient, ReadSession, Long) => T): Source[T, Future[NotUsed]] = Source.fromMaterializer { (mat, attr) => - val client = reader(mat.system, attr).client + val rdr = reader(mat.system, attr) + val client = rdr.client + val allocatorBytes = rdr.settings.arrowAllocatorBytes readSession(client, projectId, datasetId, tableId, DataFormat.ARROW, readOptions, maxNumStreams) .map { session => session.schema match { - case ReadSession.Schema.ArrowSchema(schema) => fx(schema, client, session) + case ReadSession.Schema.ArrowSchema(schema) => fx(schema, client, session, allocatorBytes) case other => throw new IllegalArgumentException(s"Only Arrow format is supported, received: $other") } } diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/GrpcBigQueryStorageReader.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/GrpcBigQueryStorageReader.scala index 19370175f..a18f23e4d 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/GrpcBigQueryStorageReader.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/GrpcBigQueryStorageReader.scala @@ -23,7 +23,8 @@ import com.google.cloud.bigquery.storage.v1.storage.BigQueryReadClient /** * Holds the gRPC scala reader client instance. */ -final class GrpcBigQueryStorageReader private (settings: BigQueryStorageSettings, sys: ClassicActorSystemProvider) { +final class GrpcBigQueryStorageReader private[scaladsl] (val settings: BigQueryStorageSettings, + sys: ClassicActorSystemProvider) { @ApiMayChange final val client = BigQueryReadClient(PekkoGrpcSettings.fromBigQuerySettings(settings)(sys))(sys) diff --git a/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorageSpec.scala b/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorageSpec.scala index fd29b32bb..29a476548 100644 --- a/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorageSpec.scala +++ b/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryArrowStorageSpec.scala @@ -32,7 +32,8 @@ class BigQueryArrowStorageSpec "BigQueryArrowStorage.readArrow" should { - val reader = new SimpleRowReader(ArrowSchema(serializedSchema = GCPSerializedArrowSchema)) + val reader = new SimpleRowReader(ArrowSchema(serializedSchema = GCPSerializedArrowSchema), + BigQueryStorageSettings.DefaultArrowAllocatorBytes) val expectedRecords = reader.read(ArrowRecordBatch(GCPSerializedArrowTenRecordBatch, 10)) "stream the results for a query in records merged" in { @@ -83,7 +84,7 @@ class BigQueryArrowStorageSpec .futureValue .head - val rowReader = new SimpleRowReader(schema) + val rowReader = new SimpleRowReader(schema, BigQueryStorageSettings.DefaultArrowAllocatorBytes) val records = rowReader.read(recordBatch) records shouldBe expectedRecords diff --git a/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryStorageSpec.scala b/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryStorageSpec.scala index efff7c244..24601819c 100644 --- a/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryStorageSpec.scala +++ b/google-cloud-bigquery-storage/src/test/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/scaladsl/BigQueryStorageSpec.scala @@ -103,7 +103,8 @@ class BigQueryStorageSpec } "stream the results for a query using arrow deserializer" in { - val reader = new SimpleRowReader(ArrowSchema(serializedSchema = GCPSerializedArrowSchema)) + val reader = new SimpleRowReader(ArrowSchema(serializedSchema = GCPSerializedArrowSchema), + BigQueryStorageSettings.DefaultArrowAllocatorBytes) val expectedRecords = reader.read(ArrowRecordBatch(GCPSerializedArrowTenRecordBatch, 10)) implicit val um: ArrowByteStringDecoder = new ArrowByteStringDecoder(ArrowSchema(GCPSerializedArrowSchema)) From c1531ea28b715126f90a2826283a7b691707d6cf Mon Sep 17 00:00:00 2001 From: PJ Fanning Date: Mon, 10 Aug 2026 09:55:21 +0100 Subject: [PATCH 3/3] close SimpleRowReader --- .../googlecloud/bigquery/storage/impl/ArrowSource.scala | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala index cb9e5986b..3fb802092 100644 --- a/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala +++ b/google-cloud-bigquery-storage/src/main/scala/org/apache/pekko/stream/connectors/googlecloud/bigquery/storage/impl/ArrowSource.scala @@ -49,8 +49,11 @@ object ArrowSource { allocatorBytes: Long): Seq[Source[BigQueryRecord, NotUsed]] = read(client, session) .map { a => - a.map(new SimpleRowReader(session.schema.arrowSchema.get, allocatorBytes).read(_)) - .mapConcat(c => c) + a.map { batch => + val reader = new SimpleRowReader(session.schema.arrowSchema.get, allocatorBytes) + try reader.read(batch) + finally reader.close() + }.mapConcat(c => c) } def read(client: BigQueryReadClient, session: ReadSession): Seq[Source[ArrowRecordBatch, NotUsed]] =