diff --git a/asap-tools/experiments/experiment_utils/services/clickhouse_service.py b/asap-tools/experiments/experiment_utils/services/clickhouse_service.py index a423c94..b3e4f6a 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,87 @@ 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 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) + 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 +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}...") - 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, - ) - if isinstance(result, subprocess.CompletedProcess) and result.returncode != 0: - raise RuntimeError( - f"ClickBench data load failed: {result.stderr.strip()[:200]}" - ) + 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 @@ -503,42 +558,15 @@ 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 - ) - ) - - 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]}" - ) + self._load_json_batched(remote_data_file, url, table, max_rows, "Custom") 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 0000000..fb9dac5 --- /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()