LabDataStorageEvaluation/common/queries/json_queries.py

"""Q1-Q7 over the JSON-LD FULL variant: ijson streaming parser (spec 08).

Each coupon file is streamed event-by-event; scalar coupon fields appear
before the bulk arrays in the serialization, so filtering queries (Q2) can
abandon a non-matching file before parsing its megabytes of points.
"""

from __future__ import annotations

from pathlib import Path

import ijson

from common.queries import Q1_TRACK_CODE
from common.track_summary import summarize_track


class CouponDoc:
	__slots__ = ("coupon_code", "batch_code", "run_code", "au", "hardness", "tracks")

	def __init__(self) -> None:
		self.coupon_code = ""
		self.batch_code = ""
		self.run_code = ""
		self.au = 0.0
		self.hardness: list[float] = []
		self.tracks: list[dict] = []


def _scan(path: Path, au_range: tuple[float, float] | None = None) -> CouponDoc | None:
	"""Stream one coupon file; with au_range set, bail out early on mismatch."""
	doc = CouponDoc()
	track: dict | None = None
	with open(path, "rb") as fh:
		for prefix, event, value in ijson.parse(fh, use_float=True):
			if prefix == "coupon_code":
				doc.coupon_code = value
			elif prefix == "batch_code":
				doc.batch_code = value
			elif prefix == "run_code":
				doc.run_code = value
				if doc.run_code == "RESERVE":
					return None
			elif prefix == "au_wtpct_mean":
				doc.au = value
				if au_range and not au_range[0] <= value <= au_range[1]:
					return None
			elif prefix == "nanoindentation.indents.item.hardness_gpa":
				doc.hardness.append(value)
			elif prefix == "tracks.item" and event == "start_map":
				track = {"cycles": [], "cofs": []}
			elif track is not None:
				if prefix == "tracks.item.track_code":
					track["track_code"] = value
				elif prefix == "tracks.item.environment":
					track["environment"] = value
				elif prefix == "tracks.item.load_mn":
					track["load_mn"] = value
				elif prefix == "tracks.item.speed_mm_s":
					track["speed_mm_s"] = value
				elif prefix == "tracks.item.cycles.item.cycle":
					track["cycles"].append(value)
				elif prefix == "tracks.item.cycles.item.cof":
					track["cofs"].append(value)
				elif prefix == "tracks.item.wear.wear_volume_um3":
					track["wear"] = value
				elif prefix == "tracks.item" and event == "end_map":
					doc.tracks.append(track)
					track = None
	return doc


def _files(json_root: Path) -> list[Path]:
	return sorted((json_root / "full").glob("*.jsonld"))


def _mean(values: list[float]) -> float:
	total = 0.0
	for v in values:
		total += v
	return total / len(values)


def _cof_ss(doc: CouponDoc) -> float:
	return _mean([summarize_track(t["cofs"])[0] for t in doc.tracks])


def q1(json_root: Path) -> list[tuple]:
	coupon_code = Q1_TRACK_CODE.rsplit("-", 1)[0]
	doc = _scan(json_root / "full" / f"{coupon_code}.jsonld")
	for track in doc.tracks:
		if track["track_code"] == Q1_TRACK_CODE:
			return [(int(c), cof) for c, cof in zip(track["cycles"], track["cofs"])]
	raise LookupError(f"track {Q1_TRACK_CODE} not found")


def q2(json_root: Path) -> list[tuple]:
	out = []
	for path in _files(json_root):
		doc = _scan(path, au_range=(9.5, 10.5))
		if doc is None:
			continue
		out.append((doc.coupon_code, doc.tracks[0]["environment"], _cof_ss(doc)))
	return out


def q3(json_root: Path) -> list[tuple]:
	out = []
	for path in _files(json_root):
		doc = _scan(path)
		if doc is None:
			continue
		out.append((doc.coupon_code, doc.batch_code, _mean(doc.hardness), _cof_ss(doc)))
	return out


def q4(json_root: Path) -> list[tuple]:
	out = []
	for path in _files(json_root):
		doc = _scan(path)
		if doc is None:
			continue
		for track in doc.tracks:
			if track["environment"] != "dry_n2" or track["load_mn"] != 100.0:
				continue
			cof_ss = summarize_track(track["cofs"])[0]
			if cof_ss > 0.20:
				out.append((track["track_code"], cof_ss))
	return out


def q5(json_root: Path) -> list[tuple]:
	sums: dict[str, list[float]] = {}
	for path in _files(json_root):
		doc = _scan(path)
		if doc is None:
			continue
		acc = sums.setdefault(doc.batch_code, [0.0, 0.0])
		for track in doc.tracks:
			acc[0] += summarize_track(track["cofs"])[2]
			acc[1] += 1
	return [(batch, acc[0] / acc[1]) for batch, acc in sorted(sums.items())]


def q6(json_root: Path) -> list[tuple]:
	out = []
	for path in _files(json_root):
		doc = _scan(path)
		if doc is None:
			continue
		load = doc.tracks[0]["load_mn"]
		out.append((doc.coupon_code, load, _mean([t["wear"] for t in doc.tracks])))
	return out


def q7(json_root: Path) -> list[tuple]:
	groups: dict[tuple[str, int], list[float]] = {}
	for path in _files(json_root):
		doc = _scan(path)
		if doc is None:
			continue
		for track in doc.tracks:
			bucket = int(round(track["speed_mm_s"] * track["load_mn"]))
			run_in = summarize_track(track["cofs"])[2]
			acc = groups.setdefault((track["environment"], bucket), [0.0, 0.0])
			for cof in track["cofs"][run_in:]:
				acc[0] += cof
				acc[1] += 1
	return [(env, bucket, acc[0] / acc[1]) for (env, bucket), acc in sorted(groups.items())]