From d931bd06359f347b4422f945d33a93995c0ba0f2 Mon Sep 17 00:00:00 2001 From: Harold Wang Date: Tue, 3 May 2022 14:59:07 -0700 Subject: [PATCH] fix col // op in eager mode Summary: add support that allows column and dataframe floor division operations in OSS (eager mode). Differential Revision: D36112289 fbshipit-source-id: 7420fb4c9890e52d6874aa46b1bedad457acb224 --- torcharrow/test/test_dataframe.py | 7 +++++++ torcharrow/velox_rt/numerical_column_cpu.py | 8 ++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/torcharrow/test/test_dataframe.py b/torcharrow/test/test_dataframe.py index b4e1f1932..d4184db3c 100644 --- a/torcharrow/test/test_dataframe.py +++ b/torcharrow/test/test_dataframe.py @@ -637,6 +637,13 @@ def base_test_operators(self): colx = ta.column([3, 4, 5], device=self.device) self.assertEqual(list(dfx % colx), [(0.0, 2.0), (3.0, 3.0), (4.0, 3.0)]) + # // + dfx = ta.dataframe({"a": [3, 4, 6], "b": [6, 8, 7]}, device=self.device) + self.assertEqual(list(dfx["a"] // dfx), [(1, 0), (1, 0), (1, 0)]) + self.assertEqual(list(dfx["b"] // dfx), [(2, 1), (2, 1), (1, 1)]) + self.assertEqual(list(dfx // dfx["a"]), [(1, 2), (1, 2), (1, 1)]) + self.assertEqual(list(dfx // dfx["b"]), [(0, 1), (0, 1), (0, 1)]) + def base_test_python_comparison_ops(self): # Use a dtype of list to prevent fast path through numerical # column operators to ensure we are testing the generic python diff --git a/torcharrow/velox_rt/numerical_column_cpu.py b/torcharrow/velox_rt/numerical_column_cpu.py index a14576885..7d59eaff4 100644 --- a/torcharrow/velox_rt/numerical_column_cpu.py +++ b/torcharrow/velox_rt/numerical_column_cpu.py @@ -373,8 +373,8 @@ def __floordiv__(self, other): This behavior is different from __truediv__, but is consistent with Pytorch. """ return self._rethrow_zero_division_error( - lambda: self._checked_arithmetic_op_call( - other, "floordiv", operator.floordiv + lambda: self._checked_arithmetic_op_call_with_df( + other, "floordiv", operator.floordiv, "__rfloordiv__" ) ) @@ -390,8 +390,8 @@ def __rfloordiv__(self, other): This behavior is different from __rtruediv__, but is consistent with Pytorch. """ return self._rethrow_zero_division_error( - lambda: self._checked_arithmetic_op_call( - other, "rfloordiv", Column._swap(operator.floordiv) + lambda: self._checked_arithmetic_op_call_with_df( + other, "rfloordiv", Column._swap(operator.floordiv), "__floordiv__" ) )