Skip to content

Commit 1ad876f

Browse files
kavyabhat02kavyabhat02
committed
fix formatting
1 parent 14dbd06 commit 1ad876f

2 files changed

Lines changed: 85 additions & 46 deletions

File tree

‎asap-tools/execution-utilities/benchmark/export_to_database.py‎

Lines changed: 40 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,7 @@ def load_h2o_clickhouse(
195195
print(f"Loaded {total:,} rows into ClickHouse (h2o_groupby)")
196196
return True
197197

198+
198199
def load_h2o_elasticsearch(
199200
es_host: str,
200201
es_port: int,
@@ -229,27 +230,30 @@ def load_h2o_elasticsearch(
229230
es.indices.delete(index=index_name)
230231

231232
print(f"Creating index: {index_name}")
232-
es.indices.create(index=index_name, body={
233-
"settings": {
234-
"number_of_shards": 1,
235-
"number_of_replicas": 0,
236-
"refresh_interval": "30s",
237-
},
238-
"mappings": {
239-
"properties": {
240-
"timestamp": {"type": "date", "format": "epoch_millis"},
241-
"id1": {"type": "keyword"},
242-
"id2": {"type": "keyword"},
243-
"id3": {"type": "keyword"},
244-
"id4": {"type": "long"},
245-
"id5": {"type": "long"},
246-
"id6": {"type": "long"},
247-
"v1": {"type": "long"},
248-
"v2": {"type": "long"},
249-
"v3": {"type": "double"},
250-
}
233+
es.indices.create(
234+
index=index_name,
235+
body={
236+
"settings": {
237+
"number_of_shards": 1,
238+
"number_of_replicas": 0,
239+
"refresh_interval": "30s",
240+
},
241+
"mappings": {
242+
"properties": {
243+
"timestamp": {"type": "date", "format": "epoch_millis"},
244+
"id1": {"type": "keyword"},
245+
"id2": {"type": "keyword"},
246+
"id3": {"type": "keyword"},
247+
"id4": {"type": "long"},
248+
"id5": {"type": "long"},
249+
"id6": {"type": "long"},
250+
"v1": {"type": "long"},
251+
"v2": {"type": "long"},
252+
"v3": {"type": "double"},
253+
}
254+
},
251255
},
252-
})
256+
)
253257

254258
if not os.path.exists(file_path):
255259
print(f"ERROR: Data file not found: {file_path}")
@@ -278,9 +282,9 @@ def generate_docs():
278282
"id4": int(parts[3] or 0),
279283
"id5": int(parts[4] or 0),
280284
"id6": int(parts[5] or 0),
281-
"v1": int(parts[6] or 0),
282-
"v2": int(parts[7] or 0),
283-
"v3": float(parts[8] or 0.0),
285+
"v1": int(parts[6] or 0),
286+
"v2": int(parts[7] or 0),
287+
"v3": float(parts[8] or 0.0),
284288
},
285289
}
286290

@@ -302,6 +306,7 @@ def generate_docs():
302306
print(f"✓ Import complete! Index: {index_name}")
303307
return True
304308

309+
305310
def load_custom(
306311
clickhouse_url: str,
307312
file_path: str,
@@ -441,12 +446,20 @@ def main():
441446
)
442447

443448
# Elasticsearch-specific flags
444-
es_group = parser.add_argument_group("Elasticsearch options (--database elasticsearch)")
449+
es_group = parser.add_argument_group(
450+
"Elasticsearch options (--database elasticsearch)"
451+
)
445452
es_group.add_argument("--es-host", default="localhost", help="Elasticsearch host")
446-
es_group.add_argument("--es-port", type=int, default=9200, help="Elasticsearch port")
447-
es_group.add_argument("--es-index", default="h2o_benchmark", help="Elasticsearch index name")
453+
es_group.add_argument(
454+
"--es-port", type=int, default=9200, help="Elasticsearch port"
455+
)
456+
es_group.add_argument(
457+
"--es-index", default="h2o_benchmark", help="Elasticsearch index name"
458+
)
448459
es_group.add_argument("--es-api-key", default=None, help="Elasticsearch API key")
449-
es_group.add_argument("--es-bulk-size", type=int, default=5000, help="Bulk insert batch size")
460+
es_group.add_argument(
461+
"--es-bulk-size", type=int, default=5000, help="Bulk insert batch size"
462+
)
450463

451464
args = parser.parse_args()
452465

‎asap-tools/execution-utilities/benchmark/run_benchmark.py‎

Lines changed: 45 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,9 @@ def run_query(
108108
if api_key:
109109
headers["Authorization"] = f"ApiKey {api_key}"
110110
body = {"query": query.strip().rstrip(";"), "fetch_size": fetch_size}
111-
response = session.post(endpoint_url, headers=headers, json=body, timeout=timeout)
111+
response = session.post(
112+
endpoint_url, headers=headers, json=body, timeout=timeout
113+
)
112114
else:
113115
encoded_query = urllib.parse.quote(query)
114116
separator = "&" if "?" in endpoint_url else "?"
@@ -132,7 +134,10 @@ def run_query(
132134
if hits:
133135
col_names = list(hits[0].get("_source", {}).keys())
134136
formatted_rows = [
135-
", ".join(f"{k}={hit.get('_source', {}).get(k)}" for k in col_names)
137+
", ".join(
138+
f"{k}={hit.get('_source', {}).get(k)}"
139+
for k in col_names
140+
)
136141
for hit in hits
137142
]
138143
result_text = "\n".join(formatted_rows)
@@ -144,12 +149,18 @@ def run_query(
144149
elif "rows" in data:
145150
rows = data.get("rows", [])
146151
columns = data.get("columns", [])
147-
col_names = [c.get("name", f"col{i}") for i, c in enumerate(columns)]
152+
col_names = [
153+
c.get("name", f"col{i}") for i, c in enumerate(columns)
154+
]
148155
formatted_rows = [
149-
", ".join(
150-
f"{col_names[i]}={v}" if i < len(col_names) else str(v)
151-
for i, v in enumerate(row)
152-
) if isinstance(row, (list, tuple)) else str(row)
156+
(
157+
", ".join(
158+
f"{col_names[i]}={v}" if i < len(col_names) else str(v)
159+
for i, v in enumerate(row)
160+
)
161+
if isinstance(row, (list, tuple))
162+
else str(row)
163+
)
153164
for row in rows
154165
]
155166
result_text = "\n".join(formatted_rows)
@@ -267,8 +278,13 @@ def run_benchmark(
267278
last_result, last_error, last_row_count = None, None, 0
268279
for _ in range(repeat):
269280
lat, result, row_count, error = run_query(
270-
sql, endpoint_url, session, timeout, debug,
271-
database=database, api_key=api_key,
281+
sql,
282+
endpoint_url,
283+
session,
284+
timeout,
285+
debug,
286+
database=database,
287+
api_key=api_key,
272288
)
273289
trial_latencies.append(lat)
274290
last_result, last_error, last_row_count = result, error, row_count
@@ -436,13 +452,21 @@ def main():
436452
)
437453

438454
# Elasticsearch flags
439-
es_group = parser.add_argument_group("Elasticsearch options (--database elasticsearch)")
440-
es_group.add_argument("--elastic-host", default=DEFAULT_ELASTIC_HOST,
441-
help="Elasticsearch host")
442-
es_group.add_argument("--elastic-port", type=int, default=DEFAULT_ELASTIC_PORT,
443-
help="Elasticsearch port")
444-
es_group.add_argument("--elastic-api-key", default=None,
445-
help="Elasticsearch API key")
455+
es_group = parser.add_argument_group(
456+
"Elasticsearch options (--database elasticsearch)"
457+
)
458+
es_group.add_argument(
459+
"--elastic-host", default=DEFAULT_ELASTIC_HOST, help="Elasticsearch host"
460+
)
461+
es_group.add_argument(
462+
"--elastic-port",
463+
type=int,
464+
default=DEFAULT_ELASTIC_PORT,
465+
help="Elasticsearch port",
466+
)
467+
es_group.add_argument(
468+
"--elastic-api-key", default=None, help="Elasticsearch API key"
469+
)
446470

447471
# Shared flags
448472
parser.add_argument(
@@ -498,8 +522,8 @@ def main():
498522
if use_elastic
499523
else args.clickhouse_url
500524
)
501-
asap_url = (
502-
args.asap_url or (DEFAULT_ASAP_ELASTIC_URL if use_elastic else DEFAULT_ASAP_CLICKHOUSE_URL)
525+
asap_url = args.asap_url or (
526+
DEFAULT_ASAP_ELASTIC_URL if use_elastic else DEFAULT_ASAP_CLICKHOUSE_URL
503527
)
504528

505529
output_dir = Path(args.output_dir)
@@ -542,7 +566,9 @@ def main():
542566
)
543567

544568
if args.mode == "both" and not args.no_plot:
545-
_plot_comparison(asap_csv, baseline_csv, output_dir / f"{prefix}_comparison.png")
569+
_plot_comparison(
570+
asap_csv, baseline_csv, output_dir / f"{prefix}_comparison.png"
571+
)
546572

547573

548574
if __name__ == "__main__":

0 commit comments

Comments
 (0)