security(weave_query): Upgrade Pyarrow to 17.0.0 (#3617)

* chore: resolve ArrowInvalid error from ambiguous null specification

* chore: update all ListArray.from_arrays calls to not use mask

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

* chore: remove debug to trigger ci

* chore: add dummy commit to trigger frontend-tests

* chore: remove dummy commit changes

Author: domphan-wandb

weave_query/requirements.legacy.txt

@@ -6,8 +6,7 @@ typing_extensions>=4.0.0
 
 # Definitely need arrow
 # TODO: Colab has 9.0.0, can we support?
-# TODO: 17.0.0 breaks a bunch of tests - can we move this requirement to just the engine?
-pyarrow>=14.0.1,<17.0.0
+pyarrow==17.0.0
 
 # pydantic integration, and required by openai anyway
 openai>=1.0.0

weave_query/weave_query/arrow/arrow.py

@@ -340,15 +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),
+            mask=pa.compute.is_null(data)
         )
     else:
         # We have a column of refs
@@ -474,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,10 +33,11 @@ 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):
-        return pa.ListArray.from_arrays(
-            offsets_starting_at_zero(array),
+        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()),
+            mask=pa.compute.invert(array.is_valid())
         )
     elif array.type == pa.string():
         return pc.dictionary_encode(array)

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
 
@@ -198,6 +199,7 @@ def _concatenate_lists(
         l2_arrow_data.flatten(), l2.object_type.object_type, l2._artifact
     )
     concatted_values = _concatenate(self_values, other_values, depth=depth + 1)
+
     new_offsets = pa_concat_arrays(
         [
             offsets_starting_at_zero(l1_arrow_data)[:-1],
@@ -207,17 +209,18 @@ def _concatenate_lists(
             ).cast(pa.int32()),
         ]
     )
+
+    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)
+        ])
+    )
+
     return ArrowWeaveList(
-        pa.ListArray.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),
-                ]
-            ),
-        ),
+        result_array,
         types.List(
             types.merge_types(l1.object_type.object_type, l2.object_type.object_type)
         ),
@@ -512,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,
@@ -539,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,7 +308,7 @@ def none_unboxer(iterator: typing.Iterable):
             mapper._object_type,
             py_objs_already_mapped,
         )
-        return pa.ListArray.from_arrays(
+        return safe_list_array_from_arrays(
             offsets, new_objs, mask=pa.array(mask, type=pa.bool_())
         )
     elif pa.types.is_temporal(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,
@@ -800,14 +801,16 @@ def _map_column(
             items: ArrowWeaveList = ArrowWeaveList(
                 arr.flatten(), self.object_type.object_type, self._artifact
             )._map_column(fn, pre_fn, path + (PathItemList(),))
-            # print("SELF OBJECT TYPE", self.object_type)
-            # print("SELF ARROW DATA TYPE", self._arrow_data.type)
+
+            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)
+            )
+
             with_mapped_children = ArrowWeaveList(
-                pa.ListArray.from_arrays(
-                    offsets_starting_at_zero(self._arrow_data),
-                    items._arrow_data,
-                    mask=pa.compute.is_null(arr),
-                ),
+                result_array,
                 self.object_type.__class__(items.object_type),
                 self._artifact,
                 invalid_reason=items._invalid_reason,

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, mask=pa.compute.is_null(a)
-    )
+    unflattened = safe_list_array_from_arrays(new_offsets, new_data, mask=pa.compute.is_null(a))
 
     return ArrowWeaveList(
         unflattened,