fix(weave_query): Check for offsets with nulls before using a mask

Author: domphan-wandb

weave_query/weave_query/arrow/arrow.py

@@ -340,14 +340,18 @@ def rewrite_weavelist_refs(arrow_data, object_type, source_artifact, target_arti
         data = arrow_data
         if isinstance(arrow_data, pa.ChunkedArray):
             data = arrow_data.combine_chunks()
-        return pa.ListArray.from_arrays(
-            offsets_starting_at_zero(data),
+
+        new_offsets = offsets_starting_at_zero(data)
+
+        return safe_list_array_from_arrays(
+            new_offsets,
             rewrite_weavelist_refs(
                 data.flatten(),
                 object_type.object_type,
                 source_artifact,
                 target_artifact,
-            )
+            ),
+            mask=pa.compute.is_null(data)
         )
     else:
         # We have a column of refs
@@ -473,6 +477,33 @@ def safe_coalesce(*arrs: pa.Array):
         result = pa.compute.if_else(safe_is_null(result), arr, result)
     return result
 
+def safe_list_array_from_arrays(offsets, values, mask=None):
+    # In PyArrow 17.0.0, ListArray.from_arrays() was updated to check for ambiguity
+    # between null offsets and masks via offsets.null_count(). When concatenating
+    # offset arrays, PyArrow uses lazy evaluation which can cause this check to see
+    # an intermediate state and incorrectly raise an "ambiguous to specify both 
+    # validity map and offsets with nulls" error.
+
+    # Convert offsets to PyArrow array if it isn't already
+    if not isinstance(offsets, pa.Array):
+        offsets = pa.array(offsets)
+
+    has_offsets_with_nulls = offsets.null_count > 0
+
+    if has_offsets_with_nulls:
+        # Can't use both a mask and an offset with nulls
+        result_array = pa.ListArray.from_arrays(
+            offsets,
+            values
+        )
+        return pc.if_else(mask, None, result_array)
+    else:
+        return pa.ListArray.from_arrays(
+            offsets,
+            values,
+            mask=mask
+        )
+
 
 def arrow_zip(*arrs: pa.Array) -> pa.Array:
     n_arrs = len(arrs)

weave_query/weave_query/arrow/arrow_tags.py

@@ -5,7 +5,10 @@
 
 from weave_query import weave_types as types
 from weave_query.arrow import convert
-from weave_query.arrow.arrow import offsets_starting_at_zero
+from weave_query.arrow.arrow import (
+    offsets_starting_at_zero,
+    safe_list_array_from_arrays
+)
 from weave_query.language_features.tagging import (
     process_opdef_output_type,
     tag_store,
@@ -30,11 +33,12 @@ def recursively_encode_pyarrow_strings_as_dictionaries(array: pa.Array) -> pa.Ar
             mask=pa.compute.invert(array.is_valid()),
         )
     elif pa.types.is_list(array.type):
-        result_array = pa.ListArray.from_arrays(
-            offsets_starting_at_zero(array),
-            recursively_encode_pyarrow_strings_as_dictionaries(array.flatten())
+        new_offsets = offsets_starting_at_zero(array)
+        return safe_list_array_from_arrays(
+            new_offsets,
+            recursively_encode_pyarrow_strings_as_dictionaries(array.flatten()),
+            mask=pa.compute.invert(array.is_valid())
         )
-        return pc.if_else(pa.compute.invert(array.is_valid()), None, result_array)
     elif array.type == pa.string():
         return pc.dictionary_encode(array)
     else:

weave_query/weave_query/arrow/concat.py

@@ -25,6 +25,7 @@
     make_vec_none,
     offsets_starting_at_zero,
     unsafe_awl_construction,
+    safe_list_array_from_arrays
 )
 from weave_query.language_features.tagging import tagged_value_type
 
@@ -209,21 +210,15 @@ def _concatenate_lists(
         ]
     )
 
-    result_array = pa.ListArray.from_arrays(
+    result_array = safe_list_array_from_arrays(
         new_offsets,
         concatted_values._arrow_data,
+        mask=pa_concat_arrays([
+            pa.compute.is_null(l1._arrow_data),
+            pa.compute.is_null(l2._arrow_data)
+        ])
     )
 
-    # In pyarrow 17.0.0, we can only use either null offsets OR provide a mask,
-    # so track the null values to convert empty lists (that were null before
-    # conversion) and convert those empty lists to None before returning the AWL
-    combined_nulls = pa_concat_arrays([
-        pa.compute.is_null(l1._arrow_data),
-        pa.compute.is_null(l2._arrow_data)
-    ])
-
-    result_array = pc.if_else(combined_nulls, None, result_array)
-
     return ArrowWeaveList(
         result_array,
         types.List(
@@ -520,7 +515,6 @@ def _concatenate(
                             pa.array(np.zeros(len(other), dtype=np.int8)),
                         ]
                     ),
-                    # offsets=pa.array(np.arange(len(self), dtype=np.int32)),
                     offsets=pa_concat_arrays(
                         [
                             self_offsets,
@@ -547,7 +541,6 @@ def _concatenate(
                         pa.compute.equal(other_type_codes, other_i).cast(pa.int8()),
                     ]
                 ),
-                # offsets=pa.array(np.arange(len(other), dtype=np.int32)),
                 offsets=pa_concat_arrays(
                     [
                         pa.array(np.zeros(len(self), dtype=np.int32)),

weave_query/weave_query/arrow/convert.py

@@ -21,6 +21,7 @@
     ArrowWeaveList,
     PathType,
     unsafe_awl_construction,
+    safe_list_array_from_arrays
 )
 from weave_query.language_features.tagging import tag_store, tagged_value_type
 
@@ -307,9 +308,9 @@ def none_unboxer(iterator: typing.Iterable):
             mapper._object_type,
             py_objs_already_mapped,
         )
-        result_array = pa.ListArray.from_arrays(offsets, new_objs)
-
-        return pc.if_else(pa.array(mask, type=pa.bool_()), None, result_array)
+        return safe_list_array_from_arrays(
+            offsets, new_objs, mask=pa.array(mask, type=pa.bool_())
+        )
     elif pa.types.is_temporal(pyarrow_type):
         if py_objs_already_mapped:
             return pa.array(py_objs, type=pyarrow_type)

weave_query/weave_query/arrow/list_.py

@@ -32,6 +32,7 @@
     offsets_starting_at_zero,
     pretty_print_arrow_type,
     safe_is_null,
+    safe_list_array_from_arrays
 )
 from weave_query.language_features.tagging import (
     tag_store,
@@ -802,12 +803,14 @@ def _map_column(
             )._map_column(fn, pre_fn, path + (PathItemList(),))
             # print("SELF OBJECT TYPE", self.object_type)
             # print("SELF ARROW DATA TYPE", self._arrow_data.type)
-            result_array = pa.ListArray.from_arrays(
-                offsets_starting_at_zero(self._arrow_data),
-                items._arrow_data
+
+            new_offsets = offsets_starting_at_zero(self._arrow_data)
+            result_array = safe_list_array_from_arrays(
+                new_offsets,
+                items._arrow_data,
+                pa.compute.is_null(arr)
             )
 
-            result_array = pc.if_else(pa.compute.is_null(arr), None, result_array)
             with_mapped_children = ArrowWeaveList(
                 result_array,
                 self.object_type.__class__(items.object_type),

weave_query/weave_query/ops_arrow/arraylist_ops.py

@@ -4,7 +4,11 @@
 from pyarrow import compute as pc
 
 from weave_query import weave_types as types
-from weave_query.arrow.arrow import ArrowWeaveListType, arrow_as_array
+from weave_query.arrow.arrow import (
+    ArrowWeaveListType,
+    arrow_as_array,
+    safe_list_array_from_arrays
+)
 from weave_query.arrow.list_ import ArrowWeaveList
 from weave_query.decorator_arrow_op import arrow_op
 from weave_query.language_features.tagging import tagged_value_type
@@ -350,9 +354,7 @@ def dropna(self):
     cumulative_non_null_counts = pa.compute.cumulative_sum(non_null)
     new_offsets = cumulative_non_null_counts.take(pa.compute.subtract(end_indexes, 1))
     new_offsets = pa.concat_arrays([start_indexes[:1], new_offsets])
-    unflattened = pa.ListArray.from_arrays(new_offsets, new_data)
-    unflattened = pa.compute.if_else(pa.compute.is_null(a), None, unflattened)
-
+    unflattened = safe_list_array_from_arrays(new_offsets, new_data, mask=pa.compute.is_null(a))
 
     return ArrowWeaveList(
         unflattened,