edponce commented on a change in pull request #10349:
URL: https://github.com/apache/arrow/pull/10349#discussion_r696160814
##########
File path: python/pyarrow/tests/test_compute.py
##########
@@ -1299,6 +1301,98 @@ def test_arithmetic_multiply():
assert result.equals(expected)
+@pytest.mark.parametrize("ty", ["round", "mround"])
+def test_round_to_integer(ty):
+ if ty == "round":
+ func = pc.round
+ RoundOptions = pc.RoundOptions
+ elif ty == "mround":
+ func = pc.mround
+ RoundOptions = pc.MRoundOptions
+
+ approx_equals = partial(np.isclose, equal_nan=True)
+ round_modes = [
+ "down", "up", "towards_zero", "towards_infinity", "half_down",
+ "half_up", "half_towards_zero", "half_towards_infinity",
+ "half_to_even", "half_to_odd",
+ ]
+ arr = pa.array([3.2, 3.5, 3.7, 4.5, -3.2, -3.5, -3.7, None])
+ for round_mode in round_modes:
+ options = RoundOptions(round_mode=round_mode)
+ result = func(arr, options=options).to_numpy(zero_copy_only=False)
+ if round_mode == "down":
+ expected_ = [3, 3, 3, 4, -4, -4, -4, None]
+ elif round_mode == "up":
+ expected_ = [4, 4, 4, 5, -3, -3, -3, None]
+ elif round_mode == "towards_zero":
+ expected_ = [3, 3, 3, 4, -3, -3, -3, None]
+ elif round_mode == "towards_infinity":
+ expected_ = [4, 4, 4, 5, -4, -4, -4, None]
+ elif round_mode == "half_down":
+ expected_ = [3, 3, 4, 4, -3, -4, -4, None]
+ elif round_mode == "half_up":
+ expected_ = [3, 4, 4, 5, -3, -3, -4, None]
+ elif round_mode == "half_towards_zero":
+ expected_ = [3, 3, 4, 4, -3, -3, -4, None]
+ elif round_mode == "half_towards_infinity":
+ expected_ = [3, 4, 4, 5, -3, -4, -4, None]
+ elif round_mode == "half_to_even":
+ expected_ = [3, 4, 4, 4, -3, -4, -4, None]
+ elif round_mode == "half_to_odd":
+ expected_ = [3, 3, 4, 5, -3, -3, -4, None]
+
+ expected = pa.array(expected_).to_numpy(zero_copy_only=False)
+ assert all(map(approx_equals, result, expected))
+
+
+def test_round():
+ approx_equals = partial(np.isclose, equal_nan=True)
+ round_ndigits = [-2, -1, 0, 1, 2]
+ arr = pa.array([320, 3.5, 3.075, 4.5, -3.212, -35.1234, -3.045, None])
+ for ndigits in round_ndigits:
+ options = pc.RoundOptions(
+ ndigits=ndigits, round_mode="half_towards_infinity")
+ result = pc.round(arr, options=options).to_numpy(zero_copy_only=False)
+ if ndigits == -2:
+ expected_ = [300, 0, 0, 0, 0, 0, 0, None]
+ elif ndigits == -1:
+ expected_ = [320, 0, 0, 0, 0, -40, 0, None]
+ elif ndigits == 0:
+ expected_ = [320, 4, 3, 5, -3, -35, -3, None]
+ elif ndigits == 1:
+ expected_ = [320, 3.5, 3.1, 4.5, -3.2, -35.1, -3, None]
+ elif ndigits == 2:
+ expected_ = [320, 3.5, 3.08, 4.5, -3.21, -35.12, -3.05, None]
+
+ expected = pa.array(expected_).to_numpy(zero_copy_only=False)
+ assert all(map(approx_equals, result, expected))
+
+
+def test_mround():
+ approx_equals = partial(np.isclose, equal_nan=True)
+ multiples = [-2, -0.05, 0.1, 0, 10, 100]
+ arr = pa.array([320, 3.5, 3.075, 4.5, -3.212, -35.1234, -3.045, None])
+ for multiple in multiples:
+ options = pc.MRoundOptions(
+ multiple=multiple, round_mode="half_towards_infinity")
+ result = pc.mround(arr, options=options).to_numpy(zero_copy_only=False)
+ if multiple == -2:
+ expected_ = [320, 4, 4, 4, -4, -36, -4, None]
+ elif multiple == -0.05:
+ expected_ = [320, 3.5, 3.1, 4.5, -3.2, -35.1, -3.05, None]
+ elif multiple == 0.1:
+ expected_ = [320, 3.5, 3.1, 4.5, -3.2, -35.1, -3, None]
+ elif multiple == 0:
+ expected_ = [0, 0, 0, 0, 0, 0, 0, None]
+ elif multiple == 10:
+ expected_ = [320, 0, 0, 0, 0, -40, 0, None]
+ elif multiple == 100:
+ expected_ = [300, 0, 0, 0, 0, 0, 0, None]
+
+ expected = pa.array(expected_).to_numpy(zero_copy_only=False)
Review comment:
`to_numpy` has as default `zero_copy_only=True` but the implicit array
conversion `__array__` uses `zero_copy_only=False`. I followed your suggestion
and used implicit casts Thanks for this suggestion!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
