hbase batch mutate

TongchengOpenSource · Dec 9, 2024 · 8056f8e · 8056f8e
1 parent 16c17a1
commit 8056f8e
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 7 deletions.
diff --git a/...e/src/main/java/org/apache/seatunnel/connectors/seatunnel/hbase/catalog/HbaseCatalog.java b/...e/src/main/java/org/apache/seatunnel/connectors/seatunnel/hbase/catalog/HbaseCatalog.java
@@ -69,7 +69,14 @@ public void open() throws CatalogException {
 
     @Override
     public void close() throws CatalogException {
-        hbaseClient.close();
+        if (this.hbaseClient != null) {
+            try {
+                this.hbaseClient.close();
+            } catch (Exception e) {
+                throw new CatalogException("Failed to close HBase connection.", e);
+            }
+            log.info("Current HBase connection is closed.");
+        }
     }
 
     @Override

diff --git a/...ase/src/main/java/org/apache/seatunnel/connectors/seatunnel/hbase/client/HbaseClient.java b/...ase/src/main/java/org/apache/seatunnel/connectors/seatunnel/hbase/client/HbaseClient.java
@@ -354,6 +354,16 @@ public void mutate(Put put) throws IOException {
         hbaseMutator.mutate(put);
     }
 
+    /**
+     * Mutate Put List.
+     *
+     * @param putList Hbase put list
+     * @throws IOException exception
+     */
+    public void batchMutate(List<Put> putList) throws IOException {
+        hbaseMutator.mutate(putList);
+    }
+
     /**
      * Scan a table.
      *

diff --git a/...e/src/main/java/org/apache/seatunnel/connectors/seatunnel/hbase/sink/HbaseSinkWriter.java b/...e/src/main/java/org/apache/seatunnel/connectors/seatunnel/hbase/sink/HbaseSinkWriter.java
@@ -22,7 +22,7 @@
 import org.apache.seatunnel.api.table.type.SeaTunnelDataType;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
-import org.apache.seatunnel.common.exception.CommonErrorCodeDeprecated;
+import org.apache.seatunnel.common.exception.CommonErrorCode;
 import org.apache.seatunnel.connectors.seatunnel.hbase.client.HbaseClient;
 import org.apache.seatunnel.connectors.seatunnel.hbase.config.HbaseParameters;
 import org.apache.seatunnel.connectors.seatunnel.hbase.exception.HbaseConnectorException;
@@ -36,8 +36,8 @@
 
 import java.io.IOException;
 import java.nio.charset.Charset;
+import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -56,10 +56,14 @@ public class HbaseSinkWriter
 
     private List<Integer> rowkeyColumnIndexes;
 
+    public static final int BATCH_NUM = 10000;
+
     private int versionColumnIndex;
 
     private String defaultFamilyName = "value";
 
+    private List<Put> putList;
+
     public HbaseSinkWriter(
             SeaTunnelRowType seaTunnelRowType,
             HbaseParameters hbaseParameters,
@@ -76,12 +80,22 @@ public HbaseSinkWriter(
         }
 
         this.hbaseClient = HbaseClient.createInstance(hbaseParameters);
+        putList = new ArrayList<>();
     }
 
     @Override
     public void write(SeaTunnelRow element) throws IOException {
         Put put = convertRowToPut(element);
-        hbaseClient.mutate(put);
+        if (!put.isEmpty()) {
+            putList.add(put);
+        }
+        if (putList.size() >= BATCH_NUM) {
+            try {
+                hbaseClient.batchMutate(putList);
+            } finally {
+                putList.clear();
+            }
+        }
     }
 
     @Override
@@ -94,6 +108,15 @@ public void abortPrepare() {}
 
     @Override
     public void close() throws IOException {
+
+        if (!putList.isEmpty()) {
+            try {
+                hbaseClient.batchMutate(putList);
+            } finally {
+                putList.clear();
+            }
+        }
+
         if (hbaseClient != null) {
             hbaseClient.close();
         }
@@ -120,7 +143,6 @@ private Put convertRowToPut(SeaTunnelRow row) {
                         .collect(Collectors.toList());
         for (Integer writeColumnIndex : writeColumnIndexes) {
             String fieldName = seaTunnelRowType.getFieldName(writeColumnIndex);
-            Map<String, String> configurationFamilyNames = hbaseParameters.getFamilyNames();
             String familyName =
                     hbaseParameters.getFamilyNames().getOrDefault(fieldName, defaultFamilyName);
             byte[] bytes = convertColumnToBytes(row, writeColumnIndex);
@@ -184,8 +206,7 @@ private byte[] convertColumnToBytes(SeaTunnelRow row, int index) {
                         String.format(
                                 "Hbase connector does not support this column type [%s]",
                                 fieldType.getSqlType());
-                throw new HbaseConnectorException(
-                        CommonErrorCodeDeprecated.UNSUPPORTED_DATA_TYPE, errorMsg);
+                throw new HbaseConnectorException(CommonErrorCode.UNSUPPORTED_DATA_TYPE, errorMsg);
         }
     }
 }