Skip to content

Fix Polars JSON and NDJSON reader type hints for loaders - #1762

Open
digit50 wants to merge 1 commit into
apache:mainfrom
digit50:fix/polars-json-reader-type-hints
Open

digit50 wants to merge 1 commit into
apache:mainfrom
digit50:fix/polars-json-reader-type-hints

Conversation

@digit50

@digit50 digit50 commented Oct 15, 2026

Copy link
Copy Markdown

closes #1757

PolarsJSONReader and PolarsNDJSONReader annotated schema and schema_overrides with polars._typing.SchemaDefinition. That alias is built from forward references that polars only imports under TYPE_CHECKING, so typing.get_type_hints on either reader raised NameError: name 'DataTypeClass' is not defined. Since get_required_arguments and get_optional_arguments call get_type_hints, @load_from.json, @load_from.ndjson, from_.json and from_.ndjson all failed before reading any file.

Changes

  • In hamilton/plugins/polars_post_1_0_0_extensions.py, schema and schema_overrides on both readers are now annotated Mapping[str, Any] | Sequence[str | tuple[str, Any]] | None. This resolves at runtime and still covers what polars accepts: a mapping of column name to dtype, or a sequence of names or (name, dtype) pairs. I followed the existing PolarsSpreadsheetReader.schema_overrides annotation. The values are still passed to polars unchanged.
  • I removed the SchemaDefinition import because nothing else uses it.

How I tested this

I added the following to tests/plugins/test_polars_extensions.py:

  • test_getting_type_hints_json_readers, parametrized over both readers. It checks that typing.get_type_hints succeeds and that get_required_arguments and get_optional_arguments return the expected fields.
  • test_polars_json_readers_materialize, for json and ndjson. It writes a file, materializes it through from_.json / from_.ndjson on a driver and checks that a downstream node sees both rows.
  • test_polars_ndjson_load_from. It uses @load_from.ndjson with schema_overrides=value({"a": pl.Float64}) and checks that the override reaches polars.

All five test cases fail on main with the NameError from the issue and pass with this change. With polars 1.44.2, pytest tests/function_modifiers tests/plugins/test_polars*.py gives 453 passed and 2 skipped. With polars 2.0.0, the new tests also pass. The csv, parquet and feather reader tests fail on 2.0.0 both with and without this change, because of arguments that polars 2.0 removed (rechunk, row_count_offset).

Notes

polars_pre_1_0_0_extension.py has the same annotation on its JSON reader. I left it alone because that module targets polars versions before 1.0 and this issue is about the post-1.0 readers.

Checklist

  • PR has an informative and human-readable title (this will be pulled into the release notes)
  • Changes are limited to a single goal (no scope creep)
  • Code passed the pre-commit check & code is left cleaner/nicer than when first encountered.
  • Any change in functionality is tested
  • New functions are documented (with a description, list of inputs, and expected output)
  • Placeholder code is flagged / future TODOs are captured in comments
  • Project documentation has been updated if adding/changing functionality.

PolarsJSONReader and PolarsNDJSONReader annotated schema and schema_overrides with polars' SchemaDefinition, whose forward references cannot be resolved by typing.get_type_hints in this module. That broke @load_from.json/ndjson and from_.json/ndjson with a NameError. The fields now use a runtime-resolvable Mapping or Sequence annotation, and the values are still passed to polars unchanged.

Closes apache#1757

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

@load_from.ndjson and from_.ndjson fail with NameError on 'DataTypeClass'

1 participant