forked from rapidsai/cudf
-
Notifications
You must be signed in to change notification settings - Fork 0
/
expressions.pyx
120 lines (103 loc) · 4.62 KB
/
expressions.pyx
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
# Copyright (c) 2022-2023, NVIDIA CORPORATION.
from enum import Enum
from cython.operator cimport dereference
from libc.stdint cimport int64_t
from libcpp.memory cimport unique_ptr
from libcpp.utility cimport move
from cudf._lib.cpp cimport expressions as libcudf_exp
from cudf._lib.cpp.libcpp.memory cimport make_unique
from cudf._lib.cpp.types cimport size_type
# Necessary for proper casting, see below.
ctypedef int32_t underlying_type_ast_operator
# Aliases for simplicity
ctypedef unique_ptr[libcudf_exp.expression] expression_ptr
class ASTOperator(Enum):
ADD = libcudf_exp.ast_operator.ADD
SUB = libcudf_exp.ast_operator.SUB
MUL = libcudf_exp.ast_operator.MUL
DIV = libcudf_exp.ast_operator.DIV
TRUE_DIV = libcudf_exp.ast_operator.TRUE_DIV
FLOOR_DIV = libcudf_exp.ast_operator.FLOOR_DIV
MOD = libcudf_exp.ast_operator.MOD
PYMOD = libcudf_exp.ast_operator.PYMOD
POW = libcudf_exp.ast_operator.POW
EQUAL = libcudf_exp.ast_operator.EQUAL
NULL_EQUAL = libcudf_exp.ast_operator.NULL_EQUAL
NOT_EQUAL = libcudf_exp.ast_operator.NOT_EQUAL
LESS = libcudf_exp.ast_operator.LESS
GREATER = libcudf_exp.ast_operator.GREATER
LESS_EQUAL = libcudf_exp.ast_operator.LESS_EQUAL
GREATER_EQUAL = libcudf_exp.ast_operator.GREATER_EQUAL
BITWISE_AND = libcudf_exp.ast_operator.BITWISE_AND
BITWISE_OR = libcudf_exp.ast_operator.BITWISE_OR
BITWISE_XOR = libcudf_exp.ast_operator.BITWISE_XOR
LOGICAL_AND = libcudf_exp.ast_operator.LOGICAL_AND
NULL_LOGICAL_AND = libcudf_exp.ast_operator.NULL_LOGICAL_AND
LOGICAL_OR = libcudf_exp.ast_operator.LOGICAL_OR
NULL_LOGICAL_OR = libcudf_exp.ast_operator.NULL_LOGICAL_OR
# Unary operators
IDENTITY = libcudf_exp.ast_operator.IDENTITY
IS_NULL = libcudf_exp.ast_operator.IS_NULL
SIN = libcudf_exp.ast_operator.SIN
COS = libcudf_exp.ast_operator.COS
TAN = libcudf_exp.ast_operator.TAN
ARCSIN = libcudf_exp.ast_operator.ARCSIN
ARCCOS = libcudf_exp.ast_operator.ARCCOS
ARCTAN = libcudf_exp.ast_operator.ARCTAN
SINH = libcudf_exp.ast_operator.SINH
COSH = libcudf_exp.ast_operator.COSH
TANH = libcudf_exp.ast_operator.TANH
ARCSINH = libcudf_exp.ast_operator.ARCSINH
ARCCOSH = libcudf_exp.ast_operator.ARCCOSH
ARCTANH = libcudf_exp.ast_operator.ARCTANH
EXP = libcudf_exp.ast_operator.EXP
LOG = libcudf_exp.ast_operator.LOG
SQRT = libcudf_exp.ast_operator.SQRT
CBRT = libcudf_exp.ast_operator.CBRT
CEIL = libcudf_exp.ast_operator.CEIL
FLOOR = libcudf_exp.ast_operator.FLOOR
ABS = libcudf_exp.ast_operator.ABS
RINT = libcudf_exp.ast_operator.RINT
BIT_INVERT = libcudf_exp.ast_operator.BIT_INVERT
NOT = libcudf_exp.ast_operator.NOT
class TableReference(Enum):
LEFT = libcudf_exp.table_reference.LEFT
RIGHT = libcudf_exp.table_reference.RIGHT
# Note that this function only currently supports numeric literals. libcudf
# expressions don't really support other types yet though, so this isn't
# restrictive at the moment.
cdef class Literal(Expression):
def __cinit__(self, value):
if isinstance(value, int):
self.c_scalar.reset(new numeric_scalar[int64_t](value, True))
self.c_obj = <expression_ptr> move(make_unique[libcudf_exp.literal](
<numeric_scalar[int64_t] &>dereference(self.c_scalar)
))
elif isinstance(value, float):
self.c_scalar.reset(new numeric_scalar[double](value, True))
self.c_obj = <expression_ptr> move(make_unique[libcudf_exp.literal](
<numeric_scalar[double] &>dereference(self.c_scalar)
))
elif isinstance(value, str):
self.c_scalar.reset(new string_scalar(value.encode(), True))
self.c_obj = <expression_ptr> move(make_unique[libcudf_exp.literal](
<string_scalar &>dereference(self.c_scalar)
))
cdef class ColumnReference(Expression):
def __cinit__(self, size_type index):
self.c_obj = <expression_ptr>move(make_unique[libcudf_exp.column_reference](
index
))
cdef class Operation(Expression):
def __cinit__(self, op, Expression left, Expression right=None):
cdef libcudf_exp.ast_operator op_value = <libcudf_exp.ast_operator>(
<underlying_type_ast_operator> op.value
)
if right is None:
self.c_obj = <expression_ptr> move(make_unique[libcudf_exp.operation](
op_value, dereference(left.c_obj)
))
else:
self.c_obj = <expression_ptr> move(make_unique[libcudf_exp.operation](
op_value, dereference(left.c_obj), dereference(right.c_obj)
))