Add New Bitwise Int Operators

temporian/core/operators/binary/BUILD

@@ -14,6 +14,7 @@ py_library(
         ":arithmetic",
         ":logical",
         ":relational",
+        ":bitwise",
     ],
 )
 
@@ -73,3 +74,18 @@ py_library(
         "//temporian/core/data:schema",
     ],
 )
+
+py_library(
+    name = "bitwise",
+    srcs = ["bitwise.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":base",
+        "//temporian/core:compilation",
+        "//temporian/core:operator_lib",
+        "//temporian/core:typing",
+        "//temporian/core/data:dtype",
+        "//temporian/core/data:node",
+        "//temporian/core/data:schema",
+    ],
+)

temporian/core/operators/binary/__init__.py

@@ -56,3 +56,16 @@
     logical_or,
     logical_xor,
 )
+
+from temporian.core.operators.binary.bitwise import (
+    BitwiseAndOperator,
+    BitwiseOrOperator,
+    BitwiseXorOperator,
+    LeftShiftOperator,
+    RightShiftOperator,
+    bitwise_and,
+    bitwise_or,
+    bitwise_xor,
+    left_shift,
+    right_shift,
+)

temporian/core/operators/binary/bitwise.py

@@ -0,0 +1,152 @@
+# Copyright 2021 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Binary logic operators classes and public API function definitions."""
+
+from temporian.core import operator_lib
+from temporian.core.compilation import compile
+from temporian.core.data.dtype import DType
+from temporian.core.data.node import EventSetNode
+from temporian.core.data.schema import FeatureSchema
+from temporian.core.operators.binary.base import BaseBinaryOperator
+from temporian.core.typing import EventSetOrNode
+
+
+class BaseBitwiseOperator(BaseBinaryOperator):
+    OP_NAME = ""
+
+    def __init__(self, input_1: EventSetNode, input_2: EventSetNode):
+        super().__init__(input_1, input_2)
+
+        # Check that all features are Integer
+        # Note: Assuming that input_1 and input_2 features have the same dtype
+        for feature in input_1.schema.features:
+            if feature.dtype != DType.INT32 and feature.dtype != DType.INT64:
+                raise ValueError(
+                    "Bitwise operators only support INT32 or INT64 types, but "
+                    f"feature {feature.name} has dtype {feature.dtype}"
+                )
+
+    def output_feature_dtype(
+        self, feature_1: FeatureSchema, feature_2: FeatureSchema
+    ) -> DType:
+        if feature_1.dtype == DType.INT64 or feature_2.dtype == DType.INT64:
+            return DType.INT64
+        else:
+            return DType.INT32
+
+    @classmethod
+    def operator_def_key(cls) -> str:
+        return cls.OP_NAME.upper()
+
+    @property
+    def prefix(self) -> str:
+        return self.OP_NAME.lower()
+
+
+class BitwiseAndOperator(BaseBitwiseOperator):
+    OP_NAME = "bitwise_and"
+
+
+class BitwiseOrOperator(BaseBitwiseOperator):
+    OP_NAME = "bitwise_or"
+
+
+class BitwiseXorOperator(BaseBitwiseOperator):
+    OP_NAME = "bitwise_xor"
+
+
+class LeftShiftOperator(BaseBitwiseOperator):
+    OP_NAME = "left_shift"
+
+
+class RightShiftOperator(BaseBitwiseOperator):
+    OP_NAME = "right_shift"
+
+
+@compile
+def bitwise_and(
+    input_1: EventSetOrNode,
+    input_2: EventSetOrNode,
+) -> EventSetOrNode:
+    assert isinstance(input_1, EventSetNode)
+    assert isinstance(input_2, EventSetNode)
+
+    return BitwiseAndOperator(
+        input_1=input_1,
+        input_2=input_2,
+    ).outputs["output"]
+
+
+@compile
+def bitwise_or(
+    input_1: EventSetOrNode,
+    input_2: EventSetOrNode,
+) -> EventSetOrNode:
+    assert isinstance(input_1, EventSetNode)
+    assert isinstance(input_2, EventSetNode)
+
+    return BitwiseOrOperator(
+        input_1=input_1,
+        input_2=input_2,
+    ).outputs["output"]
+
+
+@compile
+def bitwise_xor(
+    input_1: EventSetOrNode,
+    input_2: EventSetOrNode,
+) -> EventSetOrNode:
+    assert isinstance(input_1, EventSetNode)
+    assert isinstance(input_2, EventSetNode)
+
+    return BitwiseXorOperator(
+        input_1=input_1,
+        input_2=input_2,
+    ).outputs["output"]
+
+
+@compile
+def left_shift(
+    input_1: EventSetOrNode,
+    input_2: EventSetOrNode,
+) -> EventSetOrNode:
+    assert isinstance(input_1, EventSetNode)
+    assert isinstance(input_2, EventSetNode)
+
+    return LeftShiftOperator(
+        input_1=input_1,
+        input_2=input_2,
+    ).outputs["output"]
+
+
+@compile
+def right_shift(
+    input_1: EventSetOrNode,
+    input_2: EventSetOrNode,
+) -> EventSetOrNode:
+    assert isinstance(input_1, EventSetNode)
+    assert isinstance(input_2, EventSetNode)
+
+    return RightShiftOperator(
+        input_1=input_1,
+        input_2=input_2,
+    ).outputs["output"]
+
+
+operator_lib.register_operator(BitwiseAndOperator)
+operator_lib.register_operator(BitwiseOrOperator)
+operator_lib.register_operator(BitwiseXorOperator)
+operator_lib.register_operator(LeftShiftOperator)
+operator_lib.register_operator(RightShiftOperator)

temporian/core/operators/scalar/BUILD

@@ -13,6 +13,7 @@ py_library(
     deps = [
         ":arithmetic_scalar",
         ":relational_scalar",
+        ":bitwise_scalar"
     ],
 )
 
@@ -57,3 +58,17 @@ py_library(
         "//temporian/core/data:schema",
     ],
 )
+
+py_library(
+    name = "bitwise_scalar",
+    srcs = ["bitwise_scalar.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":base",
+        "//temporian/core:compilation",
+        "//temporian/core:operator_lib",
+        "//temporian/core:typing",
+        "//temporian/core/data:dtype",
+        "//temporian/core/data:node",
+    ],
+)

temporian/core/operators/scalar/__init__.py

@@ -47,3 +47,16 @@
     LessEqualScalarOperator,
     LessScalarOperator,
 )
+
+from temporian.core.operators.scalar.bitwise_scalar import (
+    bitwise_and_scalar,
+    bitwise_or_scalar,
+    bitwise_xor_scalar,
+    left_shift_scalar,
+    right_shift_scalar,
+    BitwiseAndScalarOperator,
+    BitwiseOrScalarOperator,
+    BitwiseXorScalarOperator,
+    LeftShiftScalarOperator,
+    RightShiftScalarOperator,
+)

temporian/core/operators/scalar/bitwise_scalar.py

@@ -0,0 +1,141 @@
+# Copyright 2021 Google LLC.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Event/scalar arithmetic operators classes and public API definitions."""
+
+from typing import Union
+
+from temporian.core import operator_lib
+from temporian.core.compilation import compile
+from temporian.core.data.dtype import DType
+from temporian.core.data.node import EventSetNode
+from temporian.core.operators.scalar.base import (
+    BaseScalarOperator,
+)
+from temporian.core.typing import EventSetOrNode
+
+class BaseBitwiseScalarOperator(BaseScalarOperator):
+    DEF_KEY = ""
+    def __init__(
+        self,
+        input: EventSetNode,
+        value: int,
+        is_value_first: bool = False,
+    ):
+        # Check that all features are Integer
+        for feature in input.schema.features:
+            if feature.dtype != DType.INT32 and feature.dtype != DType.INT64:
+                raise ValueError(
+                    "Bitwise operators only support INT32 or INT64 types, but "
+                    f"feature {feature.name} has dtype {feature.dtype}"
+                )
+        # Check that the value is an integer
+        if DType.from_python_value(value) not in [DType.INT32, DType.INT64]:
+            raise ValueError(
+                "Bitwise operators only support int values, but value"
+                f" has type {DType.from_python_value(value)}"
+            )
+        super().__init__(input, value, is_value_first)
+
+
+class BitwiseAndScalarOperator(BaseBitwiseScalarOperator):
+    DEF_KEY = "BITWISE_AND_SCALAR"
+
+
+class BitwiseOrScalarOperator(BaseBitwiseScalarOperator):
+    DEF_KEY = "BITWISE_OR_SCALAR"
+
+
+class BitwiseXorScalarOperator(BaseBitwiseScalarOperator):
+    DEF_KEY = "BITWISE_XOR_SCALAR"
+
+
+class LeftShiftScalarOperator(BaseBitwiseScalarOperator):
+    DEF_KEY = "LEFT_SHIFT_SCALAR"
+
+
+class RightShiftScalarOperator(BaseBitwiseScalarOperator):
+    DEF_KEY = "RIGHT_SHIFT_SCALAR"
+
+
+@compile
+def bitwise_and_scalar(
+    input: EventSetOrNode,
+    value: int,
+) -> EventSetOrNode:
+    assert isinstance(input, EventSetNode)
+
+    return BitwiseAndScalarOperator(
+        input=input,
+        value=value,
+    ).outputs["output"]
+
+
+@compile
+def bitwise_or_scalar(
+    input: EventSetOrNode,
+    value: int,
+) -> EventSetOrNode:
+    assert isinstance(input, EventSetNode)
+
+    return BitwiseOrScalarOperator(
+        input=input,
+        value=value,
+    ).outputs["output"]
+
+
+@compile
+def bitwise_xor_scalar(
+    input: EventSetOrNode,
+    value: int,
+) -> EventSetOrNode:
+    assert isinstance(input, EventSetNode)
+
+    return BitwiseXorScalarOperator(
+        input=input,
+        value=value,
+    ).outputs["output"]
+
+
+@compile
+def left_shift_scalar(
+    input: EventSetOrNode,
+    value: int,
+) -> EventSetOrNode:
+    assert isinstance(input, EventSetNode)
+
+    return LeftShiftScalarOperator(
+        input=input,
+        value=value,
+    ).outputs["output"]
+
+
+@compile
+def right_shift_scalar(
+    input: EventSetOrNode,
+    value: int,
+) -> EventSetOrNode:
+    assert isinstance(input, EventSetNode)
+
+    return RightShiftScalarOperator(
+        input=input,
+        value=value,
+    ).outputs["output"]
+
+
+operator_lib.register_operator(BitwiseAndScalarOperator)
+operator_lib.register_operator(BitwiseOrScalarOperator)
+operator_lib.register_operator(BitwiseXorScalarOperator)
+operator_lib.register_operator(LeftShiftScalarOperator)
+operator_lib.register_operator(RightShiftScalarOperator)