feat: add insert_memory_limit option

4paradigm · Nov 30, 2023 · 857cfd6 · 857cfd6
1 parent c36849e
commit 857cfd6
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 0 deletions.
diff --git a/java/openmldb-jdbc/src/main/java/com/_4paradigm/openmldb/sdk/SdkOption.java b/java/openmldb-jdbc/src/main/java/com/_4paradigm/openmldb/sdk/SdkOption.java
@@ -45,13 +45,15 @@ public class SdkOption {
     private int glogLevel = 0;
     private String glogDir = "";
     private int maxSqlCacheSize = 50;
+    private int insertMemoryUsageLimit = 0;  // [0-100], the default value 0 means unlimited
 
     private void buildBaseOptions(BasicRouterOptions opt) {
         opt.setEnable_debug(getEnableDebug());
         opt.setRequest_timeout(getRequestTimeout());
         opt.setGlog_level(getGlogLevel());
         opt.setGlog_dir(getGlogDir());
         opt.setMax_sql_cache_size(getMaxSqlCacheSize());
+        opt.setInsert_memory_usage_limit(getInsertMemoryUsageLimit());
     }
 
     public SQLRouterOptions buildSQLRouterOptions() throws SqlException {

diff --git a/...connector/src/main/java/com/_4paradigm/openmldb/spark/write/OpenmldbDataSingleWriter.java b/...connector/src/main/java/com/_4paradigm/openmldb/spark/write/OpenmldbDataSingleWriter.java
@@ -44,6 +44,7 @@ public OpenmldbDataSingleWriter(OpenmldbWriteConfig config, int partitionId, lon
             SdkOption option = new SdkOption();
             option.setZkCluster(config.zkCluster);
             option.setZkPath(config.zkPath);
+            option.setInsertMemoryUsageLimit(config.memoryLimit);
             SqlClusterExecutor executor = new SqlClusterExecutor(option);
             String dbName = config.dbName;
             String tableName = config.tableName;

diff --git a/...park-connector/src/main/java/com/_4paradigm/openmldb/spark/write/OpenmldbWriteConfig.java b/...park-connector/src/main/java/com/_4paradigm/openmldb/spark/write/OpenmldbWriteConfig.java
@@ -24,13 +24,15 @@
 // Must serializable
 public class OpenmldbWriteConfig implements Serializable {
     public final String dbName, tableName, zkCluster, zkPath, writerType;
+    public final int insertMemoryUsageLimit;
 
     public OpenmldbWriteConfig(String dbName, String tableName, SdkOption option, String writerType) {
         this.dbName = dbName;
         this.tableName = tableName;
         this.zkCluster = option.getZkCluster();
         this.zkPath = option.getZkPath();
         this.writerType = writerType;
+        this.insertMemoryUsageLimit = option.getInsertMemoryUsageLimit();
         // TODO(hw): other configs in SdkOption
     }
 }
diff --git a/src/sdk/sql_router.h b/src/sdk/sql_router.h
@@ -49,6 +49,7 @@ struct BasicRouterOptions {
     int glog_level = 0;
     // empty means to stderr
     std::string glog_dir = "";
+    int insert_memory_usage_limit = 0;
 };
 
 struct SQLRouterOptions : BasicRouterOptions {