From afd33e0f43779d8ae39e27f33ca3a09500833679 Mon Sep 17 00:00:00 2001 From: Akanksha Akkihal Date: Tue, 28 Jul 2026 11:46:44 -0700 Subject: [PATCH 1/3] fix(tools): batch ClickHouse JSON-lines bulk loads Bulk-loading a JSON-lines dataset piped the entire file through a single INSERT ... FORMAT JSONEachRow, which is unreliable for the multi-GB datasets the ClickHouse benchmark now uses. Load in bounded batches via a dedicated loader script instead, poll ClickHouse HTTP until it is actually reachable before loading, and validate each line so malformed input fails with the offending line number rather than an opaque client error. Also let init_sql_file own DROP/CREATE for its own objects, so a schema that defines dependent materialized views is not broken by a standalone DROP TABLE. Co-authored-by: Cursor --- .../services/clickhouse_service.py | 169 ++++++++++-------- .../experiment_utils/services/json/loader.py | 151 ++++++++++++++++ 2 files changed, 250 insertions(+), 70 deletions(-) create mode 100644 asap-tools/experiments/experiment_utils/services/json/loader.py diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index a423c944..6e9c6eb6 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -5,6 +5,7 @@ import os import shlex import subprocess +import time from typing import Optional from jinja2 import Template @@ -192,8 +193,10 @@ class ClickHouseDataLoaderService(BaseService): # H2O loader script (relative to _ASSETS_DIR). H2O_LOADER_SCRIPT = "h2o/loader.py" + JSON_LOADER_SCRIPT = "json/loader.py" H2O_BATCH_SIZE = 50_000 + JSON_BATCH_SIZE = 100_000 DEFAULT_TABLES = { "clickbench": "hits", @@ -287,13 +290,15 @@ def start( url = f"http://localhost:{self.clickhouse_http_port}/" - print(f"Dropping table {table!r}...") - self._exec_sql(f"DROP TABLE IF EXISTS {table}", url) - if init_sql_file is not None: + # init SQL owns DROP/CREATE for the raw table and any MVs / agg + # tables (e.g. netflow_init.sql). Skipping the standalone DROP + # avoids failing on dependents that still reference ``table``. print(f"Running init SQL from {init_sql_file!r}...") self._exec_sql_file(init_sql_file, url) elif dataset_name in self.BUILTIN_DDL_FILES: + print(f"Dropping table {table!r}...") + self._exec_sql(f"DROP TABLE IF EXISTS {table}", url) local_ddl = os.path.join( self._ASSETS_DIR, self.BUILTIN_DDL_FILES[dataset_name] ) @@ -309,6 +314,8 @@ def start( f"No built-in DDL for dataset_name={dataset_name!r}; pass init_sql_file" ) + self._ensure_clickhouse_http_ready(url) + if dataset_name == "clickbench": self._load_clickbench(remote_data_file, url, table, max_rows) elif dataset_name == "h2o": @@ -403,6 +410,84 @@ def _exec_sql_file(self, remote_sql_file: str, url: str) -> None: for stmt in (s.strip() for s in result.stdout.split(";") if s.strip()): self._exec_sql(stmt, url) + def _ensure_clickhouse_http_ready(self, url: str, max_retries: int = 30) -> None: + """Wait until ClickHouse HTTP /ping returns Ok. before bulk load.""" + ping_url = url.rstrip("/") + "/ping" + for attempt in range(max_retries): + result = self.provider.execute_command( + node_idx=self.node_offset, + cmd=f"curl -sS {shlex.quote(ping_url)}", + cmd_dir=None, + nohup=False, + popen=False, + ignore_errors=True, + ) + if ( + isinstance(result, subprocess.CompletedProcess) + and result.returncode == 0 + and result.stdout.strip() == "Ok." + ): + return + print( + f"Waiting for ClickHouse HTTP before data load... " + f"({attempt + 1}/{max_retries})" + ) + time.sleep(2) + + logs = self.provider.execute_command( + node_idx=self.node_offset, + cmd="docker logs clickhouse-server --tail 30 2>&1", + cmd_dir=None, + nohup=False, + popen=False, + ignore_errors=True, + ) + log_tail = "" + if isinstance(logs, subprocess.CompletedProcess): + log_tail = (logs.stdout or logs.stderr or "").strip()[-500:] + raise RuntimeError( + "ClickHouse HTTP is not ready before data load. " + f"Check clickhouse_logs/ on the experiment node. Recent logs: {log_tail}" + ) + + def _load_json_batched( + self, + remote_data_file: str, + url: str, + table: str, + max_rows: int, + dataset_label: str, + batch_size: int = JSON_BATCH_SIZE, + ) -> None: + """Load JSON-lines via batched HTTP INSERTs (safe for multi-GB files).""" + local_script = os.path.join(self._ASSETS_DIR, self.JSON_LOADER_SCRIPT) + remote_script = f"/tmp/json_loader_{os.getpid()}.py" + self._rsync_to_remote(local_script, remote_script) + try: + cmd = ( + "python3 {} --data-file {} --table {} " + "--batch-size {} --max-rows {} --container {}" + ).format( + shlex.quote(remote_script), + shlex.quote(remote_data_file), + shlex.quote(table), + batch_size, + max_rows, + shlex.quote(ClickHouseService.CONTAINER_NAME), + ) + result = self.provider.execute_command( + node_idx=self.node_offset, + cmd=cmd, + cmd_dir=None, + nohup=False, + popen=False, + ) + if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: + detail = (result.stderr or result.stdout or "").strip()[:300] + raise RuntimeError(f"{dataset_label} data load failed: {detail}") + finally: + self._remote_rm(remote_script) + def _check_row_count(self, table: str, url: str) -> int: """Return the row count for a table on the remote node, or 0 on error.""" cmd = "curl -sS {} --data {}".format( @@ -429,40 +514,9 @@ def _load_clickbench( ) -> None: """Stream a JSON-lines file (optionally gzipped) into ClickHouse.""" print(f"Loading ClickBench data from {remote_data_file!r}...") - file_lower = remote_data_file.lower() - is_gz = file_lower.endswith(".json.gz") or file_lower.endswith(".jsonl.gz") - insert_sql = shlex.quote(f"INSERT INTO {table} FORMAT JSONEachRow") - - if is_gz: - if max_rows > 0: - reader = "zcat {} | head -n {}".format( - shlex.quote(remote_data_file), max_rows - ) - else: - reader = "zcat {}".format(shlex.quote(remote_data_file)) - else: - if max_rows > 0: - reader = "head -n {} {}".format(max_rows, shlex.quote(remote_data_file)) - else: - reader = "cat {}".format(shlex.quote(remote_data_file)) - - cmd = ( - "{} | docker exec -i clickhouse-server clickhouse-client --query {}".format( - reader, insert_sql - ) - ) - - result = self.provider.execute_command( - node_idx=self.node_offset, - cmd=cmd, - cmd_dir=None, - nohup=False, - popen=False, + self._load_json_batched( + remote_data_file, url, table, max_rows, "ClickBench" ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: - raise RuntimeError( - f"ClickBench data load failed: {result.stderr.strip()[:200]}" - ) def _load_h2o( self, remote_data_file: str, url: str, batch_size: int, max_rows: int @@ -503,42 +557,17 @@ def _load_custom( """Stream a custom JSON-lines file (plain or gzipped) into ClickHouse.""" print(f"Loading custom data from {remote_data_file!r} into {table!r}...") file_lower = remote_data_file.lower() - is_gz = file_lower.endswith(".json.gz") or file_lower.endswith(".jsonl.gz") - is_json = file_lower.endswith(".json") or file_lower.endswith(".jsonl") - insert_sql = shlex.quote(f"INSERT INTO {table} FORMAT JSONEachRow") - - if is_gz: - if max_rows > 0: - reader = "zcat {} | head -n {}".format( - shlex.quote(remote_data_file), max_rows - ) - else: - reader = "zcat {}".format(shlex.quote(remote_data_file)) - elif is_json: - if max_rows > 0: - reader = "head -n {} {}".format(max_rows, shlex.quote(remote_data_file)) - else: - reader = "cat {}".format(shlex.quote(remote_data_file)) - else: + is_json = ( + file_lower.endswith(".json") + or file_lower.endswith(".jsonl") + or file_lower.endswith(".json.gz") + or file_lower.endswith(".jsonl.gz") + ) + if not is_json: raise ValueError( f"Unsupported file format for {remote_data_file!r}. " "Use dataset_name='h2o' for CSV files." ) - - cmd = ( - "{} | docker exec -i clickhouse-server clickhouse-client --query {}".format( - reader, insert_sql - ) + self._load_json_batched( + remote_data_file, url, table, max_rows, "Custom" ) - - result = self.provider.execute_command( - node_idx=self.node_offset, - cmd=cmd, - cmd_dir=None, - nohup=False, - popen=False, - ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: - raise RuntimeError( - f"Custom data load failed: {result.stderr.strip()[:200]}" - ) diff --git a/asap-tools/experiments/experiment_utils/services/json/loader.py b/asap-tools/experiments/experiment_utils/services/json/loader.py new file mode 100644 index 00000000..fb9dac5f --- /dev/null +++ b/asap-tools/experiments/experiment_utils/services/json/loader.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +""" +Load a JSON-lines file into ClickHouse in bounded-size batches. + +Uses ``docker exec … clickhouse-client`` on the experiment node so multi-GB +files stream without loading the whole payload into curl/HTTP memory. + +Usage: + python3 loader.py \\ + --data-file /path/to/netflow.jsonl \\ + --table netflow_table \\ + --batch-size 100000 \\ + --max-rows 0 +""" + +import argparse +import gzip +import json +import subprocess +import sys + + +DEFAULT_CONTAINER = "clickhouse-server" +PROGRESS_ROW_INTERVAL = 500_000 + + +def _open_data_file(path: str): + lower = path.lower() + if lower.endswith(".gz"): + return gzip.open(path, "rt", encoding="utf-8") + return open(path, "r", encoding="utf-8") + + +def _validate_line(line: str, line_no: int) -> str: + stripped = line.strip() + if not stripped: + return "" + if "\0" in stripped: + preview = stripped[:80].replace("\0", "\\0") + raise RuntimeError( + f"Null byte in JSON line {line_no} (file may be corrupt — " + f"regenerate on a native Linux filesystem, not WSL /mnt/c): {preview!r}" + ) + try: + json.loads(stripped) + except json.JSONDecodeError as exc: + preview = stripped[:120] + raise RuntimeError( + f"Invalid JSON on line {line_no}: {exc}. Preview: {preview!r}" + ) from exc + return stripped + + +def flush_batch( + table: str, + lines: list[str], + container: str, +) -> None: + body = "\n".join(lines).encode("utf-8") + result = subprocess.run( + [ + "docker", + "exec", + "-i", + container, + "clickhouse-client", + "--query", + f"INSERT INTO {table} FORMAT JSONEachRow", + ], + input=body, + capture_output=True, + ) + if result.returncode != 0: + detail = (result.stderr or result.stdout or b"").decode( + "utf-8", errors="replace" + )[:500] + raise RuntimeError(f"ClickHouse insert failed: {detail}") + + +def load( + data_file: str, + table: str, + batch_size: int, + max_rows: int, + container: str = DEFAULT_CONTAINER, +) -> None: + batch: list[str] = [] + total = 0 + file_line_no = 0 + next_progress_report = PROGRESS_ROW_INTERVAL + + def flush(lines: list[str]) -> None: + nonlocal total + if not lines: + return + if max_rows > 0: + lines = lines[: max_rows - total] + if not lines: + return + flush_batch(table, lines, container) + total += len(lines) + + with _open_data_file(data_file) as fin: + for raw_line in fin: + file_line_no += 1 + if max_rows > 0 and total >= max_rows: + break + stripped = _validate_line(raw_line, file_line_no) + if not stripped: + continue + batch.append(stripped) + if len(batch) >= batch_size: + flush(batch) + batch = [] + if total >= next_progress_report: + print(f" Inserted {total:,} rows...", flush=True) + next_progress_report = total + PROGRESS_ROW_INTERVAL + + flush(batch) + print(f"JSON load complete: {total:,} rows into {table!r}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-file", required=True, help="Path to JSON-lines file") + parser.add_argument("--table", required=True, help="Target table name") + parser.add_argument( + "--batch-size", type=int, default=100_000, help="Rows per INSERT batch" + ) + parser.add_argument( + "--max-rows", type=int, default=0, help="Max rows to load (0 = all)" + ) + parser.add_argument( + "--container", + default=DEFAULT_CONTAINER, + help=f"ClickHouse docker container name (default: {DEFAULT_CONTAINER})", + ) + args = parser.parse_args() + if args.batch_size < 1: + parser.error("--batch-size must be >= 1") + try: + load(args.data_file, args.table, args.batch_size, args.max_rows, args.container) + except (RuntimeError, OSError) as exc: + # The caller surfaces only the head of stderr, so print the message + # itself rather than letting a traceback bury it. + print(f"ERROR: {exc}", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() From 089de18428003350552b1ca77810bb1177df2597 Mon Sep 17 00:00:00 2001 From: Akanksha Akkihal Date: Tue, 28 Jul 2026 12:11:54 -0700 Subject: [PATCH 2/3] docs(tools): correct _load_json_batched docstring --- .../experiments/experiment_utils/services/clickhouse_service.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index 6e9c6eb6..95014e71 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -459,7 +459,7 @@ def _load_json_batched( dataset_label: str, batch_size: int = JSON_BATCH_SIZE, ) -> None: - """Load JSON-lines via batched HTTP INSERTs (safe for multi-GB files).""" + """Load JSON-lines in bounded batches via clickhouse-client.""" local_script = os.path.join(self._ASSETS_DIR, self.JSON_LOADER_SCRIPT) remote_script = f"/tmp/json_loader_{os.getpid()}.py" self._rsync_to_remote(local_script, remote_script) From 726640aadd446b3d7ce73a2b6daed86fabbb6b8d Mon Sep 17 00:00:00 2001 From: Akanksha Akkihal Date: Tue, 28 Jul 2026 12:21:20 -0700 Subject: [PATCH 3/3] docs(tools): correct _load_json_batched docstring and apply black formatting --- .../experiment_utils/services/clickhouse_service.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index 95014e71..b3e4f6ac 100644 --- a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py +++ b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py @@ -482,7 +482,10 @@ def _load_json_batched( nohup=False, popen=False, ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: + if ( + isinstance(result, subprocess.CompletedProcess) + and result.returncode != 0 + ): detail = (result.stderr or result.stdout or "").strip()[:300] raise RuntimeError(f"{dataset_label} data load failed: {detail}") finally: @@ -514,9 +517,7 @@ def _load_clickbench( ) -> None: """Stream a JSON-lines file (optionally gzipped) into ClickHouse.""" print(f"Loading ClickBench data from {remote_data_file!r}...") - self._load_json_batched( - remote_data_file, url, table, max_rows, "ClickBench" - ) + self._load_json_batched(remote_data_file, url, table, max_rows, "ClickBench") def _load_h2o( self, remote_data_file: str, url: str, batch_size: int, max_rows: int @@ -568,6 +569,4 @@ def _load_custom( f"Unsupported file format for {remote_data_file!r}. " "Use dataset_name='h2o' for CSV files." ) - self._load_json_batched( - remote_data_file, url, table, max_rows, "Custom" - ) + self._load_json_batched(remote_data_file, url, table, max_rows, "Custom")