Skip to content

Commit 61dcc1d

Browse files
perf: Optimize oneOf validation and add schema/validator caching (#1197)
## Summary Follow-up to #1196 targeting the remaining validation bottlenecks, identified via cProfile on a 135MB Biacore T200 `.bme` file (917 measurements, 35M data points). - **Fast `oneOf` for typed arrays** — `tDimensionArray` and `tMeasureArray` are `oneOf` over 3 typed array refs (number/boolean/string). The default `oneOf` validator validates against the first matching branch, then re-validates against ALL remaining branches to confirm uniqueness. We short-circuit by inspecting the first element's type to select the correct branch directly. **(~6s saved)** - **Fast `oneOf` for array-vs-object** — `oneOf[tDimensionArray, tFunction]` patterns: if the instance is a list, it must be the array branch (tFunction is an object). Skip the redundant validation. **(917 calls saved)** - **Schema + validator caching** — Cache schema JSON, `RefResolver`, and `FastDraft202012Validator` instances per schema path to avoid redundant disk I/O and object construction on repeated validations. **(~1.7s saved on first call)** - **Eliminate redundant `fields()` call** — `unstructure()` called `fields()` twice per dataclass (once for iteration, once to build a set for custom_information_document check). Reuse the tuple. ### Benchmark (135MB .bme file) | Phase | Before #1196 | After #1196 | After this PR | |---|---|---|---| | unstructure | 13s | 0.2s | 0.2s | | validate | 147s | 12.8s | **5.0s** | | **Total** | **~176s** | **~38s** | **~25s** | ## Test plan - [x] Full test suite passes (1,119 tests) - [x] Lint + mypy clean - [x] Verified oneOf fast validator correctly handles: number/string/boolean arrays, nullable arrays, mixed-type rejection, non-array oneOf fallback, array-vs-object disambiguation - [x] Benchmarked end-to-end against 135MB `.bme` file 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent fe88e2b commit 61dcc1d

2 files changed

Lines changed: 129 additions & 16 deletions

File tree

src/allotropy/allotrope/converter.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,8 @@ def unstructure(obj: Any) -> Any:
183183

184184
if is_dataclass(obj) and not isinstance(obj, type):
185185
result: dict[str, Any] = {}
186-
for f in fields(obj):
186+
dc_fields = fields(obj)
187+
for f in dc_fields:
187188
value = getattr(obj, f.name)
188189
if value is None:
189190
# Keep None for required fields (no default) to preserve
@@ -194,10 +195,9 @@ def unstructure(obj: Any) -> Any:
194195
json_key = f.metadata.get("json_name", default_json_name(f.name))
195196
result[json_key] = unstructure(value)
196197
# Handle dynamically-attached custom_information_document (not in fields())
197-
field_names = {f.name for f in fields(obj)}
198198
if (
199199
hasattr(obj, "custom_information_document")
200-
and "custom_information_document" not in field_names
200+
and not any(f.name == "custom_information_document" for f in dc_fields)
201201
and not isinstance(obj.custom_information_document, list)
202202
):
203203
result[

src/allotropy/allotrope/schemas.py

Lines changed: 126 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -110,9 +110,97 @@ def _fast_items_validator(
110110
return
111111

112112

113+
# ---------------------------------------------------------------------------
114+
# Optimized oneOf validator for homogeneous arrays
115+
# ---------------------------------------------------------------------------
116+
# `tDimensionArray` and `tMeasureArray` are defined as oneOf over typed array
117+
# variants (e.g., tNumberArray | tBooleanArray | tStringArray). The default
118+
# oneOf validator validates against the first matching branch, then
119+
# re-validates against ALL remaining branches to confirm uniqueness.
120+
# For large arrays this is extremely wasteful — we can determine the correct
121+
# branch by inspecting the first element's type.
122+
123+
_ARRAY_ONEOF_TYPE_MAP: dict[type, int] = {
124+
int: 0, # number branch
125+
float: 0, # number branch
126+
str: 2, # string branch
127+
bool: 1, # boolean branch
128+
}
129+
130+
_original_one_of = jsonschema.validators.Draft202012Validator.VALIDATORS["oneOf"]
131+
132+
133+
def _is_typed_array_oneof(one_of: list[Any]) -> bool:
134+
"""Check if this oneOf matches the tDimensionArray/tMeasureArray pattern."""
135+
if len(one_of) != 3:
136+
return False
137+
for branch in one_of:
138+
if not isinstance(branch, dict) or "$ref" not in branch:
139+
return False
140+
ref = branch["$ref"]
141+
# All cube.schema oneOf refs look like #/$defs/tNumberArray, etc.
142+
if not isinstance(ref, str):
143+
return False
144+
name = ref.rsplit("/", 1)[-1] if "/" in ref else ""
145+
if name not in (
146+
"tNumberArray",
147+
"tBooleanArray",
148+
"tStringArray",
149+
"tNumberOrNullArray",
150+
"tBooleanOrNullArray",
151+
"tStringOrNullArray",
152+
):
153+
return False
154+
return True
155+
156+
157+
def _fast_one_of_validator(
158+
validator: Any, one_of: Any, instance: Any, schema: Any
159+
) -> Any:
160+
"""oneOf validator that short-circuits for typed array schemas."""
161+
if isinstance(instance, list) and isinstance(one_of, list):
162+
# Fast path for tDimensionArray/tMeasureArray oneOf (3 typed array refs).
163+
if instance and _is_typed_array_oneof(one_of):
164+
first = instance[0]
165+
element = first
166+
if element is None:
167+
for item in instance:
168+
if item is not None:
169+
element = item
170+
break
171+
if element is not None:
172+
branch_idx = _ARRAY_ONEOF_TYPE_MAP.get(type(element))
173+
if branch_idx is not None:
174+
yield from validator.descend(
175+
instance, one_of[branch_idx], schema_path=branch_idx
176+
)
177+
return
178+
# Fast path for oneOf[tDimensionArray, tFunction]: if instance is a list,
179+
# it must be the array branch (tFunction is an object).
180+
elif len(one_of) == 2:
181+
for idx, branch in enumerate(one_of):
182+
if isinstance(branch, dict) and "$ref" in branch:
183+
ref_name = branch["$ref"].rsplit("/", 1)[-1]
184+
if ref_name in (
185+
"tDimensionArray",
186+
"tMeasureArray",
187+
"tNumberArray",
188+
"tNumberOrNullArray",
189+
):
190+
yield from validator.descend(
191+
instance, one_of[idx], schema_path=idx
192+
)
193+
return
194+
195+
yield from _original_one_of(validator, one_of, instance, schema)
196+
197+
113198
FastDraft202012Validator = jsonschema.validators.extend( # type: ignore[no-untyped-call]
114199
jsonschema.validators.Draft202012Validator,
115-
validators={"items": _fast_items_validator},
200+
validators={
201+
"items": _fast_items_validator,
202+
"oneOf": _fast_one_of_validator,
203+
},
116204
)
117205

118206

@@ -157,33 +245,58 @@ def _get_schema_store() -> dict[str, dict[str, Any]]:
157245
return store
158246

159247

248+
_schema_cache: dict[str, dict[str, Any]] = {}
249+
250+
160251
def _get_schema_by_path(schema_path: Path) -> dict[str, Any]:
161-
"""Load a schema from the schemas/ directory."""
252+
"""Load a schema from the schemas/ directory, with caching."""
253+
key = str(schema_path)
254+
cached = _schema_cache.get(key)
255+
if cached is not None:
256+
return cached
162257
full_path = SCHEMA_DIR_PATH / schema_path
163258
with open(full_path, encoding=DEFAULT_ENCODING) as f:
164-
return json.load(f) # type: ignore[no-any-return]
259+
schema: dict[str, Any] = json.load(f)
260+
_schema_cache[key] = schema
261+
return schema
262+
263+
264+
_validator_cache: dict[str, Any] = {}
265+
266+
267+
def _get_validator(schema_path: Path) -> Any:
268+
"""Get a cached validator for a schema path."""
269+
key = str(schema_path)
270+
cached = _validator_cache.get(key)
271+
if cached is not None:
272+
return cached
273+
schema = _get_schema_by_path(schema_path)
274+
store = _get_schema_store()
275+
resolver = jsonschema.RefResolver(
276+
base_uri=schema.get("$id", ""),
277+
referrer=schema,
278+
store=store,
279+
)
280+
validator = FastDraft202012Validator(
281+
schema, resolver=resolver, format_checker=FORMAT_CHECKER
282+
)
283+
_validator_cache[key] = validator
284+
return validator
165285

166286

167287
def validate_asm_schema(asm_dict: dict[str, Any]) -> None:
168288
"""Validate an ASM dict against schemas with cross-schema $ref resolution."""
169289
try:
170290
schema_path = get_schema_path_from_asm(asm_dict)
171-
schema = _get_schema_by_path(schema_path)
172291
except Exception as e:
173292
msg = f"Failed to retrieve schema for model: {e}"
174293
raise AllotropeSerializationError(msg) from e
175294

176295
try:
177-
store = _get_schema_store()
178-
resolver = jsonschema.RefResolver(
179-
base_uri=schema.get("$id", ""),
180-
referrer=schema,
181-
store=store,
182-
)
183-
validator = FastDraft202012Validator(
184-
schema, resolver=resolver, format_checker=FORMAT_CHECKER
185-
)
296+
validator = _get_validator(schema_path)
186297
validator.validate(asm_dict)
298+
except AllotropeValidationError:
299+
raise
187300
except Exception as e:
188301
msg = f"Failed to validate allotrope model against schema: {e}"
189302
raise AllotropeValidationError(msg) from e

0 commit comments

Comments
 (0)