@@ -516,6 +516,9 @@ def _compile_orderby(self, node: nodes.OrderByNode) -> algebra_pb2.Rel:
516516 "pow" : 69 ,
517517 "cov" : 70 ,
518518 "corr" : 71 ,
519+ "bitwise_and" : 72 ,
520+ "bitwise_or" : 73 ,
521+ "bitwise_xor" : 74 ,
519522 }
520523
521524 _OP_TO_EXTENSION = {
@@ -923,16 +926,35 @@ def _compile_fillna_op(
923926 @_compile_op .register (comparison_ops .GtOp )
924927 @_compile_op .register (comparison_ops .LeOp )
925928 @_compile_op .register (comparison_ops .GeOp )
926- @_compile_op .register (bool_ops .AndOp )
927- @_compile_op .register (bool_ops .OrOp )
928- @_compile_op .register (bool_ops .XorOp )
929929 def _compile_basic_binops (
930930 self , op : Any , inputs : Sequence [ex .Expression ], child : nodes .BigFrameNode
931931 ) -> algebra_pb2 .Expression :
932932 op_class = type (op )
933933 ext_name = self ._OP_TO_EXTENSION [op_class ]
934934 return self ._compile_basic_binop (ext_name , inputs , child )
935935
936+ @_compile_op .register (bool_ops .AndOp )
937+ @_compile_op .register (bool_ops .OrOp )
938+ @_compile_op .register (bool_ops .XorOp )
939+ def _compile_logical_binops (
940+ self , op : Any , inputs : Sequence [ex .Expression ], child : nodes .BigFrameNode
941+ ) -> algebra_pb2 .Expression :
942+ import bigframes .dtypes as dtypes
943+ input_dtype = self ._get_expression_dtype (inputs [0 ], child )
944+ if input_dtype == dtypes .INT_DTYPE :
945+ if isinstance (op , bool_ops .AndOp ):
946+ ext_name = "bitwise_and"
947+ elif isinstance (op , bool_ops .OrOp ):
948+ ext_name = "bitwise_or"
949+ elif isinstance (op , bool_ops .XorOp ):
950+ ext_name = "bitwise_xor"
951+ else :
952+ raise NotImplementedError (f"Unsupported binary bitwise op: { type (op )} " )
953+ else :
954+ op_class = type (op )
955+ ext_name = self ._OP_TO_EXTENSION [op_class ]
956+ return self ._compile_basic_binop (ext_name , inputs , child )
957+
936958 def _compile_basic_binop (
937959 self , ext_name : str , inputs : Sequence [ex .Expression ], child : nodes .BigFrameNode
938960 ) -> algebra_pb2 .Expression :
0 commit comments