mirror of
https://github.com/openobserve/openobserve.git
synced 2026-11-04 08:09:27 +00:00
feat(search): support two-field GROUP BY via Tantivy index optimization (#12348)
This commit is contained in:
mac119
@@ -0,0 +1,376 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Performance Benchmark: Single GROUP BY vs Two-field GROUP BY via Tantivy
|
||||
========================================================================
|
||||
Tests the optimization from issue #12039.
|
||||
|
||||
Usage:
|
||||
python3 scripts/bench_group_by.py [--base-url URL] [--records N] [--rounds N]
|
||||
|
||||
Phases:
|
||||
1. Create stream with index_fields configured
|
||||
2. Ingest test data (high cardinality fields)
|
||||
3. Wait for data flush & index build
|
||||
4. Run benchmark queries: single vs two-field GROUP BY
|
||||
5. Output comparison report
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import string
|
||||
import sys
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from statistics import mean, median, stdev
|
||||
from urllib.request import Request, urlopen
|
||||
from urllib.error import URLError
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config
|
||||
# ---------------------------------------------------------------------------
|
||||
DEFAULT_BASE_URL = "http://localhost:5080"
|
||||
DEFAULT_USER = "admin@test.com"
|
||||
DEFAULT_PASS = "Admin123!"
|
||||
ORG = "default"
|
||||
STREAM = "bench_group_by"
|
||||
|
||||
# Data generation params
|
||||
USERIDS = [f"user_{i:04d}" for i in range(500)] # 500 unique users
|
||||
SEARCH_PHRASES = [f"phrase_{i:03d}" for i in range(200)] # 200 unique phrases
|
||||
STATUSES = ["200", "301", "400", "403", "404", "500"]
|
||||
METHODS = ["GET", "POST", "PUT", "DELETE", "PATCH"]
|
||||
SERVICES = [f"svc_{i:02d}" for i in range(50)] # 50 services
|
||||
|
||||
|
||||
def make_auth_header(user, password):
|
||||
import base64
|
||||
cred = base64.b64encode(f"{user}:{password}".encode()).decode()
|
||||
return {"Authorization": f"Basic {cred}", "Content-Type": "application/json"}
|
||||
|
||||
|
||||
def api_request(base_url, path, method="GET", data=None, headers=None):
|
||||
url = f"{base_url}{path}"
|
||||
body = json.dumps(data).encode() if data else None
|
||||
req = Request(url, data=body, headers=headers or {}, method=method)
|
||||
try:
|
||||
with urlopen(req, timeout=120) as resp:
|
||||
return resp.status, json.loads(resp.read().decode())
|
||||
except URLError as e:
|
||||
body_text = ""
|
||||
if hasattr(e, "read"):
|
||||
try:
|
||||
body_text = e.read().decode()[:200]
|
||||
except Exception:
|
||||
pass
|
||||
return 0, {"error": str(e), "body": body_text}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 1: Setup stream with index fields
|
||||
# ---------------------------------------------------------------------------
|
||||
def setup_stream(base_url, headers):
|
||||
print("\n=== Phase 1: Setup stream & index fields ===")
|
||||
|
||||
# Update stream settings to enable index on key fields
|
||||
index_fields = ["userid", "searchphrase", "service", "method", "status"]
|
||||
settings = {
|
||||
"index_fields": {"add": index_fields, "remove": []},
|
||||
"full_text_search_keys": {"add": ["message"], "remove": []},
|
||||
}
|
||||
status, resp = api_request(
|
||||
base_url,
|
||||
f"/api/{ORG}/streams/{STREAM}/settings",
|
||||
method="PUT",
|
||||
data=settings,
|
||||
headers=headers,
|
||||
)
|
||||
if status == 200:
|
||||
print(f" Stream settings updated: index_fields = {index_fields}")
|
||||
return True
|
||||
else:
|
||||
print(f" Settings response: {status} {resp}")
|
||||
return False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 2: Generate & ingest data
|
||||
# ---------------------------------------------------------------------------
|
||||
def generate_batch(batch_size, base_ts):
|
||||
records = []
|
||||
for i in range(batch_size):
|
||||
# Keep within last 4 hours to avoid ZO_INGEST_ALLOWED_UPTO rejection
|
||||
ts = base_ts + timedelta(seconds=random.randint(0, 4 * 3600))
|
||||
records.append({
|
||||
"_timestamp": int(ts.timestamp() * 1_000_000),
|
||||
"userid": random.choice(USERIDS),
|
||||
"searchphrase": random.choice(SEARCH_PHRASES),
|
||||
"status": random.choice(STATUSES),
|
||||
"method": random.choice(METHODS),
|
||||
"service": random.choice(SERVICES),
|
||||
"response_time": random.randint(1, 5000),
|
||||
"message": f"request from {random.choice(USERIDS)} searching {random.choice(SEARCH_PHRASES)}",
|
||||
})
|
||||
return records
|
||||
|
||||
|
||||
def ingest_data(base_url, headers, total_records, batch_size=5000):
|
||||
print(f"\n=== Phase 2: Ingesting {total_records:,} records ===")
|
||||
base_ts = datetime.now(timezone.utc) - timedelta(hours=4)
|
||||
ingested = 0
|
||||
start = time.time()
|
||||
|
||||
# First batch: create the stream
|
||||
first_batch = min(batch_size, total_records)
|
||||
records = generate_batch(first_batch, base_ts)
|
||||
status, resp = api_request(
|
||||
base_url,
|
||||
f"/api/{ORG}/{STREAM}/_json",
|
||||
method="POST",
|
||||
data=records,
|
||||
headers=headers,
|
||||
)
|
||||
if status == 200:
|
||||
successful = resp.get("status", [{}])[0].get("successful", 0)
|
||||
failed = resp.get("status", [{}])[0].get("failed", 0)
|
||||
if failed > 0:
|
||||
error = resp.get("status", [{}])[0].get("error", "")
|
||||
print(f" WARNING: first batch {successful} ok, {failed} failed: {error}")
|
||||
ingested += first_batch
|
||||
|
||||
# Now configure index fields (stream must exist) - retry up to 5 times
|
||||
for attempt in range(5):
|
||||
time.sleep(1) # Wait a moment for stream to be fully created
|
||||
if setup_stream(base_url, headers):
|
||||
break
|
||||
print(f" Retrying settings ({attempt+2}/5)...")
|
||||
|
||||
# Continue ingesting remaining data
|
||||
while ingested < total_records:
|
||||
this_batch = min(batch_size, total_records - ingested)
|
||||
records = generate_batch(this_batch, base_ts)
|
||||
status, resp = api_request(
|
||||
base_url,
|
||||
f"/api/{ORG}/{STREAM}/_json",
|
||||
method="POST",
|
||||
data=records,
|
||||
headers=headers,
|
||||
)
|
||||
ingested += this_batch
|
||||
elapsed = time.time() - start
|
||||
rate = ingested / elapsed if elapsed > 0 else 0
|
||||
print(f"\r Ingested: {ingested:>10,} / {total_records:,} ({rate:,.0f} rec/s)", end="", flush=True)
|
||||
|
||||
elapsed = time.time() - start
|
||||
print(f"\n Done in {elapsed:.1f}s ({ingested/elapsed:,.0f} rec/s)")
|
||||
return ingested
|
||||
|
||||
|
||||
def flush_data(base_url, headers):
|
||||
print("\n=== Phase 3: Flushing data & waiting for index build ===")
|
||||
api_request(base_url, "/node/flush", method="PUT", headers=headers)
|
||||
print(" Flush requested, waiting 30s for compaction + index build...")
|
||||
for i in range(30, 0, -1):
|
||||
print(f"\r Waiting... {i:2d}s", end="", flush=True)
|
||||
time.sleep(1)
|
||||
print("\r Ready. ")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 4: Benchmark queries
|
||||
# ---------------------------------------------------------------------------
|
||||
def run_query(base_url, headers, sql, start_time, end_time):
|
||||
"""Run a search query and return (elapsed_ms, num_hits, scan_size)."""
|
||||
payload = {
|
||||
"query": {
|
||||
"sql": sql,
|
||||
"start_time": start_time,
|
||||
"end_time": end_time,
|
||||
"from": 0,
|
||||
"size": 1000,
|
||||
}
|
||||
}
|
||||
t0 = time.time()
|
||||
status, resp = api_request(
|
||||
base_url,
|
||||
f"/api/{ORG}/_search?type=logs",
|
||||
method="POST",
|
||||
data=payload,
|
||||
headers=headers,
|
||||
)
|
||||
elapsed_ms = (time.time() - t0) * 1000
|
||||
|
||||
hits = 0
|
||||
scan_size = 0
|
||||
if status == 200:
|
||||
hits = resp.get("total", 0)
|
||||
scan_size = resp.get("scan_size", 0)
|
||||
else:
|
||||
print(f"\n Query error: {status} {resp}")
|
||||
|
||||
return elapsed_ms, hits, scan_size
|
||||
|
||||
|
||||
def benchmark(base_url, headers, rounds):
|
||||
print(f"\n=== Phase 4: Running benchmark ({rounds} rounds per query) ===\n")
|
||||
|
||||
now = datetime.now(timezone.utc)
|
||||
start_us = int((now - timedelta(hours=5)).timestamp() * 1_000_000)
|
||||
end_us = int(now.timestamp() * 1_000_000)
|
||||
|
||||
queries = {
|
||||
# --- Single GROUP BY ---
|
||||
"1-field GROUP BY (userid, LIMIT 10)": (
|
||||
f'SELECT userid, COUNT(*) as cnt FROM "{STREAM}" '
|
||||
f"GROUP BY userid ORDER BY cnt DESC LIMIT 10"
|
||||
),
|
||||
"1-field GROUP BY (searchphrase, LIMIT 10)": (
|
||||
f'SELECT searchphrase, COUNT(*) as cnt FROM "{STREAM}" '
|
||||
f"GROUP BY searchphrase ORDER BY cnt DESC LIMIT 10"
|
||||
),
|
||||
"1-field GROUP BY (service, LIMIT 20)": (
|
||||
f'SELECT service, COUNT(*) as cnt FROM "{STREAM}" '
|
||||
f"GROUP BY service ORDER BY cnt DESC LIMIT 20"
|
||||
),
|
||||
# --- Two-field GROUP BY ---
|
||||
"2-field GROUP BY (userid+searchphrase, LIMIT 10)": (
|
||||
f'SELECT userid, searchphrase, COUNT(*) as cnt FROM "{STREAM}" '
|
||||
f"GROUP BY userid, searchphrase ORDER BY cnt DESC LIMIT 10"
|
||||
),
|
||||
"2-field GROUP BY (userid+service, LIMIT 10)": (
|
||||
f'SELECT userid, service, COUNT(*) as cnt FROM "{STREAM}" '
|
||||
f"GROUP BY userid, service ORDER BY cnt DESC LIMIT 10"
|
||||
),
|
||||
"2-field GROUP BY (service+method, LIMIT 20)": (
|
||||
f'SELECT service, method, COUNT(*) as cnt FROM "{STREAM}" '
|
||||
f"GROUP BY service, method ORDER BY cnt DESC LIMIT 20"
|
||||
),
|
||||
# --- Baseline: full scan aggregation (non-indexed field) ---
|
||||
"Baseline: GROUP BY non-indexed (response_time buckets)": (
|
||||
f'SELECT CASE WHEN response_time < 100 THEN \'fast\' '
|
||||
f"WHEN response_time < 1000 THEN 'medium' ELSE 'slow' END as speed, "
|
||||
f'COUNT(*) as cnt FROM "{STREAM}" GROUP BY speed ORDER BY cnt DESC LIMIT 10'
|
||||
),
|
||||
}
|
||||
|
||||
results = {}
|
||||
for label, sql in queries.items():
|
||||
print(f" {label}")
|
||||
print(f" SQL: {sql[:100]}...")
|
||||
latencies = []
|
||||
for r in range(rounds):
|
||||
ms, hits, scan_size = run_query(base_url, headers, sql, start_us, end_us)
|
||||
latencies.append(ms)
|
||||
print(f" Round {r+1}: {ms:>8.1f} ms (hits={hits})", flush=True)
|
||||
|
||||
results[label] = {
|
||||
"latencies": latencies,
|
||||
"min": min(latencies),
|
||||
"max": max(latencies),
|
||||
"mean": mean(latencies),
|
||||
"median": median(latencies),
|
||||
"stdev": stdev(latencies) if len(latencies) > 1 else 0,
|
||||
}
|
||||
print()
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase 5: Report
|
||||
# ---------------------------------------------------------------------------
|
||||
def print_report(results, total_records):
|
||||
print("\n" + "=" * 90)
|
||||
print(f" BENCHMARK REPORT ({total_records:,} records)")
|
||||
print("=" * 90)
|
||||
print(f"{'Query':<55} {'Min':>8} {'Mean':>8} {'Median':>8} {'Max':>8} {'StdDev':>8}")
|
||||
print("-" * 90)
|
||||
|
||||
for label, r in results.items():
|
||||
print(
|
||||
f"{label:<55} "
|
||||
f"{r['min']:>7.1f}ms "
|
||||
f"{r['mean']:>7.1f}ms "
|
||||
f"{r['median']:>7.1f}ms "
|
||||
f"{r['max']:>7.1f}ms "
|
||||
f"{r['stdev']:>7.1f}ms"
|
||||
)
|
||||
|
||||
# Comparison section
|
||||
print("\n" + "-" * 90)
|
||||
print(" COMPARISON: 1-field vs 2-field GROUP BY")
|
||||
print("-" * 90)
|
||||
|
||||
single_keys = [k for k in results if k.startswith("1-field")]
|
||||
double_keys = [k for k in results if k.startswith("2-field")]
|
||||
|
||||
if single_keys and double_keys:
|
||||
avg_single = mean([results[k]["median"] for k in single_keys])
|
||||
avg_double = mean([results[k]["median"] for k in double_keys])
|
||||
ratio = avg_double / avg_single if avg_single > 0 else float("inf")
|
||||
print(f" Avg median latency (1-field): {avg_single:>8.1f} ms")
|
||||
print(f" Avg median latency (2-field): {avg_double:>8.1f} ms")
|
||||
print(f" Ratio (2-field / 1-field): {ratio:>8.2f}x")
|
||||
|
||||
baseline_keys = [k for k in results if k.startswith("Baseline")]
|
||||
if baseline_keys and double_keys:
|
||||
avg_baseline = mean([results[k]["median"] for k in baseline_keys])
|
||||
avg_double = mean([results[k]["median"] for k in double_keys])
|
||||
speedup = avg_baseline / avg_double if avg_double > 0 else float("inf")
|
||||
print(f"\n Avg median latency (baseline full-scan): {avg_baseline:>8.1f} ms")
|
||||
print(f" Avg median latency (2-field tantivy): {avg_double:>8.1f} ms")
|
||||
print(f" Speedup (tantivy vs full-scan): {speedup:>8.2f}x")
|
||||
|
||||
print("=" * 90)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Main
|
||||
# ---------------------------------------------------------------------------
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Benchmark: GROUP BY via Tantivy")
|
||||
parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
|
||||
parser.add_argument("--user", default=DEFAULT_USER)
|
||||
parser.add_argument("--password", default=DEFAULT_PASS)
|
||||
parser.add_argument("--records", type=int, default=500_000, help="Total records to ingest")
|
||||
parser.add_argument("--batch-size", type=int, default=5000, help="Batch size per request")
|
||||
parser.add_argument("--rounds", type=int, default=5, help="Rounds per query")
|
||||
parser.add_argument("--skip-ingest", action="store_true", help="Skip data ingestion")
|
||||
args = parser.parse_args()
|
||||
|
||||
headers = make_auth_header(args.user, args.password)
|
||||
|
||||
# Health check
|
||||
status, _ = api_request(args.base_url, "/healthz", headers=headers)
|
||||
if status != 200:
|
||||
print(f"ERROR: OpenObserve not reachable at {args.base_url}")
|
||||
sys.exit(1)
|
||||
print(f"OpenObserve is running at {args.base_url}")
|
||||
|
||||
if not args.skip_ingest:
|
||||
ingest_data(args.base_url, headers, args.records, args.batch_size)
|
||||
flush_data(args.base_url, headers)
|
||||
else:
|
||||
print("\n Skipping ingestion (--skip-ingest)")
|
||||
|
||||
results = benchmark(args.base_url, headers, args.rounds)
|
||||
print_report(results, args.records)
|
||||
|
||||
# Save raw results
|
||||
out_file = "scripts/bench_group_by_results.json"
|
||||
with open(out_file, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||||
"total_records": args.records,
|
||||
"rounds": args.rounds,
|
||||
"results": {k: {kk: vv for kk, vv in v.items()} for k, v in results.items()},
|
||||
},
|
||||
f,
|
||||
indent=2,
|
||||
)
|
||||
print(f"\nRaw results saved to {out_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,184 @@
|
||||
<!DOCTYPE html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8">
|
||||
<title>GROUP BY Benchmark - Performance Chart</title>
|
||||
<script src="https://cdn.jsdelivr.net/npm/chart.js@4.4.0/dist/chart.umd.min.js"></script>
|
||||
<style>
|
||||
body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; margin: 0; padding: 20px; background: #f5f7fa; }
|
||||
h1 { text-align: center; color: #333; margin-bottom: 5px; }
|
||||
.subtitle { text-align: center; color: #666; margin-bottom: 30px; font-size: 14px; }
|
||||
.chart-container { background: white; border-radius: 12px; padding: 24px; margin: 0 auto 30px; max-width: 1100px; box-shadow: 0 2px 12px rgba(0,0,0,0.08); }
|
||||
.chart-container h2 { margin-top: 0; color: #444; font-size: 16px; }
|
||||
canvas { width: 100% !important; }
|
||||
.summary { background: white; border-radius: 12px; padding: 24px; margin: 0 auto; max-width: 1100px; box-shadow: 0 2px 12px rgba(0,0,0,0.08); }
|
||||
.summary table { width: 100%; border-collapse: collapse; font-size: 14px; }
|
||||
.summary th, .summary td { padding: 10px 12px; text-align: left; border-bottom: 1px solid #eee; }
|
||||
.summary th { background: #f8f9fa; font-weight: 600; color: #555; }
|
||||
.summary tr:hover { background: #f0f7ff; }
|
||||
.tag { display: inline-block; padding: 2px 8px; border-radius: 4px; font-size: 12px; font-weight: 500; }
|
||||
.tag-1f { background: #e3f2fd; color: #1565c0; }
|
||||
.tag-2f { background: #fff3e0; color: #e65100; }
|
||||
.tag-bl { background: #f3e5f5; color: #7b1fa2; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<h1>GROUP BY Performance Benchmark</h1>
|
||||
<div class="subtitle">10,000,000 records · 5 rounds per query · OpenObserve + Tantivy</div>
|
||||
|
||||
<div class="chart-container">
|
||||
<h2>Latency per Round (all queries)</h2>
|
||||
<canvas id="lineChart" height="100"></canvas>
|
||||
</div>
|
||||
|
||||
<div class="chart-container">
|
||||
<h2>Median Latency Comparison</h2>
|
||||
<canvas id="barChart" height="60"></canvas>
|
||||
</div>
|
||||
|
||||
<div class="summary">
|
||||
<h2 style="margin-top:0; color:#444;">Summary Table</h2>
|
||||
<table id="summaryTable">
|
||||
<thead>
|
||||
<tr><th>Query</th><th>Type</th><th>Min</th><th>Median</th><th>Mean</th><th>Max</th><th>StdDev</th></tr>
|
||||
</thead>
|
||||
<tbody></tbody>
|
||||
</table>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
const data = {
|
||||
"1-field GROUP BY (userid, LIMIT 10)": {
|
||||
"latencies": [130.56, 82.82, 88.39, 89.55, 83.41],
|
||||
"min": 82.82, "max": 130.56, "mean": 94.95, "median": 88.39, "stdev": 20.13
|
||||
},
|
||||
"1-field GROUP BY (searchphrase, LIMIT 10)": {
|
||||
"latencies": [161.52, 79.49, 84.71, 84.64, 87.53],
|
||||
"min": 79.49, "max": 161.52, "mean": 99.58, "median": 84.71, "stdev": 34.75
|
||||
},
|
||||
"1-field GROUP BY (service, LIMIT 20)": {
|
||||
"latencies": [151.94, 80.97, 85.08, 85.87, 78.16],
|
||||
"min": 78.16, "max": 151.94, "mean": 96.40, "median": 85.08, "stdev": 31.20
|
||||
},
|
||||
"2-field GROUP BY (userid+searchphrase, LIMIT 10)": {
|
||||
"latencies": [188.55, 153.51, 152.81, 148.32, 145.67],
|
||||
"min": 145.67, "max": 188.55, "mean": 157.77, "median": 152.81, "stdev": 17.51
|
||||
},
|
||||
"2-field GROUP BY (userid+service, LIMIT 10)": {
|
||||
"latencies": [159.39, 108.85, 105.67, 110.62, 103.72],
|
||||
"min": 103.72, "max": 159.39, "mean": 117.65, "median": 108.85, "stdev": 23.49
|
||||
},
|
||||
"2-field GROUP BY (service+method, LIMIT 20)": {
|
||||
"latencies": [152.64, 87.91, 97.47, 93.10, 96.85],
|
||||
"min": 87.91, "max": 152.64, "mean": 105.59, "median": 96.85, "stdev": 26.57
|
||||
},
|
||||
"Baseline: full-scan (non-indexed)": {
|
||||
"latencies": [158.97, 121.34, 115.34, 111.15, 106.62],
|
||||
"min": 106.62, "max": 158.97, "mean": 122.68, "median": 115.34, "stdev": 21.00
|
||||
}
|
||||
};
|
||||
|
||||
const colors = {
|
||||
'1-field': ['#1976D2', '#42A5F5', '#90CAF9'],
|
||||
'2-field': ['#E65100', '#FF9800', '#FFCC80'],
|
||||
'Baseline': ['#7B1FA2']
|
||||
};
|
||||
|
||||
const labels = ['Round 1', 'Round 2', 'Round 3', 'Round 4', 'Round 5'];
|
||||
const datasets = [];
|
||||
let colorIdx1 = 0, colorIdx2 = 0;
|
||||
|
||||
Object.entries(data).forEach(([key, val]) => {
|
||||
let color, dash;
|
||||
if (key.startsWith('1-field')) {
|
||||
color = colors['1-field'][colorIdx1++];
|
||||
dash = [];
|
||||
} else if (key.startsWith('2-field')) {
|
||||
color = colors['2-field'][colorIdx2++];
|
||||
dash = [5, 5];
|
||||
} else {
|
||||
color = colors['Baseline'][0];
|
||||
dash = [10, 5];
|
||||
}
|
||||
datasets.push({
|
||||
label: key,
|
||||
data: val.latencies,
|
||||
borderColor: color,
|
||||
backgroundColor: color + '20',
|
||||
borderWidth: 2.5,
|
||||
borderDash: dash,
|
||||
pointRadius: 4,
|
||||
pointHoverRadius: 6,
|
||||
tension: 0.3,
|
||||
fill: false
|
||||
});
|
||||
});
|
||||
|
||||
// Line chart
|
||||
new Chart(document.getElementById('lineChart'), {
|
||||
type: 'line',
|
||||
data: { labels, datasets },
|
||||
options: {
|
||||
responsive: true,
|
||||
plugins: {
|
||||
legend: { position: 'bottom', labels: { usePointStyle: true, padding: 15 } },
|
||||
tooltip: { callbacks: { label: ctx => `${ctx.dataset.label}: ${ctx.parsed.y.toFixed(1)} ms` } }
|
||||
},
|
||||
scales: {
|
||||
y: { title: { display: true, text: 'Latency (ms)' }, beginAtZero: false },
|
||||
x: { title: { display: true, text: 'Round' } }
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Bar chart - median comparison
|
||||
const barLabels = Object.keys(data);
|
||||
const barValues = Object.values(data).map(v => v.median);
|
||||
const barColors = barLabels.map(l => {
|
||||
if (l.startsWith('1-field')) return '#1976D2';
|
||||
if (l.startsWith('2-field')) return '#E65100';
|
||||
return '#7B1FA2';
|
||||
});
|
||||
|
||||
new Chart(document.getElementById('barChart'), {
|
||||
type: 'bar',
|
||||
data: {
|
||||
labels: barLabels.map(l => l.replace('GROUP BY ', '').replace('Baseline: full-scan (non-indexed)', 'Baseline (full-scan)')),
|
||||
datasets: [{
|
||||
label: 'Median Latency (ms)',
|
||||
data: barValues,
|
||||
backgroundColor: barColors.map(c => c + 'CC'),
|
||||
borderColor: barColors,
|
||||
borderWidth: 1.5,
|
||||
borderRadius: 6
|
||||
}]
|
||||
},
|
||||
options: {
|
||||
responsive: true,
|
||||
plugins: {
|
||||
legend: { display: false },
|
||||
tooltip: { callbacks: { label: ctx => `${ctx.parsed.y.toFixed(1)} ms` } }
|
||||
},
|
||||
scales: {
|
||||
y: { title: { display: true, text: 'Median Latency (ms)' }, beginAtZero: true },
|
||||
x: { ticks: { maxRotation: 20 } }
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
// Summary table
|
||||
const tbody = document.querySelector('#summaryTable tbody');
|
||||
Object.entries(data).forEach(([key, val]) => {
|
||||
const type = key.startsWith('1-field') ? '<span class="tag tag-1f">1-field</span>' :
|
||||
key.startsWith('2-field') ? '<span class="tag tag-2f">2-field</span>' :
|
||||
'<span class="tag tag-bl">Baseline</span>';
|
||||
tbody.innerHTML += `<tr>
|
||||
<td>${key}</td><td>${type}</td>
|
||||
<td>${val.min.toFixed(1)} ms</td><td><strong>${val.median.toFixed(1)} ms</strong></td>
|
||||
<td>${val.mean.toFixed(1)} ms</td><td>${val.max.toFixed(1)} ms</td><td>${val.stdev.toFixed(1)} ms</td>
|
||||
</tr>`;
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,105 @@
|
||||
{
|
||||
"timestamp": "2026-05-29T07:28:07.919275+00:00",
|
||||
"total_records": 10000000,
|
||||
"rounds": 5,
|
||||
"results": {
|
||||
"1-field GROUP BY (userid, LIMIT 10)": {
|
||||
"latencies": [
|
||||
130.55777549743652,
|
||||
82.8249454498291,
|
||||
88.39106559753418,
|
||||
89.55097198486328,
|
||||
83.40668678283691
|
||||
],
|
||||
"min": 82.8249454498291,
|
||||
"max": 130.55777549743652,
|
||||
"mean": 94.9462890625,
|
||||
"median": 88.39106559753418,
|
||||
"stdev": 20.12677171641535
|
||||
},
|
||||
"1-field GROUP BY (searchphrase, LIMIT 10)": {
|
||||
"latencies": [
|
||||
161.5164279937744,
|
||||
79.48851585388184,
|
||||
84.70940589904785,
|
||||
84.6409797668457,
|
||||
87.53085136413574
|
||||
],
|
||||
"min": 79.48851585388184,
|
||||
"max": 161.5164279937744,
|
||||
"mean": 99.57723617553711,
|
||||
"median": 84.70940589904785,
|
||||
"stdev": 34.746509489872174
|
||||
},
|
||||
"1-field GROUP BY (service, LIMIT 20)": {
|
||||
"latencies": [
|
||||
151.93629264831543,
|
||||
80.97243309020996,
|
||||
85.07585525512695,
|
||||
85.87050437927246,
|
||||
78.1550407409668
|
||||
],
|
||||
"min": 78.1550407409668,
|
||||
"max": 151.93629264831543,
|
||||
"mean": 96.40202522277832,
|
||||
"median": 85.07585525512695,
|
||||
"stdev": 31.20206191889121
|
||||
},
|
||||
"2-field GROUP BY (userid+searchphrase, LIMIT 10)": {
|
||||
"latencies": [
|
||||
188.5545253753662,
|
||||
153.50747108459473,
|
||||
152.80795097351074,
|
||||
148.31876754760742,
|
||||
145.67303657531738
|
||||
],
|
||||
"min": 145.67303657531738,
|
||||
"max": 188.5545253753662,
|
||||
"mean": 157.7723503112793,
|
||||
"median": 152.80795097351074,
|
||||
"stdev": 17.508143656809874
|
||||
},
|
||||
"2-field GROUP BY (userid+service, LIMIT 10)": {
|
||||
"latencies": [
|
||||
159.39092636108398,
|
||||
108.85119438171387,
|
||||
105.6673526763916,
|
||||
110.61620712280273,
|
||||
103.71589660644531
|
||||
],
|
||||
"min": 103.71589660644531,
|
||||
"max": 159.39092636108398,
|
||||
"mean": 117.6483154296875,
|
||||
"median": 108.85119438171387,
|
||||
"stdev": 23.48904646906988
|
||||
},
|
||||
"2-field GROUP BY (service+method, LIMIT 20)": {
|
||||
"latencies": [
|
||||
152.63843536376953,
|
||||
87.91232109069824,
|
||||
97.46885299682617,
|
||||
93.09720993041992,
|
||||
96.84538841247559
|
||||
],
|
||||
"min": 87.91232109069824,
|
||||
"max": 152.63843536376953,
|
||||
"mean": 105.59244155883789,
|
||||
"median": 96.84538841247559,
|
||||
"stdev": 26.57323200224028
|
||||
},
|
||||
"Baseline: GROUP BY non-indexed (response_time buckets)": {
|
||||
"latencies": [
|
||||
158.96916389465332,
|
||||
121.33550643920898,
|
||||
115.3421401977539,
|
||||
111.15288734436035,
|
||||
106.62150382995605
|
||||
],
|
||||
"min": 106.62150382995605,
|
||||
"max": 158.96916389465332,
|
||||
"mean": 122.68424034118652,
|
||||
"median": 115.3421401977539,
|
||||
"stdev": 20.995861503532634
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -24,6 +24,8 @@ pub enum IndexOptimizeMode {
|
||||
SimpleHistogram(i64, u64, usize),
|
||||
SimpleMultiHistogram(i64, i64, u64, String),
|
||||
SimpleTopN(String, usize, bool),
|
||||
/// Two-field GROUP BY with count(*): (field1, field2, limit, ascend)
|
||||
SimpleTopNMulti(Vec<String>, usize, bool),
|
||||
SimpleDistinct(String, usize, bool),
|
||||
}
|
||||
|
||||
@@ -41,6 +43,10 @@ impl IndexOptimizeMode {
|
||||
IndexOptimizeMode::SimpleTopN(field, limit, ascend) => {
|
||||
format!("t(f{field},l:{limit},a:{ascend})")
|
||||
}
|
||||
IndexOptimizeMode::SimpleTopNMulti(fields, limit, ascend) => {
|
||||
let fields_str = fields.join(",");
|
||||
format!("tm(f:{fields_str},l:{limit},a:{ascend})")
|
||||
}
|
||||
IndexOptimizeMode::SimpleDistinct(field, limit, ascend) => {
|
||||
format!("d(f:{field},l:{limit},a:{ascend})")
|
||||
}
|
||||
@@ -70,6 +76,13 @@ impl std::fmt::Display for IndexOptimizeMode {
|
||||
IndexOptimizeMode::SimpleTopN(field, limit, ascend) => {
|
||||
write!(f, "topn(field: {field}, limit: {limit}, ascend: {ascend})")
|
||||
}
|
||||
IndexOptimizeMode::SimpleTopNMulti(fields, limit, ascend) => {
|
||||
let fields_str = fields.join(", ");
|
||||
write!(
|
||||
f,
|
||||
"topn_multi(fields: [{fields_str}], limit: {limit}, ascend: {ascend})"
|
||||
)
|
||||
}
|
||||
IndexOptimizeMode::SimpleDistinct(field, limit, ascend) => {
|
||||
write!(
|
||||
f,
|
||||
@@ -89,6 +102,9 @@ impl From<cluster_rpc::IdxOptimizeMode> for IndexOptimizeMode {
|
||||
Some(cluster_rpc::idx_optimize_mode::Mode::SimpleDistinct(select)) => {
|
||||
IndexOptimizeMode::SimpleDistinct(select.field, select.limit as usize, select.asc)
|
||||
}
|
||||
Some(cluster_rpc::idx_optimize_mode::Mode::SimpleTopnMulti(select)) => {
|
||||
IndexOptimizeMode::SimpleTopNMulti(select.fields, select.limit as usize, select.asc)
|
||||
}
|
||||
None => panic!("Invalid IndexOptimizeMode"),
|
||||
}
|
||||
}
|
||||
@@ -106,6 +122,17 @@ impl From<IndexOptimizeMode> for cluster_rpc::IdxOptimizeMode {
|
||||
},
|
||||
)),
|
||||
},
|
||||
IndexOptimizeMode::SimpleTopNMulti(fields, limit, asc) => {
|
||||
cluster_rpc::IdxOptimizeMode {
|
||||
mode: Some(cluster_rpc::idx_optimize_mode::Mode::SimpleTopnMulti(
|
||||
cluster_rpc::SimpleTopNMulti {
|
||||
fields,
|
||||
limit: limit as u32,
|
||||
asc,
|
||||
},
|
||||
)),
|
||||
}
|
||||
}
|
||||
IndexOptimizeMode::SimpleDistinct(field, limit, asc) => cluster_rpc::IdxOptimizeMode {
|
||||
mode: Some(cluster_rpc::idx_optimize_mode::Mode::SimpleDistinct(
|
||||
cluster_rpc::SimpleDistinct {
|
||||
|
||||
@@ -120,6 +120,7 @@ message IdxOptimizeMode {
|
||||
oneof mode {
|
||||
SimpleTopN simple_topn = 4;
|
||||
SimpleDistinct simple_distinct = 5;
|
||||
SimpleTopNMulti simple_topn_multi = 6;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -129,6 +130,12 @@ message SimpleTopN {
|
||||
bool asc = 3;
|
||||
}
|
||||
|
||||
message SimpleTopNMulti {
|
||||
repeated string fields = 1;
|
||||
uint32 limit = 2;
|
||||
bool asc = 3;
|
||||
}
|
||||
|
||||
message SimpleDistinct {
|
||||
string field = 1;
|
||||
uint32 limit = 2;
|
||||
|
||||
@@ -3354,7 +3354,7 @@ pub struct SuperClusterInfo {
|
||||
}
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct IdxOptimizeMode {
|
||||
#[prost(oneof = "idx_optimize_mode::Mode", tags = "4, 5")]
|
||||
#[prost(oneof = "idx_optimize_mode::Mode", tags = "4, 5, 6")]
|
||||
pub mode: ::core::option::Option<idx_optimize_mode::Mode>,
|
||||
}
|
||||
/// Nested message and enum types in `IdxOptimizeMode`.
|
||||
@@ -3365,6 +3365,8 @@ pub mod idx_optimize_mode {
|
||||
SimpleTopn(super::SimpleTopN),
|
||||
#[prost(message, tag = "5")]
|
||||
SimpleDistinct(super::SimpleDistinct),
|
||||
#[prost(message, tag = "6")]
|
||||
SimpleTopnMulti(super::SimpleTopNMulti),
|
||||
}
|
||||
}
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
@@ -3377,6 +3379,15 @@ pub struct SimpleTopN {
|
||||
pub asc: bool,
|
||||
}
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct SimpleTopNMulti {
|
||||
#[prost(string, repeated, tag = "1")]
|
||||
pub fields: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
|
||||
#[prost(uint32, tag = "2")]
|
||||
pub limit: u32,
|
||||
#[prost(bool, tag = "3")]
|
||||
pub asc: bool,
|
||||
}
|
||||
#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)]
|
||||
pub struct SimpleDistinct {
|
||||
#[prost(string, tag = "1")]
|
||||
pub field: ::prost::alloc::string::String,
|
||||
|
||||
@@ -39,6 +39,7 @@ mod distinct;
|
||||
mod histogram;
|
||||
mod select;
|
||||
mod topn;
|
||||
mod topn_multi;
|
||||
mod utils;
|
||||
|
||||
use crate::service::search::datafusion::{
|
||||
@@ -49,6 +50,7 @@ use crate::service::search::datafusion::{
|
||||
histogram::{is_simple_histogram, is_simple_multi_histogram},
|
||||
select::is_simple_select,
|
||||
topn::is_simple_topn,
|
||||
topn_multi::is_simple_topn_multi,
|
||||
utils::is_complex_plan,
|
||||
},
|
||||
};
|
||||
@@ -145,13 +147,18 @@ impl TreeNodeRewriter for FollowerIndexOptimizer {
|
||||
return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop));
|
||||
}
|
||||
|
||||
// check if the query is simple topn or simple distinct
|
||||
// check if the query is simple topn, topn_multi or simple distinct
|
||||
if config::cluster::LOCAL_NODE.is_single_node() {
|
||||
if let Some(index_optimize_mode) =
|
||||
is_simple_topn(Arc::clone(&plan), self.index_fields.clone())
|
||||
{
|
||||
*self.index_optimizer_mode.lock() = Some(index_optimize_mode);
|
||||
return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop));
|
||||
} else if let Some(index_optimize_mode) =
|
||||
is_simple_topn_multi(Arc::clone(&plan), self.index_fields.clone())
|
||||
{
|
||||
*self.index_optimizer_mode.lock() = Some(index_optimize_mode);
|
||||
return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop));
|
||||
} else if let Some(index_optimize_mode) =
|
||||
is_simple_distinct(Arc::clone(&plan), self.index_fields.clone())
|
||||
{
|
||||
@@ -269,6 +276,13 @@ impl TreeNodeRewriter for LeaderIndexOptimizer {
|
||||
let mut rewriter = IndexOptimizerRewrite::new(index_optimize_mode);
|
||||
let plan = plan.rewrite(&mut rewriter)?.data;
|
||||
return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop));
|
||||
} else if let Some(index_optimize_mode) =
|
||||
is_simple_topn_multi(Arc::clone(&plan), index_fields.clone())
|
||||
{
|
||||
// Check for SimpleTopNMulti (two-field GROUP BY)
|
||||
let mut rewriter = IndexOptimizerRewrite::new(index_optimize_mode);
|
||||
let plan = plan.rewrite(&mut rewriter)?.data;
|
||||
return Ok(Transformed::new(plan, true, TreeNodeRecursion::Stop));
|
||||
} else if let Some(index_optimize_mode) =
|
||||
is_simple_distinct(Arc::clone(&plan), index_fields.clone())
|
||||
{
|
||||
|
||||
+229
@@ -0,0 +1,229 @@
|
||||
// Copyright 2026 OpenObserve Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
|
||||
use config::meta::inverted_index::IndexOptimizeMode;
|
||||
use datafusion::{
|
||||
common::{
|
||||
Result,
|
||||
tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor},
|
||||
},
|
||||
physical_plan::{
|
||||
ExecutionPlan, aggregates::AggregateExec, projection::ProjectionExec,
|
||||
sorts::sort_preserving_merge::SortPreservingMergeExec,
|
||||
},
|
||||
};
|
||||
|
||||
use crate::service::search::datafusion::optimizer::physical_optimizer::{
|
||||
index_optimizer::utils::is_complex_plan,
|
||||
utils::{get_column_name, is_column},
|
||||
};
|
||||
|
||||
#[rustfmt::skip]
|
||||
/// SimpleTopNMulti(Vec<String>, usize, bool):
|
||||
/// Supports two-field GROUP BY queries like:
|
||||
/// SELECT userid, searchphrase, COUNT(*) as cnt FROM hits GROUP BY userid, searchphrase ORDER BY cnt DESC LIMIT 10;
|
||||
/// SELECT userid, searchphrase, COUNT(*) as cnt FROM hits GROUP BY userid, searchphrase LIMIT 10;
|
||||
/// condition: select two index_fields and only have count(*) as aggregate function,
|
||||
/// group by both index_fields in select clause, order by count(*), and have limit
|
||||
pub(crate) fn is_simple_topn_multi(plan: Arc<dyn ExecutionPlan>, index_fields: HashSet<String>) -> Option<IndexOptimizeMode> {
|
||||
let mut visitor = SimpleTopnMultiVisitor::new(index_fields);
|
||||
let _ = plan.visit(&mut visitor);
|
||||
if let Some((fields, fetch, ascend)) = visitor.simple_topn_multi {
|
||||
if fields.is_empty() || fields.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
Some(IndexOptimizeMode::SimpleTopNMulti(fields, fetch, ascend))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
struct SimpleTopnMultiVisitor {
|
||||
pub simple_topn_multi: Option<(Vec<String>, usize, bool)>,
|
||||
index_fields: HashSet<String>,
|
||||
}
|
||||
|
||||
impl SimpleTopnMultiVisitor {
|
||||
pub fn new(index_fields: HashSet<String>) -> Self {
|
||||
Self {
|
||||
simple_topn_multi: None,
|
||||
index_fields,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<'n> TreeNodeVisitor<'n> for SimpleTopnMultiVisitor {
|
||||
type Node = Arc<dyn ExecutionPlan>;
|
||||
|
||||
fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
|
||||
if let Some(sort_merge) = node.as_any().downcast_ref::<SortPreservingMergeExec>() {
|
||||
// Check if this is a SortPreservingMergeExec with fetch limit
|
||||
if let Some(fetch) = sort_merge.fetch()
|
||||
&& fetch > 0
|
||||
&& !sort_merge.expr().is_empty()
|
||||
&& sort_merge.expr().len() <= 3
|
||||
// primary sort + up to 2 secondary sorts
|
||||
{
|
||||
// the first sort should be on the count(*) result, not directly on an index field
|
||||
let sort_expr = sort_merge.expr().first();
|
||||
let column_name = get_column_name(&sort_expr.expr);
|
||||
if is_column(&sort_expr.expr) && self.index_fields.contains(column_name) {
|
||||
self.simple_topn_multi = None;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
}
|
||||
|
||||
self.simple_topn_multi = Some((
|
||||
vec![], // Will be set when we find the group by fields
|
||||
fetch,
|
||||
!sort_merge.expr().first().options.descending,
|
||||
));
|
||||
return Ok(TreeNodeRecursion::Continue);
|
||||
}
|
||||
self.simple_topn_multi = None;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
} else if let Some(projection) = node.as_any().downcast_ref::<ProjectionExec>() {
|
||||
// Check ProjectionExec for the structure: [field1, field2, count(*)]
|
||||
let exprs = projection.expr();
|
||||
if exprs.len() == 3 {
|
||||
// Three expressions: two fields + count(*)
|
||||
return Ok(TreeNodeRecursion::Continue);
|
||||
}
|
||||
self.simple_topn_multi = None;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
} else if let Some(aggregate) = node.as_any().downcast_ref::<AggregateExec>() {
|
||||
// Check if the AggregateExec matches multi-field TopN pattern
|
||||
if aggregate.group_expr().expr().len() == 2 && aggregate.aggr_expr().len() == 1 {
|
||||
let mut fields = Vec::with_capacity(2);
|
||||
|
||||
// Check both group by fields are in index_fields
|
||||
for (group_expr, _) in aggregate.group_expr().expr().iter() {
|
||||
let column_name = get_column_name(group_expr);
|
||||
if is_column(group_expr) && self.index_fields.contains(column_name) {
|
||||
fields.push(column_name.to_string());
|
||||
} else {
|
||||
// One of the group by fields is not an index field
|
||||
self.simple_topn_multi = None;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
}
|
||||
}
|
||||
|
||||
// Check aggregate function is count(*)
|
||||
let aggr_expr = &aggregate.aggr_expr()[0];
|
||||
if aggr_expr.name() == "count(Int64(1))" {
|
||||
if let Some(simple_topn_multi) = &self.simple_topn_multi {
|
||||
self.simple_topn_multi =
|
||||
Some((fields, simple_topn_multi.1, simple_topn_multi.2));
|
||||
}
|
||||
return Ok(TreeNodeRecursion::Continue);
|
||||
}
|
||||
}
|
||||
self.simple_topn_multi = None;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
} else if is_complex_plan(node) {
|
||||
self.simple_topn_multi = None;
|
||||
return Ok(TreeNodeRecursion::Stop);
|
||||
}
|
||||
Ok(TreeNodeRecursion::Continue)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use arrow::array::{Int64Array, RecordBatch, StringArray};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use datafusion::{catalog::MemTable, prelude::SessionContext};
|
||||
|
||||
use super::*;
|
||||
use crate::service::search::datafusion::udf::match_all_udf;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_is_simple_topn_multi() {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("_timestamp", DataType::Int64, false),
|
||||
Field::new("userid", DataType::Utf8, false),
|
||||
Field::new("searchphrase", DataType::Utf8, false),
|
||||
Field::new("status", DataType::Utf8, false),
|
||||
]));
|
||||
|
||||
let batch = RecordBatch::try_new(
|
||||
schema.clone(),
|
||||
vec![
|
||||
Arc::new(Int64Array::from(vec![1])),
|
||||
Arc::new(StringArray::from(vec!["user1"])),
|
||||
Arc::new(StringArray::from(vec!["hello"])),
|
||||
Arc::new(StringArray::from(vec!["success"])),
|
||||
],
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let ctx = SessionContext::new();
|
||||
let provider = MemTable::try_new(schema, vec![vec![batch.clone()], vec![batch]]).unwrap();
|
||||
ctx.register_table("hits", Arc::new(provider)).unwrap();
|
||||
ctx.register_udf(match_all_udf::MATCH_ALL_UDF.clone());
|
||||
|
||||
let cases = vec![
|
||||
// Valid: two-field GROUP BY with count(*) ORDER BY cnt DESC LIMIT 10
|
||||
(
|
||||
"select userid, searchphrase, count(*) as cnt from hits where match_all('error') group by userid, searchphrase order by cnt desc limit 10",
|
||||
Some(IndexOptimizeMode::SimpleTopNMulti(
|
||||
vec!["userid".to_string(), "searchphrase".to_string()],
|
||||
10,
|
||||
false,
|
||||
)),
|
||||
),
|
||||
// Valid: two-field GROUP BY with count(*) ORDER BY cnt ASC LIMIT 5
|
||||
(
|
||||
"select userid, searchphrase, count(*) as cnt from hits where match_all('error') group by userid, searchphrase order by cnt asc limit 5",
|
||||
Some(IndexOptimizeMode::SimpleTopNMulti(
|
||||
vec!["userid".to_string(), "searchphrase".to_string()],
|
||||
5,
|
||||
true,
|
||||
)),
|
||||
),
|
||||
// Invalid: one field not in index
|
||||
(
|
||||
"select userid, status, count(*) as cnt from hits where match_all('error') group by userid, status order by cnt desc limit 10",
|
||||
None,
|
||||
),
|
||||
// Invalid: single field GROUP BY (should be handled by is_simple_topn)
|
||||
(
|
||||
"select userid, count(*) as cnt from hits where match_all('error') group by userid order by cnt desc limit 10",
|
||||
None,
|
||||
),
|
||||
// Invalid: no limit
|
||||
(
|
||||
"select userid, searchphrase, count(*) as cnt from hits where match_all('error') group by userid, searchphrase order by cnt desc",
|
||||
None,
|
||||
),
|
||||
];
|
||||
|
||||
for (sql, expected) in cases {
|
||||
let plan = ctx.state().create_logical_plan(sql).await.unwrap();
|
||||
let physical_plan = ctx.state().create_physical_plan(&plan).await.unwrap();
|
||||
|
||||
let index_fields = HashSet::from(["userid".to_string(), "searchphrase".to_string()]);
|
||||
assert_eq!(
|
||||
expected,
|
||||
is_simple_topn_multi(physical_plan, index_fields),
|
||||
"Failed for SQL: {}",
|
||||
sql
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -219,6 +219,9 @@ async fn adapt_tantivy_result(
|
||||
IndexOptimizeMode::SimpleTopN(field, limit, _ascend) => {
|
||||
create_top_n_arrow_array(&schema, result.top_n(), &field, limit)?
|
||||
}
|
||||
IndexOptimizeMode::SimpleTopNMulti(ref fields, limit, _ascend) => {
|
||||
create_top_n_multi_arrow_array(&schema, result.top_n_multi(), fields, limit)?
|
||||
}
|
||||
IndexOptimizeMode::SimpleDistinct(_field, _limit, _ascend) => {
|
||||
vec![create_distinct_arrow_array(&schema, result.distinct())?]
|
||||
}
|
||||
@@ -497,6 +500,67 @@ fn create_top_n_arrow_array(
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
fn create_top_n_multi_arrow_array(
|
||||
schema: &SchemaRef,
|
||||
top_n: Vec<(Vec<String>, u64)>,
|
||||
_fields: &[String],
|
||||
_limit: usize,
|
||||
) -> Result<Vec<Vec<Arc<dyn arrow::array::Array>>>, DataFusionError> {
|
||||
// Validate: schema should have 3 fields (field1, field2, count)
|
||||
if schema.fields().len() != 3 {
|
||||
return Err(DataFusionError::Internal(format!(
|
||||
"Expected schema with 3 fields for TopNMulti, got {}",
|
||||
schema.fields().len()
|
||||
)));
|
||||
}
|
||||
|
||||
let field1_field = &schema.fields()[0];
|
||||
let field2_field = &schema.fields()[1];
|
||||
let count_field = &schema.fields()[2];
|
||||
|
||||
let total_rows = top_n.len();
|
||||
if total_rows == 0 {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
// Unzip the composite keys and counts
|
||||
let mut field1_values = Vec::with_capacity(total_rows);
|
||||
let mut field2_values = Vec::with_capacity(total_rows);
|
||||
let mut count_values = Vec::with_capacity(total_rows);
|
||||
|
||||
for (keys, count) in top_n {
|
||||
if keys.len() >= 2 {
|
||||
field1_values.push(keys[0].clone());
|
||||
field2_values.push(keys[1].clone());
|
||||
count_values.push(count);
|
||||
}
|
||||
}
|
||||
|
||||
let total_rows = field1_values.len();
|
||||
if total_rows == 0 {
|
||||
return Ok(vec![]);
|
||||
}
|
||||
|
||||
let mut batches = Vec::new();
|
||||
let batch_size = get_batch_size();
|
||||
|
||||
for chunk_start in (0..total_rows).step_by(batch_size) {
|
||||
let chunk_end = std::cmp::min(chunk_start + batch_size, total_rows);
|
||||
|
||||
let field1_chunk = field1_values[chunk_start..chunk_end].to_vec();
|
||||
let field2_chunk = field2_values[chunk_start..chunk_end].to_vec();
|
||||
let count_chunk = count_values[chunk_start..chunk_end].to_vec();
|
||||
|
||||
let field1_array = create_field_array(field1_field, field1_chunk)?;
|
||||
let field2_array = create_field_array(field2_field, field2_chunk)?;
|
||||
let count_array = create_count_array(count_field, count_chunk)?;
|
||||
|
||||
batches.push(vec![field1_array, field2_array, count_array]);
|
||||
}
|
||||
|
||||
Ok(batches)
|
||||
}
|
||||
|
||||
fn create_distinct_arrow_array(
|
||||
schema: &SchemaRef,
|
||||
distinct_values: HashSet<String>,
|
||||
|
||||
@@ -692,6 +692,7 @@ async fn handle_tantivy_optimize(
|
||||
| Some(IndexOptimizeMode::SimpleHistogram(..))
|
||||
| Some(IndexOptimizeMode::SimpleMultiHistogram(..))
|
||||
| Some(IndexOptimizeMode::SimpleTopN(..))
|
||||
| Some(IndexOptimizeMode::SimpleTopNMulti(..))
|
||||
| Some(IndexOptimizeMode::SimpleDistinct(..))
|
||||
) {
|
||||
return Ok((vec![], file_list));
|
||||
@@ -727,7 +728,9 @@ async fn update_index_updated_at(
|
||||
|
||||
if matches!(
|
||||
idx_optimize_rule,
|
||||
Some(IndexOptimizeMode::SimpleTopN(..)) | Some(IndexOptimizeMode::SimpleMultiHistogram(..))
|
||||
Some(IndexOptimizeMode::SimpleTopN(..))
|
||||
| Some(IndexOptimizeMode::SimpleTopNMulti(..))
|
||||
| Some(IndexOptimizeMode::SimpleMultiHistogram(..))
|
||||
) {
|
||||
let ttv_secondary_index_updated_at =
|
||||
db::metas::tantivy_index::get_ttv_secondary_index_updated_at().await;
|
||||
|
||||
@@ -37,6 +37,7 @@ pub enum CacheEntry {
|
||||
Histogram(Vec<u64>), // simple histogram optimization
|
||||
MultiHistogram(Vec<(i64, String, u64)>), // multi histogram optimization
|
||||
TopN(Vec<(String, u64)>), // simple top n optimization
|
||||
TopNMulti(Vec<(Vec<String>, u64)>), // two-field top n optimization
|
||||
Distinct(HashSet<String>), // simple distinct optimization
|
||||
}
|
||||
|
||||
@@ -59,6 +60,7 @@ impl From<CacheEntry> for TantivyResult {
|
||||
TantivyResult::MultiHistogram(multi_histogram)
|
||||
}
|
||||
CacheEntry::TopN(top_n) => TantivyResult::TopN(top_n),
|
||||
CacheEntry::TopNMulti(top_n) => TantivyResult::TopNMulti(top_n),
|
||||
CacheEntry::Distinct(distinct) => TantivyResult::Distinct(distinct),
|
||||
}
|
||||
}
|
||||
@@ -95,6 +97,17 @@ impl CacheEntry {
|
||||
.sum::<usize>()
|
||||
+ std::mem::size_of::<Vec<(String, u64)>>()
|
||||
}
|
||||
CacheEntry::TopNMulti(top_n) => {
|
||||
top_n
|
||||
.iter()
|
||||
.map(|(keys, _)| {
|
||||
keys.iter().map(|s| s.capacity()).sum::<usize>()
|
||||
+ std::mem::size_of::<Vec<String>>()
|
||||
+ std::mem::size_of::<u64>()
|
||||
})
|
||||
.sum::<usize>()
|
||||
+ std::mem::size_of::<Vec<(Vec<String>, u64)>>()
|
||||
}
|
||||
CacheEntry::Distinct(distinct) => {
|
||||
distinct.iter().map(|s| s.capacity()).sum::<usize>()
|
||||
+ std::mem::size_of::<HashSet<String>>()
|
||||
|
||||
@@ -293,6 +293,10 @@ pub async fn tantivy_search(
|
||||
tantivy_result_builder.add_top_n(top_n);
|
||||
file_list_map.remove(&file_name);
|
||||
}
|
||||
TantivyResult::TopNMulti(top_n) => {
|
||||
tantivy_result_builder.add_top_n_multi(top_n);
|
||||
file_list_map.remove(&file_name);
|
||||
}
|
||||
TantivyResult::Distinct(distinct) => {
|
||||
tantivy_result_builder.add_distinct(distinct);
|
||||
file_list_map.remove(&file_name);
|
||||
@@ -494,6 +498,11 @@ async fn search_tantivy_index(
|
||||
IndexOptimizeMode::SimpleTopN(field, ..) => {
|
||||
need_fast_field.insert(field.clone());
|
||||
}
|
||||
IndexOptimizeMode::SimpleTopNMulti(fields, ..) => {
|
||||
for field in fields {
|
||||
need_fast_field.insert(field.clone());
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
@@ -559,6 +568,9 @@ async fn search_tantivy_index(
|
||||
Some(IndexOptimizeMode::SimpleTopN(field, limit, ascend)) => {
|
||||
TantivyResult::handle_simple_top_n(&searcher, query, &field, limit, ascend)
|
||||
}
|
||||
Some(IndexOptimizeMode::SimpleTopNMulti(fields, limit, ascend)) => {
|
||||
TantivyResult::handle_simple_top_n_multi(&searcher, query, &fields, limit, ascend)
|
||||
}
|
||||
Some(IndexOptimizeMode::SimpleDistinct(field, limit, ascend)) => {
|
||||
if tantivy_schema.get_field(&field).is_err() {
|
||||
log::warn!("[trace_id {trace_id_clone}] search->tantivy: {field} not index in tantivy file: {ttv_file_name}");
|
||||
@@ -579,6 +591,7 @@ async fn search_tantivy_index(
|
||||
TantivyResult::MultiHistogram(multi_histogram)
|
||||
}
|
||||
TantivyResult::TopN(top_n) => TantivyResult::TopN(top_n),
|
||||
TantivyResult::TopNMulti(top_n) => TantivyResult::TopNMulti(top_n),
|
||||
TantivyResult::Distinct(distinct) => TantivyResult::Distinct(distinct),
|
||||
TantivyResult::RowIds(row_ids) => {
|
||||
if row_ids.is_empty() || parquet_file.meta.records == 0 {
|
||||
@@ -676,6 +689,7 @@ fn get_cache_entry(tantivy_result: TantivyResult, percent: f64, parquet_rows: us
|
||||
CacheEntry::MultiHistogram(multi_histogram)
|
||||
}
|
||||
TantivyResult::TopN(top_n) => CacheEntry::TopN(top_n),
|
||||
TantivyResult::TopNMulti(top_n) => CacheEntry::TopNMulti(top_n),
|
||||
TantivyResult::Distinct(distinct) => CacheEntry::Distinct(distinct),
|
||||
TantivyResult::RowIds(_) => {
|
||||
unreachable!("unsupported tantivy search result in search_tantivy_index")
|
||||
|
||||
@@ -45,6 +45,7 @@ pub enum TantivyResult {
|
||||
Histogram(Vec<u64>), // simple histogram optimization
|
||||
MultiHistogram(Vec<(i64, String, u64)>), // multi histogram optimization (with breakdown)
|
||||
TopN(Vec<(String, u64)>), // simple top n optimization
|
||||
TopNMulti(Vec<(Vec<String>, u64)>), // two-field top n optimization
|
||||
Distinct(HashSet<String>), // simple distinct optimization
|
||||
}
|
||||
|
||||
@@ -86,6 +87,17 @@ impl TantivyResult {
|
||||
.sum::<usize>()
|
||||
+ std::mem::size_of::<Vec<(String, u64)>>()
|
||||
}
|
||||
Self::TopNMulti(top_n) => {
|
||||
top_n
|
||||
.iter()
|
||||
.map(|(keys, _)| {
|
||||
keys.iter().map(|s| s.capacity()).sum::<usize>()
|
||||
+ std::mem::size_of::<Vec<String>>()
|
||||
+ std::mem::size_of::<u64>()
|
||||
})
|
||||
.sum::<usize>()
|
||||
+ std::mem::size_of::<Vec<(Vec<String>, u64)>>()
|
||||
}
|
||||
Self::Distinct(distinct) => {
|
||||
distinct.iter().map(|s| s.capacity()).sum::<usize>()
|
||||
+ std::mem::size_of::<HashSet<String>>()
|
||||
@@ -296,6 +308,110 @@ impl TantivyResult {
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle two-field GROUP BY with count(*) using nested TermsAggregation
|
||||
pub fn handle_simple_top_n_multi(
|
||||
searcher: &Searcher,
|
||||
query: Box<dyn Query>,
|
||||
fields: &[String],
|
||||
limit: usize,
|
||||
ascend: bool,
|
||||
) -> anyhow::Result<Self> {
|
||||
if fields.len() != 2 {
|
||||
anyhow::bail!("handle_simple_top_n_multi requires exactly 2 fields");
|
||||
}
|
||||
|
||||
let order = if ascend {
|
||||
Some(CustomOrder {
|
||||
target: OrderTarget::Count,
|
||||
order: Order::Asc,
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Smarter over-fetch: outer needs enough distinct values of field[0],
|
||||
// inner only needs enough per-bucket combinations.
|
||||
// Use sqrt-based scaling to avoid O(N^2) bucket explosion.
|
||||
let effective_limit = limit.max(10);
|
||||
let outer_limit = (effective_limit * 2).max(100) as u32;
|
||||
let inner_limit = (effective_limit * 2).max(100) as u32;
|
||||
|
||||
// Build nested aggregation: outer terms on field[0], inner terms on field[1]
|
||||
let inner_aggregation = Aggregation {
|
||||
agg: AggregationVariants::Terms(TermsAggregation {
|
||||
field: fields[1].to_string(),
|
||||
size: Some(inner_limit),
|
||||
order: order.clone(),
|
||||
missing: None,
|
||||
min_doc_count: Some(1),
|
||||
show_term_doc_count_error: Some(false),
|
||||
segment_size: Some(inner_limit),
|
||||
include: None,
|
||||
exclude: None,
|
||||
}),
|
||||
sub_aggregation: Default::default(),
|
||||
};
|
||||
|
||||
let outer_aggregation = Aggregation {
|
||||
agg: AggregationVariants::Terms(TermsAggregation {
|
||||
field: fields[0].to_string(),
|
||||
size: Some(outer_limit),
|
||||
order,
|
||||
missing: None,
|
||||
min_doc_count: Some(1),
|
||||
show_term_doc_count_error: Some(false),
|
||||
include: None,
|
||||
exclude: None,
|
||||
segment_size: Some(outer_limit),
|
||||
}),
|
||||
sub_aggregation: Aggregations::from_iter(vec![("inner".to_string(), inner_aggregation)]),
|
||||
};
|
||||
|
||||
let aggregations = Aggregations::from_iter(vec![("outer".to_string(), outer_aggregation)]);
|
||||
let collector = AggregationCollector::from_aggs(aggregations, Default::default());
|
||||
|
||||
let mut res = searcher.search(&query, &collector)?;
|
||||
|
||||
// Pre-allocate with a reasonable estimate
|
||||
let estimated_results = (outer_limit as usize) * 4;
|
||||
let mut results: Vec<(Vec<String>, u64)> = Vec::with_capacity(estimated_results);
|
||||
|
||||
if let AggregationResult::BucketResult(BucketResult::Terms {
|
||||
buckets: outer_buckets,
|
||||
..
|
||||
}) = res.0.remove("outer").unwrap()
|
||||
{
|
||||
for mut outer_bucket in outer_buckets {
|
||||
let outer_key = match &outer_bucket.key {
|
||||
Key::Str(s) => s.clone(),
|
||||
Key::F64(f) => f.to_string(),
|
||||
Key::I64(i) => i.to_string(),
|
||||
Key::U64(u) => u.to_string(),
|
||||
};
|
||||
|
||||
if let Some(AggregationResult::BucketResult(BucketResult::Terms {
|
||||
buckets: inner_buckets,
|
||||
..
|
||||
})) = outer_bucket.sub_aggregation.0.remove("inner")
|
||||
{
|
||||
for inner_bucket in inner_buckets {
|
||||
let inner_key = match inner_bucket.key {
|
||||
Key::Str(s) => s,
|
||||
Key::F64(f) => f.to_string(),
|
||||
Key::I64(i) => i.to_string(),
|
||||
Key::U64(u) => u.to_string(),
|
||||
};
|
||||
results.push((vec![outer_key.clone(), inner_key], inner_bucket.doc_count));
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
anyhow::bail!("Failed to get top n multi results from tantivy");
|
||||
}
|
||||
|
||||
Ok(Self::TopNMulti(results))
|
||||
}
|
||||
|
||||
pub fn handle_simple_distinct(
|
||||
searcher: &Searcher,
|
||||
index_condition: &IndexCondition,
|
||||
@@ -346,6 +462,7 @@ pub enum TantivyMultiResultBuilder {
|
||||
Histogram(Vec<Vec<u64>>),
|
||||
MultiHistogram(Vec<Vec<(i64, String, u64)>>),
|
||||
TopN(Vec<(String, u64)>),
|
||||
TopNMulti(Vec<(Vec<String>, u64)>),
|
||||
Distinct(HashSet<String>),
|
||||
}
|
||||
|
||||
@@ -355,6 +472,7 @@ impl TantivyMultiResultBuilder {
|
||||
Some(IndexOptimizeMode::SimpleHistogram(..)) => Self::Histogram(vec![]),
|
||||
Some(IndexOptimizeMode::SimpleMultiHistogram(..)) => Self::MultiHistogram(vec![]),
|
||||
Some(IndexOptimizeMode::SimpleTopN(..)) => Self::TopN(vec![]),
|
||||
Some(IndexOptimizeMode::SimpleTopNMulti(..)) => Self::TopNMulti(vec![]),
|
||||
Some(IndexOptimizeMode::SimpleDistinct(..)) => Self::Distinct(HashSet::new()),
|
||||
Some(IndexOptimizeMode::SimpleSelect(..))
|
||||
| Some(IndexOptimizeMode::SimpleCount)
|
||||
@@ -398,6 +516,13 @@ impl TantivyMultiResultBuilder {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_top_n_multi(&mut self, top_n: Vec<(Vec<String>, u64)>) {
|
||||
match self {
|
||||
Self::TopNMulti(a) => a.extend(top_n),
|
||||
_ => unreachable!("unsupported tantivy multi result"),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn add_distinct(&mut self, distinct: HashSet<String>) {
|
||||
match self {
|
||||
Self::Distinct(a) => a.extend(distinct),
|
||||
@@ -436,6 +561,7 @@ impl TantivyMultiResultBuilder {
|
||||
TantivyMultiResult::MultiHistogram(merged)
|
||||
}
|
||||
Self::TopN(a) => TantivyMultiResult::TopN(a),
|
||||
Self::TopNMulti(a) => TantivyMultiResult::TopNMulti(a),
|
||||
Self::Distinct(a) => TantivyMultiResult::Distinct(a),
|
||||
}
|
||||
}
|
||||
@@ -446,6 +572,7 @@ pub enum TantivyMultiResult {
|
||||
Histogram(Vec<u64>),
|
||||
MultiHistogram(Vec<(i64, String, u64)>),
|
||||
TopN(Vec<(String, u64)>),
|
||||
TopNMulti(Vec<(Vec<String>, u64)>),
|
||||
Distinct(HashSet<String>),
|
||||
}
|
||||
|
||||
@@ -460,6 +587,7 @@ impl Display for TantivyMultiResult {
|
||||
write!(f, "multi_histogram hits: {}", multi_histogram.len())
|
||||
}
|
||||
Self::TopN(top_n) => write!(f, "top_n hits: {}", top_n.len()),
|
||||
Self::TopNMulti(top_n) => write!(f, "top_n_multi hits: {}", top_n.len()),
|
||||
Self::Distinct(distinct) => write!(f, "distinct hits: {}", distinct.len()),
|
||||
}
|
||||
}
|
||||
@@ -494,6 +622,13 @@ impl TantivyMultiResult {
|
||||
}
|
||||
}
|
||||
|
||||
pub fn top_n_multi(self) -> Vec<(Vec<String>, u64)> {
|
||||
match self {
|
||||
Self::TopNMulti(a) => a,
|
||||
_ => vec![],
|
||||
}
|
||||
}
|
||||
|
||||
pub fn distinct(self) -> HashSet<String> {
|
||||
match self {
|
||||
Self::Distinct(a) => a,
|
||||
|
||||
Reference in New Issue
Block a user