Fix/upsert with nulls

[kevinjqliu]

Rationale for this change

Are these changes tested?

Are there any user-facing changes?

[gemini-code-assist]

Summary of Changes

Hello @kevinjqliu, 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 resolves an issue where the upsert functionality did not correctly handle NULL values in join columns, leading to inconsistent data merges. The core change introduces null-safe equality, allowing NULL to match NULL and NaN to match NaN during the upsert process. This was achieved by modifying the internal logic for generating match filters and implementing a custom join mechanism to ensure accurate row identification, complemented by thorough testing and updated documentation.

Highlights

• Null-Safe Equality for Upsert: Implemented null-safe equality semantics for the upsert operation, ensuring that NULL values in join columns correctly match other NULL values, and NaN values match other NaN values. This aligns with SQL's <=> operator.
• Updated Match Filter Generation: The create_match_filter utility function was enhanced to generate appropriate IsNull and IsNaN expressions when constructing match filters for join columns containing null or NaN values, respectively.
• Custom Python-Based Join Logic: Refactored the row matching logic within get_rows_to_update to use a custom Python-based join instead of PyArrow's default inner join. This change was necessary because PyArrow's join semantics do not inherently support null-safe equality, which was crucial for this fix.
• Comprehensive Test Coverage: Added new unit tests for create_match_filter to validate its behavior with various combinations of null and NaN values in single and multi-column join conditions. An end-to-end integration test was also included to confirm the correct upsert functionality with nulls in join columns.
• Documentation Enhancement: Updated the docstring for the upsert function in pyiceberg/table/__init__.py to clearly explain the new null-safe equality behavior, including examples and a note on how to achieve standard SQL equality if desired.

🧠 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. ↩

[gemini-code-assist]

Code Review

This pull request correctly implements null-safe equality for upsert operations, which is a great improvement. The documentation has been updated clearly, and new tests verify the behavior with nulls and NaNs. However, I've identified a significant performance concern in the implementation of the null-safe join in get_rows_to_update. The current approach may not scale well for large datasets. My review includes a suggestion to address this by leveraging PyArrow's native capabilities for a more efficient implementation.

[gemini-code-assist]

The current implementation for the null-safe join converts the PyArrow tables to Python lists and dictionaries using to_pylist(). This can be very inefficient and memory-intensive for large tables, potentially leading to performance bottlenecks or out-of-memory errors.

PyArrow's join method supports null-safe equality since version 7.0.0 via the null_matching_behavior='equal' parameter. Using this would be much more performant as it keeps the operations within PyArrow's memory space.

I suggest reverting to the PyArrow join and adding this parameter.

Suggested change

	# Step 3: Perform an inner join to find which rows from source exist in target.
	# We use a Python-based join instead of PyArrow's join because PyArrow ignores NULL values
	# (NULL == NULL returns UNKNOWN in SQL semantics). We want null-safe equality where NULL == NULL is TRUE.
	source_keys = {tuple(row[col] for col in join_cols): row[SOURCE_INDEX_COLUMN_NAME] for row in source_index.to_pylist()}
	target_keys = {tuple(row[col] for col in join_cols): row[TARGET_INDEX_COLUMN_NAME] for row in target_index.to_pylist()}
	matching_indices = [(s, t) for key, s in source_keys.items() if (t := target_keys.get(key)) is not None]
	# Step 4: Compare all rows using Python
	to_update_indices = []
	for source_idx, target_idx in zip(
	matching_indices[SOURCE_INDEX_COLUMN_NAME].to_pylist(),
	matching_indices[TARGET_INDEX_COLUMN_NAME].to_pylist(),
	strict=True,
	):
	for source_idx, target_idx in matching_indices:
	source_row = source_table.slice(source_idx, 1)
	target_row = target_table.slice(target_idx, 1)
	# Step 3: Perform an inner join to find which rows from source exist in target.
	# PyArrow's join operator can perform null-safe joins.
	matching_indices = source_index.join(target_index, keys=list(join_cols_set), join_type="inner", null_matching_behavior="equal")
	# Step 4: Compare all rows using Python
	to_update_indices = []
	for source_idx, target_idx in zip(
	matching_indices[SOURCE_INDEX_COLUMN_NAME].to_pylist(),
	matching_indices[TARGET_INDEX_COLUMN_NAME].to_pylist(),
	strict=True,
	):
	source_row = source_table.slice(source_idx, 1)
	target_row = target_table.slice(target_idx, 1)