feat: enhance JVM configuration and memory management in benchmark scripts

Author: tae898

bindings/python/build_and_install_locally.sh

@@ -22,14 +22,29 @@ docker run --rm \
     -v "${REPO_ROOT}":/src \
     -w /src \
     maven:3.9-amazoncorretto-25 \
-    sh -c "git config --global --add safe.directory /src && ./mvnw -DskipTests -pl package -am package"
+    sh -c "if ! command -v git >/dev/null 2>&1 || ! command -v tar >/dev/null 2>&1; then \
+        if command -v yum >/dev/null 2>&1; then \
+            yum -y install git tar; \
+        elif command -v apt-get >/dev/null 2>&1; then \
+            apt-get update && apt-get install -y git tar; \
+        else \
+            echo '❌ git/tar not found and no supported package manager (yum/apt-get) available' >&2; \
+            exit 1; \
+        fi; \
+    fi; \
+    git config --global --add safe.directory /src && ./mvnw -DskipTests -pl package -am package"
 
 # 2) Copy freshly built JARs into local-jars for the Python build
 log "Staging JARs into bindings/python/local-jars/lib..."
 mkdir -p "${LOCAL_JARS_DIR}"
-HEADLESS_LIB_DIR="${REPO_ROOT}/package/target/arcadedb-26.1.1-SNAPSHOT-headless.dir/arcadedb-26.1.1-SNAPSHOT/lib"
-if [[ ! -d "${HEADLESS_LIB_DIR}" ]]; then
-    echo "❌ Headless package lib directory not found at ${HEADLESS_LIB_DIR}" >&2
+HEADLESS_PARENT_DIR=$(ls -d "${REPO_ROOT}/package/target/arcadedb-"*-headless.dir 2> /dev/null | sort | tail -n 1)
+if [[ -z "${HEADLESS_PARENT_DIR}" ]]; then
+    echo "❌ Headless package directory not found under ${REPO_ROOT}/package/target" >&2
+    exit 1
+fi
+HEADLESS_LIB_DIR=$(ls -d "${HEADLESS_PARENT_DIR}"/arcadedb-*/lib 2> /dev/null | sort | tail -n 1)
+if [[ -z "${HEADLESS_LIB_DIR}" || ! -d "${HEADLESS_LIB_DIR}" ]]; then
+    echo "❌ Headless package lib directory not found under ${HEADLESS_PARENT_DIR}" >&2
     exit 1
 fi
 cp "${HEADLESS_LIB_DIR}"/*.jar "${LOCAL_JARS_DIR}/"
@@ -43,7 +58,7 @@ cd "${PY_BINDINGS_DIR}"
 
 # 4) Install the wheel with uv (force reinstall)
 log "Installing wheel via uv pip..."
-WHEEL_PATH=$(ls -1 dist/arcadedb_embedded-*-linux_x86_64.whl | sort | tail -n 1)
+WHEEL_PATH=$(ls -1 dist/arcadedb_embedded-*.whl | sort | tail -n 1)
 if [[ -z "${WHEEL_PATH}" ]]; then
     echo "❌ No wheel found in dist/. Did the build succeed?" >&2
     exit 1

bindings/python/examples/benchmark-vector/benchmark_arcadedb_msmarco.py

@@ -89,14 +89,16 @@ def _loop():
     return stop_event
 
 
-def heap_tag_from_env() -> str:
-    """Return heap tag derived from ARCADEDB_JVM_ARGS (-Xmx), or 'default'."""
+def heap_tag_from_args(heap_size: str | None, jvm_args: str | None) -> str:
+    """Return heap tag derived from explicit args, or 'default'."""
 
-    jvm_args = os.environ.get("ARCADEDB_JVM_ARGS", "")
-    m = re.search(r"-Xmx(\S+)", jvm_args)
-    if not m:
-        return "default"
-    return m.group(1)
+    if heap_size:
+        return heap_size
+    if jvm_args:
+        m = re.search(r"-Xmx(\S+)", jvm_args)
+        if m:
+            return m.group(1)
+    return "default"
 
 
 # -------------------------
@@ -466,16 +468,38 @@ def main():
         default=100_000,
         help="Number of vectors to ingest per transaction batch",
     )
+    ap.add_argument(
+        "--heap-size",
+        default=None,
+        help=(
+            "Heap size for the embedded JVM (e.g., 8g). Prefer this over "
+            "ARCADEDB_JVM_ARGS."
+        ),
+    )
+    ap.add_argument(
+        "--jvm-args",
+        default=None,
+        help=(
+            "Extra JVM args to pass to start_jvm (e.g., thread flags). "
+            "Prefer this over ARCADEDB_JVM_ARGS."
+        ),
+    )
     args = ap.parse_args()
 
     stop_cpu = start_cpu_logger(2)
 
     np.random.seed(args.seed)
     eval_k = 50
 
-    # Import after potential JVM arg override
+    # Import and configure JVM before any DB creation
     import arcadedb_embedded as arcadedb
 
+    jvm_kwargs = {}
+    if args.heap_size is not None:
+        jvm_kwargs["heap_size"] = args.heap_size
+    if args.jvm_args is not None:
+        jvm_kwargs["jvm_args"] = args.jvm_args
+
     sources, gt_path, dim, label = resolve_dataset(Path(args.dataset_dir))
     total_rows = sum(s["count"] for s in sources)
     count = args.count if args.count is not None else total_rows
@@ -513,7 +537,7 @@ def record(name: str, result, dur, rss_start, rss_end):
     record("load_queries", {"queries": len(queries)}, dur, r0, r1)
 
     # Prepare DB path
-    heap_tag = heap_tag_from_env()
+    heap_tag = heap_tag_from_args(args.heap_size, args.jvm_args)
     param_dir = "_".join(
         [
             f"dataset={Path(args.dataset_dir).name}",
@@ -537,7 +561,8 @@ def record(name: str, result, dur, rss_start, rss_end):
 
     # Create DB
     (db, dur, r0, r1) = timed_section(
-        "create_db", lambda: arcadedb.create_database(str(db_path))
+        "create_db",
+        lambda: arcadedb.create_database(str(db_path), jvm_kwargs=jvm_kwargs),
     )
     record("create_db", {"db_path": str(db_path)}, dur, r0, r1)
 
@@ -670,7 +695,8 @@ def record(name: str, result, dur, rss_start, rss_end):
 
         # Reopen
         (db, dur, r0, r1) = timed_section(
-            "open_db", lambda: arcadedb.open_database(str(db_path))
+            "open_db",
+            lambda: arcadedb.open_database(str(db_path), jvm_kwargs=jvm_kwargs),
         )
         record("open_db", {}, dur, r0, r1)
 

bindings/python/examples/benchmark-vector/run_arcadedb_search_study.py

@@ -90,8 +90,9 @@ def run_single(
     output_root: Path,
     base_config: Dict,
     heap_tag: str | None,
+    jvm_kwargs: Dict,
+    arcadedb,
 ) -> Path:
-    import arcadedb_embedded as arcadedb
 
     sources, gt_path, dim, label = resolve_dataset(dataset_dir)
     total_rows = sum(s["count"] for s in sources)
@@ -111,7 +112,8 @@ def run_single(
     record(phases, "load_queries", {"queries": len(queries)}, dur, r0, r1)
 
     (db, dur, r0, r1) = timed_section(
-        "open_db", lambda: arcadedb.open_database(str(db_path))
+        "open_db",
+        lambda: arcadedb.open_database(str(db_path), jvm_kwargs=jvm_kwargs),
     )
     record(phases, "open_db", {}, dur, r0, r1)
 
@@ -288,6 +290,22 @@ def main() -> None:
         "--heap-tag",
         help="Heap tag to record in outputs (e.g., 8g). Defaults to heap from db path.",
     )
+    ap.add_argument(
+        "--heap-size",
+        default=None,
+        help=(
+            "Heap size for the embedded JVM (e.g., 8g). Prefer this over "
+            "ARCADEDB_JVM_ARGS."
+        ),
+    )
+    ap.add_argument(
+        "--jvm-args",
+        default=None,
+        help=(
+            "Extra JVM args to pass to start_jvm (e.g., thread flags). "
+            "Prefer this over ARCADEDB_JVM_ARGS."
+        ),
+    )
 
     args = ap.parse_args()
 
@@ -299,11 +317,19 @@ def main() -> None:
     if not db_path.exists():
         raise SystemExit(f"DB path not found: {db_path}")
 
+    import arcadedb_embedded as arcadedb
+
+    jvm_kwargs: Dict = {}
+    if args.heap_size is not None:
+        jvm_kwargs["heap_size"] = args.heap_size
+    if args.jvm_args is not None:
+        jvm_kwargs["jvm_args"] = args.jvm_args
+
     base_config = load_existing_config(db_path)
     quant = (base_config.get("quantization") or "NONE").upper()
 
     fallback_heap = run_dir_value(db_path.name, "heap")
-    heap_tag = args.heap_tag or fallback_heap
+    heap_tag = args.heap_tag or args.heap_size or fallback_heap
 
     created: List[Path] = []
     for oq in overqueries:
@@ -317,6 +343,8 @@ def main() -> None:
                 output_root=output_root,
                 base_config=base_config,
                 heap_tag=heap_tag,
+                jvm_kwargs=jvm_kwargs,
+                arcadedb=arcadedb,
             )
         )
 

bindings/python/examples/benchmark-vector/run_arcadedb_search_study.sh

@@ -58,11 +58,17 @@ for HEAP in "${HEAPS[@]}"; do
     fi
 
     echo ">>> RUN: heap=${HEAP} threads=${THREADS} overquery=${OVERQUERY_CSV}" >&2
-    ARCADEDB_JVM_ARGS="$JVM_ARGS" \
-        python "$PY" \
+    HEAP_FLAG=()
+    if [[ "$HEAP" != "default" ]]; then
+        HEAP_FLAG=("--heap-size" "$HEAP")
+    fi
+
+    python "$PY" \
         --db-path "$DB_PATH" \
         --dataset-dir "$DATASET_DIR" \
         --overquery-factors "$OVERQUERY_CSV" \
         --output-root "$OUTPUT_ROOT" \
-        --heap-tag "$HEAP"
+        --heap-tag "$HEAP" \
+        --jvm-args "$JVM_ARGS" \
+        "${HEAP_FLAG[@]}"
 done

bindings/python/examples/benchmark-vector/run_arcadedb_sweep.sh

@@ -23,8 +23,8 @@ if [[ ! -f "$BENCH_PY" ]]; then
     exit 1
 fi
 
-# Heap selection: use ARCADEDB_JVM_ARGS if set, otherwise sweep HEAP_SIZES
-PRESET_JVM_ARGS="${ARCADEDB_JVM_ARGS:-}"
+# Heap selection: use BENCH_JVM_ARGS if set, otherwise sweep HEAP_SIZES
+PRESET_JVM_ARGS="${BENCH_JVM_ARGS:-}"
 HEAP_SIZES_CSV="${HEAP_SIZES:-}"
 HEAP_SIZES=()
 
@@ -43,21 +43,21 @@ HEAP_SIZES=()
 if [[ -z "$PRESET_JVM_ARGS" ]]; then
     # Default heap sweep per dataset (override by setting HEAP_SIZES="8g,12g,16g")
     case "$(basename "$DATASET_DIR")" in
-        # *MSMARCO-100K*) HEAP_SIZES=("1g" "2g") ;;  # 1g works fine. ran with 4 threads
-        # *MSMARCO-1M*) HEAP_SIZES=("4g") ;;   # below 4g heap is already really slow. ran with 4 threads
+        *MSMARCO-100K*) HEAP_SIZES=("2g") ;; # 1g works fine. ran with 4 threads
+        # *MSMARCO-1M*) HEAP_SIZES=("4g") ;;   #
         # *MSMARCO-2M*) HEAP_SIZES=("8g") ;;   # 4 threads
         # *MSMARCO-4M*) HEAP_SIZES=("16g") ;;  # 4 threads
         # *MSMARCO-8M*) HEAP_SIZES=("32g") ;; # 4 threads
-        # *MSMARCO-16M*) HEAP_SIZES=("32g") ;;  # 4 threads
-        *MSMARCO-32M*) HEAP_SIZES=("32g") ;; # 4 threads
+        # *MSMARCO-16M*) HEAP_SIZES=("64g") ;;  # 4 threads
+        # *MSMARCO-32M*) HEAP_SIZES=("128g") ;; # 4 threads
         *) HEAP_SIZES=("default") ;;
     esac
 
     if [[ -n "$HEAP_SIZES_CSV" ]]; then
         IFS="," read -ra HEAP_SIZES <<< "$HEAP_SIZES_CSV"
     fi
 else
-    echo "Using pre-set ARCADEDB_JVM_ARGS='${PRESET_JVM_ARGS}'" >&2
+    echo "Using pre-set BENCH_JVM_ARGS='${PRESET_JVM_ARGS}'" >&2
     HEAP_SIZES=("preset")
 fi
 
@@ -106,11 +106,13 @@ for HEAP in "${HEAP_SIZES[@]}"; do
     fi
 
     BASE_ENV="${THREAD_ENV}"
+    BASE="${BASE_ENV} python \"${BENCH_PY}\" --dataset-dir \"${DATASET_DIR}\" --db-root \"${DB_ROOT}\""
     if [[ -n "$JVM_ARGS" ]]; then
-        BASE_ENV="ARCADEDB_JVM_ARGS=\"${JVM_ARGS}\" ${BASE_ENV}"
+        BASE="${BASE} --jvm-args \"${JVM_ARGS}\""
+    fi
+    if [[ "$HEAP" != "default" && "$HEAP" != "preset" ]]; then
+        BASE="${BASE} --heap-size ${HEAP}"
     fi
-
-    BASE="${BASE_ENV} python \"${BENCH_PY}\" --dataset-dir \"${DATASET_DIR}\" --db-root \"${DB_ROOT}\""
 
     for MC in "${MAX_CONNECTIONS[@]}"; do
         for BW in "${BEAM_WIDTHS[@]}"; do

bindings/python/examples/benchmark-vector/summarize_arcadedb_search_study.py

@@ -76,6 +76,7 @@ def parse_record(path: Path) -> Dict | None:
         "heap": heap_tag,
         "overquery_factor": overquery,
         "quantization": cfg.get("quantization"),
+        "add_hierarchy": cfg.get("add_hierarchy"),
         "max_connections": cfg.get("max_connections"),
         "beam_width": cfg.get("beam_width"),
         "search_s": phases.get("search", {}).get("time_sec"),
@@ -132,6 +133,7 @@ def df_to_markdown(df: pd.DataFrame) -> str:
         "heap",
         "overquery_factor",
         "quantization",
+        "add_hierarchy",
         "max_connections",
         "beam_width",
         "recall@50",

bindings/python/examples/run_with_memory_monitor.sh

@@ -6,7 +6,7 @@
 #   ./run_with_memory_monitor.sh <log_prefix> <python_command>
 #
 # Example:
-#   ./run_with_memory_monitor.sh vector_large "ARCADEDB_JVM_ARGS='-Xmx8g -Xms8g' python 06_vector_search_recommendations.py --source-db my_test_databases/movielens_graph_large_db --db-path my_test_databases/movielens_graph_large_db_vectors"
+#   ./run_with_memory_monitor.sh vector_large "python 06_vector_search_recommendations.py --source-db my_test_databases/movielens_graph_large_db --db-path my_test_databases/movielens_graph_large_db_vectors"
 #
 
 if [ $# -lt 2 ]; then

bindings/python/src/arcadedb_embedded/vector.py

@@ -241,7 +241,7 @@ def find_nearest_approximate(
         self,
         query_vector,
         k=10,
-        overquery_factor=1,
+        overquery_factor=4,
         allowed_rids=None,
     ):
         """