fix: upsert with null values in join columns

Author: mdwint

pyiceberg/table/upsert_util.py

@@ -14,38 +14,61 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-import functools
-import operator
+from math import isnan
+from typing import Any
 
 import pyarrow as pa
 from pyarrow import Table as pyarrow_table
 from pyarrow import compute as pc
 
 from pyiceberg.expressions import (
     AlwaysFalse,
+    And,
     BooleanExpression,
     EqualTo,
     In,
+    IsNaN,
+    IsNull,
     Or,
 )
 
 
 def create_match_filter(df: pyarrow_table, join_cols: list[str]) -> BooleanExpression:
     unique_keys = df.select(join_cols).group_by(join_cols).aggregate([])
+    filters = []
 
     if len(join_cols) == 1:
-        return In(join_cols[0], unique_keys[0].to_pylist())
+        column = join_cols[0]
+        values = set(unique_keys[0].to_pylist())
+
+        if None in values:
+            filters.append(IsNull(column))
+            values.remove(None)
+
+        if nans := {v for v in values if isinstance(v, float) and isnan(v)}:
+            filters.append(IsNaN(column))
+            values -= nans
+
+        filters.append(In(column, values))
+    else:
+
+        def equals(column: str, value: Any) -> BooleanExpression:
+            if value is None:
+                return IsNull(column)
+
+            if isinstance(value, float) and isnan(value):
+                return IsNaN(column)
+
+            return EqualTo(column, value)
+
+        filters = [And(*[equals(col, row[col]) for col in join_cols]) for row in unique_keys.to_pylist()]
+
+    if len(filters) == 0:
+        return AlwaysFalse()
+    elif len(filters) == 1:
+        return filters[0]
     else:
-        filters = [
-            functools.reduce(operator.and_, [EqualTo(col, row[col]) for col in join_cols]) for row in unique_keys.to_pylist()
-        ]
-
-        if len(filters) == 0:
-            return AlwaysFalse()
-        elif len(filters) == 1:
-            return filters[0]
-        else:
-            return Or(*filters)
+        return Or(*filters)
 
 
 def has_duplicate_rows(df: pyarrow_table, join_cols: list[str]) -> bool:

tests/table/test_upsert.py

@@ -710,6 +710,46 @@ def test_upsert_with_nulls(catalog: Catalog) -> None:
         schema=schema,
     )
 
+    # upsert table with null value
+    data_with_null = pa.Table.from_pylist(
+        [
+            {"foo": None, "bar": 1, "baz": False},
+        ],
+        schema=schema,
+    )
+    upd = table.upsert(data_with_null, join_cols=["foo"])
+    assert upd.rows_updated == 0
+    assert upd.rows_inserted == 1
+    assert table.scan().to_arrow() == pa.Table.from_pylist(
+        [
+            {"foo": None, "bar": 1, "baz": False},
+            {"foo": "apple", "bar": 7, "baz": False},
+            {"foo": "banana", "bar": None, "baz": False},
+        ],
+        schema=schema,
+    )
+
+    # upsert table with null and non-null values, in two join columns
+    data_with_null = pa.Table.from_pylist(
+        [
+            {"foo": None, "bar": 1, "baz": True},
+            {"foo": "lemon", "bar": None, "baz": False},
+        ],
+        schema=schema,
+    )
+    upd = table.upsert(data_with_null, join_cols=["foo", "bar"])
+    assert upd.rows_updated == 1
+    assert upd.rows_inserted == 1
+    assert table.scan().to_arrow() == pa.Table.from_pylist(
+        [
+            {"foo": "lemon", "bar": None, "baz": False},
+            {"foo": None, "bar": 1, "baz": True},
+            {"foo": "apple", "bar": 7, "baz": False},
+            {"foo": "banana", "bar": None, "baz": False},
+        ],
+        schema=schema,
+    )
+
 
 def test_transaction(catalog: Catalog) -> None:
     """Test the upsert within a Transaction. Make sure that if something fails the entire Transaction is