fix: unify Parquet schemas when reading

[shan-nvidia]

Summary

Reconcile compatible Parquet checkpoint schema drift when Data Designer reads a dataset directory. This fixes completed runs failing during profiling or dataset loading when separate batches infer different physical types for the same nested numeric field, such as 9 becoming int64 in one file and 9.5 becoming double in another.

Related Issue

Fixes #857

Changes

• Build a permissive unified Arrow schema from directory Parquet files in read_parquet_dataset().
• Pass that explicit schema to Pandas/PyArrow so compatible nested numeric drift is promoted consistently.
• Preserve the existing per-file concatenation fallback when Arrow cannot unify historical schemas.
• Add a focused regression using separate nested 9 and 9.5 checkpoints.
• Keep JSON Schema validation unchanged, including the existing Decimal-specific normalization.

This follows the same physical-schema reconciliation approach already used by Parquet export and keeps JSON Schema validation focused on whether values are valid rather than choosing a canonical Python or Parquet representation.

Checkpoint Scope

The change fixes supported Data Designer paths that use read_parquet_dataset(), including profiling, ArtifactStorage.load_dataset(), result loading, and processor-output loading. It does not rewrite completed checkpoint files. A direct pandas.read_parquet(directory) call may therefore still fail on a mixed-schema artifact directory; callers should use the Data Designer result or artifact-storage APIs. The repository already preserves a per-file fallback for schema combinations that Arrow cannot unify.

Testing

• 185 passed in focused config, storage, validator, and response-recipe tests.
• 2901 passed across the full config and engine suites.
• 1089 passed, 1 skipped in the interface suite before sandbox-only localhost bind failures.
• 31 passed when the complete OpenTelemetry test file was rerun with localhost socket access.
• Ruff formatting and lint checks passed for the changed files.
• The regression verifies the input checkpoint schemas differ and the loaded nested scores are [9.0, 9.5].

Checklist

• Follows commit message conventions
• Commits are signed off (DCO)
• Unit regression coverage added
• JSON Schema validation behavior remains unchanged from main

[greptile-apps]

Greptile Summary

The PR updates Parquet dataset loading to unify compatible schemas across directory files before reading them.

• Collects and permissively unifies schemas from sorted Parquet files.
• Falls back to reading and concatenating files individually when schema unification fails.
• Adds regression coverage for nested integer-to-float promotion across checkpoints.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
packages/data-designer-config/src/data_designer/config/utils/io_helpers.py	Adds permissive schema unification for directory-based Parquet reads, with the existing per-file concatenation behavior retained as fallback.
packages/data-designer-config/tests/config/utils/test_io_helpers.py	Adds coverage showing nested integer and fractional values from separate Parquet files are read through a unified floating-point schema.

_{Reviews (3): Last reviewed commit: "fix: unify parquet schemas when reading ..." | Re-trigger Greptile}

[github-actions]

Code Review — PR #858: fix: normalize structured JSON numbers

Summary

This PR replaces the shallow, $defs-only normalize_decimal_fields walker with
a recursive _normalize_numeric_fields traversal that canonicalizes numeric
values against their JSON Schema types after validation. The goal (fixing #857)
is sound: emitting a declared type: number field sometimes as int and
sometimes as float produces incompatible nested Parquet schemas across
checkpoints, which fails final profile/export. Coercing number → float and
integer → int at the validation boundary is the right layer to fix this.

The new implementation is a clear improvement over the old one: it follows
nested objects, arrays, prefixItems, additionalProperties, local $refs
(with cycle protection via seen_refs), and composition keywords, and it
preserves the existing Pydantic Decimal anyOf quantization by checking it
first. The rename from normalize_decimal_fields → normalize_numeric_fields
is complete — no stale references remain in src/tests. Test coverage is
good, including the two-checkpoint Parquet regression that directly reproduces
#857.

The findings below are mostly robustness/correctness edge cases in schema
shapes that StructuredResponseRecipe accepts as raw user-supplied JSON Schema
(not just Pydantic-generated schemas).

Findings

1. Boolean / non-dict subschemas crash normalization (validators.py:181, :208-220)

_normalize_numeric_fields assumes every schema reaching it is a dict, but
JSON Schema permits boolean subschemas (true/false) and, in draft-4-style
schemas, a list-valued items. These reach the recursion unguarded:

• dict branch (:211): field_schema = properties.get(key, ...) — a property
  value of true/false (e.g. {"properties": {"metadata": true}}) is passed
  straight into the recursion.
• list branch (:219): item_schema = schema.get("items", {}) — items: false
  or items: [ ... ] (tuple form) becomes field_schema.

On the next recursion, _resolve_local_ref(True/False/list) returns it
unchanged (not a dict), then _get_decimal_info_from_anyof(schema) executes
schema.get("anyOf") on a bool/list, raising
AttributeError: 'bool' object has no attribute 'get'. This fires after
successful validation, so it aborts an otherwise-valid run.

Failure scenario: a StructuredResponseRecipe built from a hand-written schema
containing {"properties": {"tags": true}} with data {"tags": {"x": 1}}
crashes in normalize_numeric_fields. (The common Pydantic case
additionalProperties: true is guarded at :211 via the isinstance(..., dict)
check, but property-level and item-level boolean schemas are not.) Note the old
code also mishandled boolean schemas, so this is a pre-existing gap the rewrite
carries forward rather than a fresh regression — but the broadened traversal
increases how often it is reached. Suggest a if not isinstance(schema, dict): return obj
guard at the top of _normalize_numeric_fields (right after the _resolve_local_ref
call would still crash — put it before _get_decimal_info_from_anyof).

2. oneOf/anyOf match returns early, skipping sibling properties/items (validators.py:188-193)

When an anyOf/oneOf alternative validates, the function recurses into that
alternative only and returns immediately. If the same schema node also carries
its own properties, items, or additionalProperties (legal in JSON Schema
2020-12, where composition keywords and structural keywords apply jointly),
those sibling keywords are never processed, so nested numbers under them are
left un-normalized.

Failure scenario: a schema node
{"anyOf": [{"required": ["a"]}, {"required": ["b"]}], "properties": {"score": {"type": "number"}}}
with value {"a": 1, "score": 9} — the first anyOf alternative validates, the
function returns, and score stays an int, reintroducing exactly the
cross-checkpoint dtype mismatch this PR fixes. Pydantic-generated schemas keep
properties inside the alternatives (via $ref), so this is primarily a
raw-schema concern, but StructuredResponseRecipe accepts raw schemas. Consider
continuing into the structural branches after handling composition, rather than
early-returning.

3. Redundant full-subtree re-validation per alternative at every level (validators.py:192) — efficiency

validator.evolve(schema=alternative).is_valid(obj) validates the entire
obj subtree against each alternative, and this runs at every recursion level
that has an anyOf/oneOf. For deeply nested data with unions at multiple
levels, the same subtrees are re-validated repeatedly (roughly
O(depth × alternatives × subtree-size)) purely to pick a normalization branch —
on top of the full validation already performed by validate() before
normalization is called. For large generated objects this adds measurable
per-record overhead on the hot generation path. If the common case is a small
number of shallow unions this is acceptable; worth a comment noting the cost, or
short-circuiting scalar obj (the only values that actually get coerced) before
evolving a validator.

4. allOf handling double-processes the object (validators.py:195-213) — cleanup

After the allOf loop normalizes a dict obj through each subschema, control
falls through to the dict branch (:207), which re-iterates the (already
normalized) object using the outer schema's properties (typically {} for
an allOf wrapper), recursing into every field again with an empty schema. It's
harmless (floats stay floats) but redundant, and it means each field is walked
twice for the common Pydantic {"allOf": [{"$ref": ...}], "description": ...}
shape. Not a bug; flagging as avoidable work / minor complexity.

Test observations

• The added tests are well-targeted. test_normalized_nested_numbers_have_compatible_parquet_schemas
  is the most valuable — it exercises the actual Structured JSON numbers can break Parquet batch reads #857 failure mode (two Parquet
  batches, 9 vs 9.5, read back as one dataset) rather than just asserting the
  Python type.
• Consider adding a case for a nullable numeric field
  (anyOf: [{type: number}, {type: null}]) with both a numeric value and
  None, since that anyOf shape is common in Pydantic Optional[float]
  output and exercises the non-Decimal anyOf path (finding DataDesigner.make_seed_reference_from_file doesn't support paths with multiple parquet partition #2's neighbor).

Structural Impact (graphify, 2.7s)

Risk: HIGH (4 import direction violation(s))

• 3 Python files, 17 AST entities, 2/78 clusters

Import Direction Violations (4)

Legal direction: interface -> engine -> config

• prune_additional_properties() (engine) --calls--> .keys() (interface)
• _validate_one_of_with_discriminator() (engine) --calls--> ValidationError (interface)
• _resolve_local_ref() (engine) --calls--> .items() (interface)
• _normalize_numeric_fields() (engine) --calls--> .items() (interface)

High-Connectivity Changes

• validate() (10 deps) in packages/data-designer-engine/src/data_designer/engine/processing/gsonschema/validators.py
• _normalize_numeric_fields() (10 deps) in packages/data-designer-engine/src/data_designer/engine/processing/gsonschema/validators.py
• validators.py (8 deps) in packages/data-designer-engine/src/data_designer/engine/processing/gsonschema/validators.py
• _resolve_local_ref() (6 deps) in packages/data-designer-engine/src/data_designer/engine/processing/gsonschema/validators.py

Cross-Package Dependencies

• .run_validate() (interface) --calls--> validate() (engine)
• prune_additional_properties() (engine) --calls--> .keys() (interface)
• prune_additional_properties() (engine) --calls--> info() (config)
• _validate_one_of_with_discriminator() (engine) --calls--> ValidationError (interface)
• _resolve_local_ref() (engine) --calls--> .items() (interface)
• _normalize_numeric_fields() (engine) --calls--> .items() (interface)

Reviewer note on the flagged violations: these appear to be analysis
artifacts from method-name collisions rather than real import-direction
violations. .keys()/.items() are builtin dict methods, and
ValidationError in this file is lazy.jsonschema.ValidationError (the
jsonschema library), not the interface package's error type — the file imports
nothing from data_designer.interface. No new cross-layer import is introduced
by this diff; validators.py remains within the engine layer and depends only
on engine/config/third-party. The "HIGH" rating is driven by the
connectivity of the touched validate()/normalization functions, which does
warrant the extra scrutiny applied above — validate() is a widely-called
boundary, so the edge-case behaviors in findings #1 and #2 have broad blast
radius.

Verdict

Approve with minor changes suggested. The core fix is correct, well-targeted
at the right layer, and the Parquet regression test is convincing. None of the
findings block the #857 fix for the Pydantic-generated-schema path that
motivated it. Findings #1 (non-dict subschema crash) and #2 (composition
sibling keywords skipped) are worth addressing before this is relied on for
arbitrary user-supplied StructuredResponseRecipe schemas, since both can
either crash a completed run or silently reintroduce the dtype mismatch the PR
sets out to eliminate. #3 and #4 are non-blocking efficiency/cleanup notes.

[shan-nvidia]

Addressed the agentic review in ffc7f94:

1. Boolean/non-dict subschemas now return the already validated value unchanged; Decimal anyOf detection also skips boolean alternatives.
2. Matching oneOf/anyOf branches are normalized without returning early, so sibling properties, items, and other structural keywords still apply.
3. Branch-level validation remains because the initial JSON Schema validation does not expose which union alternative matched. The code now documents that reason and still stops at the first matching alternative.
4. Object and array traversal now skips fields with no applicable subschema, eliminating the redundant walk through unconstrained {} wrappers.

Added coverage for boolean item/property/union schemas, both oneOf and anyOf siblings, and nullable numbers. Local verification: 2267 engine tests passed; full engine formatting and lint passed.

[nabinchha]

Thanks for putting this PR together and for adding a focused regression for #857, @shan-nvidia. A cleaner solution might be to reconcile schemas at the Parquet boundary instead of adding a recursive JSON Schema normalizer.

JSON Schema validation establishes whether a value is valid, but it does not define a canonical Python or Parquet representation—9 and 9.0 are both valid number values. Walking the schema to coerce values effectively reimplements Draft 2020-12 behavior for $ref, anyOf, patternProperties, and related applicators, creating unnecessary correctness and maintenance risks.

I tested the storage-layer approach using the exact reproduction from #857:

• The original directory read failed with the expected nested int64/double ArrowInvalid.
• Reading with pa.unify_schemas(..., promote_options="permissive") and passing the resulting schema to Arrow returned [9.0, 9.5].
• Casting and atomically rewriting the mismatched checkpoint files made the reproduction's original pd.read_parquet(directory) call succeed unchanged.

I recommend:

• Reusing the schema-unification approach already implemented by _export_parquet().
• Applying it in read_parquet_dataset() so profiling and dataset loading use an explicit unified schema.
• If raw checkpoint compatibility is required, adding a finalization step that atomically rewrites only batches whose schemas differ from the unified schema.
• Retaining the existing Decimal-specific normalization while avoiding generic JSON Schema numeric coercion.

This keeps validation focused on validation, puts physical type reconciliation in the storage layer where it belongs, and handles nested numeric drift regardless of whether the schema uses unions, references, or pattern properties.

---

This review was generated by an AI assistant.

[andreatnvidia]

Thanks for putting this together and for the thorough regression coverage. This is a useful fix for a pretty frustrating failure mode.

I found two things I think we should fix before merging, plus two JSON Schema edge cases worth handling while we're here. The blockers are the Decimal path crashing on valid strings and numeric unions narrowing floats back to integers, which can recreate the Parquet failure. I reproduced the cases below against ffc7f942.

Happy to take another look once these are addressed.

[shan-nvidia]

Thanks, @nabinchha. I agreed with the ownership boundary and refactored the PR in 137f71b. The recursive JSON Schema numeric walker and all of its validation-layer tests are removed; validator behavior is restored to main, including the existing Decimal-specific handling.

read_parquet_dataset() now builds a permissive unified Arrow schema for directory reads and passes it explicitly to Pandas/PyArrow, while retaining the existing per-file fallback when schemas cannot be unified. The regression writes the exact nested 9/9.5 checkpoint pair and verifies it loads as [9.0, 9.5].

I did not rewrite raw checkpoints because Data Designer's supported profiling/load paths already go through this helper and the repository already treats direct pandas.read_parquet(directory) compatibility separately.

[github-actions]

Fern preview: https://nvidia-preview-pr-858.docs.buildwithfern.com/nemo/datadesigner

Fern previews include the docs-website version archive with PR changes synced into latest. Notebook tutorials are rendered without execution outputs in previews.

[shan-nvidia]

Thanks for putting this together and for the thorough regression coverage. This is a useful fix for a pretty frustrating failure mode.

I found two things I think we should fix before merging, plus two JSON Schema edge cases worth handling while we're here. The blockers are the Decimal path crashing on valid strings and numeric unions narrowing floats back to integers, which can recreate the Parquet failure. I reproduced the cases below against ffc7f942.

Happy to take another look once these are addressed.

Thanks @andreatnvidia for the comments! I did a refactoring suggested by Nabin's cleaner option, and change scope is much smaller now.

[andreatnvidia]

Thanks for the contribution, @shan-nvidia! Moving schema reconciliation into the Parquet reader is a much cleaner and more focused fix, and it addresses all my earlier concerns. Looks good to me.