From a5e50a60a05dfdac9684e754668d9befd3d2c6aa Mon Sep 17 00:00:00 2001 From: Nathan Stender Date: Thu, 30 Apr 2026 10:41:06 -0400 Subject: [PATCH] perf: Optimize oneOf validation and add schema/validator caching MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up to #1196 targeting the remaining validation bottlenecks for large data cube files (profiled on a 135MB Biacore T200 .bme file with 35M float values across 917 measurements). 1. Fast oneOf for typed arrays (tDimensionArray/tMeasureArray): the default oneOf validator checks all branches then re-validates to confirm uniqueness. We short-circuit by inspecting the first element's type to pick the correct branch directly. 2. Fast oneOf for array-vs-object (tDimensionArray | tFunction): if the instance is a list it must be the array branch, skip the object check. 3. Cache schema JSON, RefResolver, and validator instances per schema path to avoid redundant disk I/O and object construction. 4. Eliminate redundant fields() call in unstructure — reuse the tuple from the iteration loop. Validation time: 12.8s → 5.0s. Total pipeline: ~38s → ~25s. Co-Authored-By: Claude Opus 4.6 --- src/allotropy/allotrope/converter.py | 6 +- src/allotropy/allotrope/schemas.py | 139 ++++++++++++++++++++++++--- 2 files changed, 129 insertions(+), 16 deletions(-) diff --git a/src/allotropy/allotrope/converter.py b/src/allotropy/allotrope/converter.py index 3bff3f625..16fa1a214 100644 --- a/src/allotropy/allotrope/converter.py +++ b/src/allotropy/allotrope/converter.py @@ -183,7 +183,8 @@ def unstructure(obj: Any) -> Any: if is_dataclass(obj) and not isinstance(obj, type): result: dict[str, Any] = {} - for f in fields(obj): + dc_fields = fields(obj) + for f in dc_fields: value = getattr(obj, f.name) if value is None: # Keep None for required fields (no default) to preserve @@ -194,10 +195,9 @@ def unstructure(obj: Any) -> Any: json_key = f.metadata.get("json_name", default_json_name(f.name)) result[json_key] = unstructure(value) # Handle dynamically-attached custom_information_document (not in fields()) - field_names = {f.name for f in fields(obj)} if ( hasattr(obj, "custom_information_document") - and "custom_information_document" not in field_names + and not any(f.name == "custom_information_document" for f in dc_fields) and not isinstance(obj.custom_information_document, list) ): result[ diff --git a/src/allotropy/allotrope/schemas.py b/src/allotropy/allotrope/schemas.py index 83e441d3e..fe7d7095f 100644 --- a/src/allotropy/allotrope/schemas.py +++ b/src/allotropy/allotrope/schemas.py @@ -110,9 +110,97 @@ def _fast_items_validator( return +# --------------------------------------------------------------------------- +# Optimized oneOf validator for homogeneous arrays +# --------------------------------------------------------------------------- +# `tDimensionArray` and `tMeasureArray` are defined as oneOf over typed array +# variants (e.g., tNumberArray | tBooleanArray | tStringArray). The default +# oneOf validator validates against the first matching branch, then +# re-validates against ALL remaining branches to confirm uniqueness. +# For large arrays this is extremely wasteful — we can determine the correct +# branch by inspecting the first element's type. + +_ARRAY_ONEOF_TYPE_MAP: dict[type, int] = { + int: 0, # number branch + float: 0, # number branch + str: 2, # string branch + bool: 1, # boolean branch +} + +_original_one_of = jsonschema.validators.Draft202012Validator.VALIDATORS["oneOf"] + + +def _is_typed_array_oneof(one_of: list[Any]) -> bool: + """Check if this oneOf matches the tDimensionArray/tMeasureArray pattern.""" + if len(one_of) != 3: + return False + for branch in one_of: + if not isinstance(branch, dict) or "$ref" not in branch: + return False + ref = branch["$ref"] + # All cube.schema oneOf refs look like #/$defs/tNumberArray, etc. + if not isinstance(ref, str): + return False + name = ref.rsplit("/", 1)[-1] if "/" in ref else "" + if name not in ( + "tNumberArray", + "tBooleanArray", + "tStringArray", + "tNumberOrNullArray", + "tBooleanOrNullArray", + "tStringOrNullArray", + ): + return False + return True + + +def _fast_one_of_validator( + validator: Any, one_of: Any, instance: Any, schema: Any +) -> Any: + """oneOf validator that short-circuits for typed array schemas.""" + if isinstance(instance, list) and isinstance(one_of, list): + # Fast path for tDimensionArray/tMeasureArray oneOf (3 typed array refs). + if instance and _is_typed_array_oneof(one_of): + first = instance[0] + element = first + if element is None: + for item in instance: + if item is not None: + element = item + break + if element is not None: + branch_idx = _ARRAY_ONEOF_TYPE_MAP.get(type(element)) + if branch_idx is not None: + yield from validator.descend( + instance, one_of[branch_idx], schema_path=branch_idx + ) + return + # Fast path for oneOf[tDimensionArray, tFunction]: if instance is a list, + # it must be the array branch (tFunction is an object). + elif len(one_of) == 2: + for idx, branch in enumerate(one_of): + if isinstance(branch, dict) and "$ref" in branch: + ref_name = branch["$ref"].rsplit("/", 1)[-1] + if ref_name in ( + "tDimensionArray", + "tMeasureArray", + "tNumberArray", + "tNumberOrNullArray", + ): + yield from validator.descend( + instance, one_of[idx], schema_path=idx + ) + return + + yield from _original_one_of(validator, one_of, instance, schema) + + FastDraft202012Validator = jsonschema.validators.extend( # type: ignore[no-untyped-call] jsonschema.validators.Draft202012Validator, - validators={"items": _fast_items_validator}, + validators={ + "items": _fast_items_validator, + "oneOf": _fast_one_of_validator, + }, ) @@ -157,33 +245,58 @@ def _get_schema_store() -> dict[str, dict[str, Any]]: return store +_schema_cache: dict[str, dict[str, Any]] = {} + + def _get_schema_by_path(schema_path: Path) -> dict[str, Any]: - """Load a schema from the schemas/ directory.""" + """Load a schema from the schemas/ directory, with caching.""" + key = str(schema_path) + cached = _schema_cache.get(key) + if cached is not None: + return cached full_path = SCHEMA_DIR_PATH / schema_path with open(full_path, encoding=DEFAULT_ENCODING) as f: - return json.load(f) # type: ignore[no-any-return] + schema: dict[str, Any] = json.load(f) + _schema_cache[key] = schema + return schema + + +_validator_cache: dict[str, Any] = {} + + +def _get_validator(schema_path: Path) -> Any: + """Get a cached validator for a schema path.""" + key = str(schema_path) + cached = _validator_cache.get(key) + if cached is not None: + return cached + schema = _get_schema_by_path(schema_path) + store = _get_schema_store() + resolver = jsonschema.RefResolver( + base_uri=schema.get("$id", ""), + referrer=schema, + store=store, + ) + validator = FastDraft202012Validator( + schema, resolver=resolver, format_checker=FORMAT_CHECKER + ) + _validator_cache[key] = validator + return validator def validate_asm_schema(asm_dict: dict[str, Any]) -> None: """Validate an ASM dict against schemas with cross-schema $ref resolution.""" try: schema_path = get_schema_path_from_asm(asm_dict) - schema = _get_schema_by_path(schema_path) except Exception as e: msg = f"Failed to retrieve schema for model: {e}" raise AllotropeSerializationError(msg) from e try: - store = _get_schema_store() - resolver = jsonschema.RefResolver( - base_uri=schema.get("$id", ""), - referrer=schema, - store=store, - ) - validator = FastDraft202012Validator( - schema, resolver=resolver, format_checker=FORMAT_CHECKER - ) + validator = _get_validator(schema_path) validator.validate(asm_dict) + except AllotropeValidationError: + raise except Exception as e: msg = f"Failed to validate allotrope model against schema: {e}" raise AllotropeValidationError(msg) from e