feat: Adds support for DATE, DATETIME and JSON BigQuery types in the Python SDK using Storage API

[enzomaruffa]

Summary

This PR adds a type_overrides parameter to BigQuery I/O operations, allowing users to specify custom BigQuery-to-Python type mappings. This addresses #25946 by enabling support for types like DATE, DATETIME, and JSON without hardcoding them in the SDK - allowing flexible usage of str to solve a lot of use cases.

Background

We were migrating a production Dataflow pipeline from Java to Python, using the Storage Write API with use_beam_io_types=True (BEAM_ROW format). Our pipeline writes to BigQuery tables containing DATE, DATETIME, and JSON columns.

The pipeline failed with:

ValueError: Encountered an unsupported type: 'DATETIME'

Solution

First we had added str types - older commits still here. Now, rather than hardcoding new type mappings (which may conflict with future official implementations), this PR adds a type_overrides parameter that lets users define their own mappings:

# Storage Write API
WriteToBigQuery(
    table='project:dataset.table',
    method=WriteToBigQuery.Method.STORAGE_WRITE_API,
    use_beam_io_types=True,
    type_overrides={
        'DATE': str,
        'DATETIME': str,
        'JSON': str,
    }
)

Files modified:

• bigquery.py: Added type_overrides to WriteToBigQuery and StorageWriteToBigQuery
• bigquery_tools.py: Added type_overrides to get_beam_typehints_from_tableschema()
• bigquery_schema_tools.py: Added type_overrides to bq_field_to_type(), generate_user_type_from_bq_schema(), convert_to_usertype()

Implementation Note

This PR initially added hardcoded mappings for DATE, DATETIME, and JSON types (similar to how GEOGRAPHY was added in #36121). Based on reviewer feedback, we refactored to the type_overrides approach instead. The commit history reflects this evolution - earlier commits add the hardcoded types, and later commits introduce type_overrides and remove the hardcoded mappings.

How It Works

The implementation merges user-provided overrides with the default type mappings:

def get_beam_typehints_from_tableschema(schema, type_overrides=None):
    effective_types = {**BIGQUERY_TYPE_TO_PYTHON_TYPE, **(type_overrides or {})}
    # ... uses effective_types for schema conversion

Testing

Added comprehensive tests for type_overrides:

• Verifies unsupported types raise ValueError without overrides
• Verifies overrides enable custom type mappings
• Tests all field modes (REQUIRED, NULLABLE, REPEATED)
• Tests both Storage Write API and read schema conversion paths

---

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

• Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
• Update CHANGES.md with noteworthy changes.
• If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

[gemini-code-assist]

Summary of Changes

Hello @enzomaruffa, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the Apache Beam Python SDK's BigQuery I/O capabilities by introducing support for DATE, DATETIME, and JSON BigQuery types. This change is crucial for users migrating Dataflow pipelines to Python, particularly those leveraging the high-performance Storage Write API with BEAM_ROW format, as it resolves previous ValueError issues encountered with these types. The implementation pragmatically maps these BigQuery types to Python strings, ensuring immediate compatibility and unblocking production pipelines.

Highlights

• BigQuery Type Support: Adds support for DATE, DATETIME, and JSON BigQuery types within the Python SDK's BigQuery I/O type mappings.
• Storage Write API Compatibility: Enables these new types to be used seamlessly with the BigQuery Storage Write API when utilizing the BEAM_ROW format (i.e., use_beam_io_types=True).
• Python Type Mapping: The DATE, DATETIME, and JSON BigQuery types are pragmatically mapped to Python's str type, consistent with existing GEOGRAPHY type handling.
• Comprehensive Testing: Extensive unit tests have been added for each new type, covering schema conversion, type mapping verification, bq_field_to_type function, convert_to_usertype, and BeamSchemaConversionDoFn processing across various modes (REQUIRED, NULLABLE, REPEATED).

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

[github-actions]

Assigning reviewers:

R: @shunping for label python.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

[Abacn]

I remember DATE/DATETIME types were already supported in Python BigQuery Storage API. cc: @ahmedabu98

If there is ineed still gap, we should use Python language type rather than str. Similar treatment was done in Python JdbcIO (Jdbc Date type and time type)

[enzomaruffa]

Hi @Abacn! From our experience, these are definitely missing. I am not sure if it has to do something also missing in under-the-hood xlang implementation that uses the Java SDK, but Beam would block us before even getting it deployed to Dataflow with these types being unsupported.

I'd love to be wrong and have a simpler solution instead of this, though 🙏

[ahmedabu98]

DATE and DATETIME aren't supported yet (only TIMESTAMP)

But I'm not very keen on defaulting them to Strings, it's not very robust. We should choose a Python native type that is closer to what a DATE/DATETIME actually is.

Some options are discussed in https://s.apache.org/beam-timestamp-strategy (under "Python Nanosecond Support" --> "DateTime"). It'll take more work but I think it's a better longterm solution for Beam.

[scwe]

Hi all 👋 I've been working on this stuff with @enzomaruffa for a while. I totally understand that this solution very hacky but realistically we don't have the resources to dedicate to building anything more robust.

So my question essentially becomes: is there any way we can get support for all BQ types bumped in priority on your end? We've been following along/waiting for over a year at this point hoping that it would be solved and would love to be able to use python dataflow in production without maintaining our own hacky fork.

[claudevdm]

Maybe a middleground is to allow users to pass type_overrides, that override the default type mappings? This way beam can add official mappings later on, and using strings now is possible without needing future breaking changes.

[enzomaruffa]

@claudevdm sounds pretty reasonable. I have added the work here doing this + updated the PR description. We'll run some more tests with these tweaks internally to make sure they work and I can report back here later.

[github-actions]

Reminder, please take a look at this pr: @shunping

[github-actions]

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @claudevdm for label python.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)