def unary_col(op, v):
"""
interpretor for executing unary operator expressions on columnars
"""
if op == "+":
return v
if op == "-":
return compute.subtract(0.0, v)
if op.lower() == "not":
return compute.invert(v)
raise Exception("unary op not implemented")
def reset_sliced_list_offset(array: pa.ListArray):
"""Return the same pyarrow.ListArray but with array.offset == 0 for compatibility with cast"""
if array.offset == 0:
return array
elif len(array) == 0:
return array.values.slice(0, 0)
else:
values_offset = array.offsets[
0] # the relevant values start at this index
new_values = array.values.slice(values_offset.as_py(
)) # get the values to start at the right position
new_offsets = pc.subtract(
array.offsets, values_offset) # update the offsets accordingly
return pa.ListArray.from_arrays(new_offsets, new_values)
def sum_of_squared_diffs_from_mean(
self,
on: KeyFn,
ignore_nulls: bool,
mean: Optional[U] = None,
) -> Optional[U]:
import pyarrow.compute as pac
if mean is None:
# If precomputed mean not given, we compute it ourselves.
mean = self.mean(on, ignore_nulls)
if mean is None:
return None
return self._apply_arrow_compute(
lambda col, skip_nulls: pac.sum(
pac.power(pac.subtract(col, mean), 2),
skip_nulls=skip_nulls,
),
on,
ignore_nulls,
)
def test_arithmetic_subtract():
left = pa.array([1, 2, 3, 4, 5])
right = pa.array([0, -1, 1, 2, 3])
result = pc.subtract(left, right)
expected = pa.array([1, 3, 2, 2, 2])
assert result.equals(expected)
