lidavidm commented on a change in pull request #10349:
URL: https://github.com/apache/arrow/pull/10349#discussion_r693986079



##########
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:
       Is there a reason why we have to unset zero_copy_only? False is the 
default value anyways.

##########
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:
       And this could just be `np.array(...)` instead of bouncing through 
PyArrow (though as noted below, you don't even need to do that).

##########
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)
+        assert all(map(approx_equals, result, expected))

Review comment:
       nit: it would be more natural to use np.all(np.isclose(arr1, arr2, 
equal_nan=True)) instead of using map/partial/all.




-- 
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


Reply via email to