[beam] branch master updated: [JdbcIO] Allow fetchSize to be set for partitioned reads (#28999)

[yhu]

This is an automated email from the ASF dual-hosted git repository.

yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new 6a57d0d49e6 [JdbcIO] Allow fetchSize to be set for partitioned reads 
(#28999)
6a57d0d49e6 is described below

commit 6a57d0d49e6864ec16e6fb2482af6f77ef8a10e8
Author: Bruno Volpato <bvolp...@google.com>
AuthorDate: Mon Oct 16 11:32:59 2023 -0400

    [JdbcIO] Allow fetchSize to be set for partitioned reads (#28999)
---
 .../jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git 
a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java 
b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
index 6e7ad865cc3..f8dad23d1fb 100644
--- a/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
+++ b/sdks/java/io/jdbc/src/main/java/org/apache/beam/sdk/io/jdbc/JdbcIO.java
@@ -360,6 +360,7 @@ public class JdbcIO {
     return new AutoValue_JdbcIO_ReadWithPartitions.Builder<T, 
PartitionColumnT>()
         .setPartitionColumnType(partitioningColumnType)
         .setNumPartitions(DEFAULT_NUM_PARTITIONS)
+        .setFetchSize(DEFAULT_FETCH_SIZE)
         .setUseBeamSchema(false)
         .build();
   }
@@ -1195,6 +1196,9 @@ public class JdbcIO {
     @Pure
     abstract @Nullable String getPartitionColumn();
 
+    @Pure
+    abstract int getFetchSize();
+
     @Pure
     abstract boolean getUseBeamSchema();
 
@@ -1233,6 +1237,8 @@ public class JdbcIO {
 
       abstract Builder<T, PartitionColumnT> setUseBeamSchema(boolean 
useBeamSchema);
 
+      abstract Builder<T, PartitionColumnT> setFetchSize(int fetchSize);
+
       abstract Builder<T, PartitionColumnT> setTable(String tableName);
 
       abstract Builder<T, PartitionColumnT> setPartitionColumnType(
@@ -1357,7 +1363,8 @@ public class JdbcIO {
                         .withRowMapper(
                             checkStateNotNull(
                                 
JdbcUtil.JdbcReadWithPartitionsHelper.getPartitionsHelper(
-                                    getPartitionColumnType()))))
+                                    getPartitionColumnType())))
+                        .withFetchSize(getFetchSize()))
                 .apply(
                     MapElements.via(
                         new SimpleFunction<
@@ -1421,6 +1428,7 @@ public class JdbcIO {
                   String.format(
                       "select * from %1$s where %2$s >= ? and %2$s < ?", 
table, partitionColumn))
               .withRowMapper(rowMapper)
+              .withFetchSize(getFetchSize())
               .withParameterSetter(
                   checkStateNotNull(
                           
JdbcUtil.JdbcReadWithPartitionsHelper.getPartitionsHelper(