customize the batch size for proton

jovezhong · Jul 17, 2024 · 5cbce10 · 5cbce10
1 parent adaed27
commit 5cbce10
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/core/dbio/database/transaction.go b/core/dbio/database/transaction.go
@@ -322,6 +322,8 @@ func InsertBatchStream(conn Connection, tx Transaction, tableFName string, ds *i
 				row = processClickhouseInsertRow(bColumns, row)
 			} else if conn.GetType() == dbio.TypeDbTrino {
 				row = processTrinoInsertRow(bColumns, row)
+			} else if conn.GetType() == dbio.TypeDbProton {
+				row = processClickhouseInsertRow(bColumns, row)
 			}
 			vals = append(vals, row...)
 		}
@@ -391,6 +393,12 @@ func InsertBatchStream(conn Connection, tx Transaction, tableFName string, ds *i
 			batchSize = cast.ToInt(conn.GetTemplateValue("variable.batch_values")) / len(columns)
 		}
 
+		if conn.GetType() == dbio.TypeDbProton {
+			batchSize = 100
+		} else {
+			batchSize = cast.ToInt(conn.GetTemplateValue("variable.batch_values")) / len(columns)
+		}
+
 		for row := range batch.Rows {
 			batchRows = append(batchRows, row)
 			count++