1# Licensed to the Apache Software Foundation (ASF) under one
2# or more contributor license agreements.  See the NOTICE file
3# distributed with this work for additional information
4# regarding copyright ownership.  The ASF licenses this file
5# to you under the Apache License, Version 2.0 (the
6# "License"); you may not use this file except in compliance
7# with the License.  You may obtain a copy of the License at
8#
9#   http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing,
12# software distributed under the License is distributed on an
13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14# KIND, either express or implied.  See the License for the
15# specific language governing permissions and limitations
16# under the License.
17
18# distutils: language = c++
19
20from libcpp.string cimport string as c_string
21from libcpp.unordered_set cimport unordered_set as c_unordered_set
22from libc.stdint cimport int64_t, int32_t, uint8_t, uintptr_t
23
24from pyarrow.includes.common cimport *
25from pyarrow.includes.libarrow cimport *
26
27cdef extern from "gandiva/node.h" namespace "gandiva" nogil:
28
29    cdef cppclass CNode" gandiva::Node":
30        c_string ToString()
31        shared_ptr[CDataType] return_type()
32
33    cdef cppclass CExpression" gandiva::Expression":
34        c_string ToString()
35        shared_ptr[CNode] root()
36        shared_ptr[CField] result()
37
38    ctypedef vector[shared_ptr[CNode]] CNodeVector" gandiva::NodeVector"
39
40    ctypedef vector[shared_ptr[CExpression]] \
41        CExpressionVector" gandiva::ExpressionVector"
42
43cdef extern from "gandiva/selection_vector.h" namespace "gandiva" nogil:
44
45    cdef cppclass CSelectionVector" gandiva::SelectionVector":
46
47        shared_ptr[CArray] ToArray()
48
49    enum CSelectionVector_Mode" gandiva::SelectionVector::Mode":
50        CSelectionVector_Mode_NONE" gandiva::SelectionVector::Mode::MODE_NONE"
51        CSelectionVector_Mode_UINT16" \
52                gandiva::SelectionVector::Mode::MODE_UINT16"
53        CSelectionVector_Mode_UINT32" \
54                gandiva::SelectionVector::Mode::MODE_UINT32"
55        CSelectionVector_Mode_UINT64" \
56                gandiva::SelectionVector::Mode::MODE_UINT64"
57
58    cdef CStatus SelectionVector_MakeInt16\
59        "gandiva::SelectionVector::MakeInt16"(
60            int64_t max_slots, CMemoryPool* pool,
61            shared_ptr[CSelectionVector]* selection_vector)
62
63    cdef CStatus SelectionVector_MakeInt32\
64        "gandiva::SelectionVector::MakeInt32"(
65            int64_t max_slots, CMemoryPool* pool,
66            shared_ptr[CSelectionVector]* selection_vector)
67
68    cdef CStatus SelectionVector_MakeInt64\
69        "gandiva::SelectionVector::MakeInt64"(
70            int64_t max_slots, CMemoryPool* pool,
71            shared_ptr[CSelectionVector]* selection_vector)
72
73cdef inline CSelectionVector_Mode _ensure_selection_mode(str name) except *:
74    uppercase = name.upper()
75    if uppercase == 'NONE':
76        return CSelectionVector_Mode_NONE
77    elif uppercase == 'UINT16':
78        return CSelectionVector_Mode_UINT16
79    elif uppercase == 'UINT32':
80        return CSelectionVector_Mode_UINT32
81    elif uppercase == 'UINT64':
82        return CSelectionVector_Mode_UINT64
83    else:
84        raise ValueError('Invalid value for Selection Mode: {!r}'.format(name))
85
86cdef inline str _selection_mode_name(CSelectionVector_Mode ctype):
87    if ctype == CSelectionVector_Mode_NONE:
88        return 'NONE'
89    elif ctype == CSelectionVector_Mode_UINT16:
90        return 'UINT16'
91    elif ctype == CSelectionVector_Mode_UINT32:
92        return 'UINT32'
93    elif ctype == CSelectionVector_Mode_UINT64:
94        return 'UINT64'
95    else:
96        raise RuntimeError('Unexpected CSelectionVector_Mode value')
97
98cdef extern from "gandiva/condition.h" namespace "gandiva" nogil:
99
100    cdef cppclass CCondition" gandiva::Condition":
101        c_string ToString()
102        shared_ptr[CNode] root()
103        shared_ptr[CField] result()
104
105cdef extern from "gandiva/arrow.h" namespace "gandiva" nogil:
106
107    ctypedef vector[shared_ptr[CArray]] CArrayVector" gandiva::ArrayVector"
108
109
110cdef extern from "gandiva/tree_expr_builder.h" namespace "gandiva" nogil:
111
112    cdef shared_ptr[CNode] TreeExprBuilder_MakeBoolLiteral \
113        "gandiva::TreeExprBuilder::MakeLiteral"(c_bool value)
114
115    cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt8Literal \
116        "gandiva::TreeExprBuilder::MakeLiteral"(uint8_t value)
117
118    cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt16Literal \
119        "gandiva::TreeExprBuilder::MakeLiteral"(uint16_t value)
120
121    cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt32Literal \
122        "gandiva::TreeExprBuilder::MakeLiteral"(uint32_t value)
123
124    cdef shared_ptr[CNode] TreeExprBuilder_MakeUInt64Literal \
125        "gandiva::TreeExprBuilder::MakeLiteral"(uint64_t value)
126
127    cdef shared_ptr[CNode] TreeExprBuilder_MakeInt8Literal \
128        "gandiva::TreeExprBuilder::MakeLiteral"(int8_t value)
129
130    cdef shared_ptr[CNode] TreeExprBuilder_MakeInt16Literal \
131        "gandiva::TreeExprBuilder::MakeLiteral"(int16_t value)
132
133    cdef shared_ptr[CNode] TreeExprBuilder_MakeInt32Literal \
134        "gandiva::TreeExprBuilder::MakeLiteral"(int32_t value)
135
136    cdef shared_ptr[CNode] TreeExprBuilder_MakeInt64Literal \
137        "gandiva::TreeExprBuilder::MakeLiteral"(int64_t value)
138
139    cdef shared_ptr[CNode] TreeExprBuilder_MakeFloatLiteral \
140        "gandiva::TreeExprBuilder::MakeLiteral"(float value)
141
142    cdef shared_ptr[CNode] TreeExprBuilder_MakeDoubleLiteral \
143        "gandiva::TreeExprBuilder::MakeLiteral"(double value)
144
145    cdef shared_ptr[CNode] TreeExprBuilder_MakeStringLiteral \
146        "gandiva::TreeExprBuilder::MakeStringLiteral"(const c_string& value)
147
148    cdef shared_ptr[CNode] TreeExprBuilder_MakeBinaryLiteral \
149        "gandiva::TreeExprBuilder::MakeBinaryLiteral"(const c_string& value)
150
151    cdef shared_ptr[CExpression] TreeExprBuilder_MakeExpression\
152        "gandiva::TreeExprBuilder::MakeExpression"(
153            shared_ptr[CNode] root_node, shared_ptr[CField] result_field)
154
155    cdef shared_ptr[CNode] TreeExprBuilder_MakeFunction \
156        "gandiva::TreeExprBuilder::MakeFunction"(
157            const c_string& name, const CNodeVector& children,
158            shared_ptr[CDataType] return_type)
159
160    cdef shared_ptr[CNode] TreeExprBuilder_MakeField \
161        "gandiva::TreeExprBuilder::MakeField"(shared_ptr[CField] field)
162
163    cdef shared_ptr[CNode] TreeExprBuilder_MakeIf \
164        "gandiva::TreeExprBuilder::MakeIf"(
165            shared_ptr[CNode] condition, shared_ptr[CNode] this_node,
166            shared_ptr[CNode] else_node, shared_ptr[CDataType] return_type)
167
168    cdef shared_ptr[CNode] TreeExprBuilder_MakeAnd \
169        "gandiva::TreeExprBuilder::MakeAnd"(const CNodeVector& children)
170
171    cdef shared_ptr[CNode] TreeExprBuilder_MakeOr \
172        "gandiva::TreeExprBuilder::MakeOr"(const CNodeVector& children)
173
174    cdef shared_ptr[CCondition] TreeExprBuilder_MakeCondition \
175        "gandiva::TreeExprBuilder::MakeCondition"(
176            shared_ptr[CNode] condition)
177
178    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionInt32 \
179        "gandiva::TreeExprBuilder::MakeInExpressionInt32"(
180            shared_ptr[CNode] node, const c_unordered_set[int32_t]& values)
181
182    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionInt64 \
183        "gandiva::TreeExprBuilder::MakeInExpressionInt64"(
184            shared_ptr[CNode] node, const c_unordered_set[int64_t]& values)
185
186    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionTime32 \
187        "gandiva::TreeExprBuilder::MakeInExpressionTime32"(
188            shared_ptr[CNode] node, const c_unordered_set[int32_t]& values)
189
190    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionTime64 \
191        "gandiva::TreeExprBuilder::MakeInExpressionTime64"(
192            shared_ptr[CNode] node, const c_unordered_set[int64_t]& values)
193
194    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionDate32 \
195        "gandiva::TreeExprBuilder::MakeInExpressionDate32"(
196            shared_ptr[CNode] node, const c_unordered_set[int32_t]& values)
197
198    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionDate64 \
199        "gandiva::TreeExprBuilder::MakeInExpressionDate64"(
200            shared_ptr[CNode] node, const c_unordered_set[int64_t]& values)
201
202    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionTimeStamp \
203        "gandiva::TreeExprBuilder::MakeInExpressionTimeStamp"(
204            shared_ptr[CNode] node, const c_unordered_set[int64_t]& values)
205
206    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionString \
207        "gandiva::TreeExprBuilder::MakeInExpressionString"(
208            shared_ptr[CNode] node, const c_unordered_set[c_string]& values)
209
210    cdef shared_ptr[CNode] TreeExprBuilder_MakeInExpressionBinary \
211        "gandiva::TreeExprBuilder::MakeInExpressionBinary"(
212            shared_ptr[CNode] node, const c_unordered_set[c_string]& values)
213
214cdef extern from "gandiva/projector.h" namespace "gandiva" nogil:
215
216    cdef cppclass CProjector" gandiva::Projector":
217
218        CStatus Evaluate(
219            const CRecordBatch& batch, CMemoryPool* pool,
220            const CArrayVector* output)
221
222        CStatus Evaluate(
223            const CRecordBatch& batch,
224            const CSelectionVector* selection,
225            CMemoryPool* pool,
226            const CArrayVector* output)
227
228        c_string DumpIR()
229
230    cdef CStatus Projector_Make \
231        "gandiva::Projector::Make"(
232            shared_ptr[CSchema] schema, const CExpressionVector& children,
233            shared_ptr[CProjector]* projector)
234
235    cdef CStatus Projector_Make \
236        "gandiva::Projector::Make"(
237            shared_ptr[CSchema] schema, const CExpressionVector& children,
238            CSelectionVector_Mode mode,
239            shared_ptr[CConfiguration] configuration,
240            shared_ptr[CProjector]* projector)
241
242cdef extern from "gandiva/filter.h" namespace "gandiva" nogil:
243
244    cdef cppclass CFilter" gandiva::Filter":
245
246        CStatus Evaluate(
247            const CRecordBatch& batch,
248            shared_ptr[CSelectionVector] out_selection)
249
250        c_string DumpIR()
251
252    cdef CStatus Filter_Make \
253        "gandiva::Filter::Make"(
254            shared_ptr[CSchema] schema, shared_ptr[CCondition] condition,
255            shared_ptr[CFilter]* filter)
256
257cdef extern from "gandiva/function_signature.h" namespace "gandiva" nogil:
258
259    cdef cppclass CFunctionSignature" gandiva::FunctionSignature":
260
261        CFunctionSignature(const c_string& base_name,
262                           vector[shared_ptr[CDataType]] param_types,
263                           shared_ptr[CDataType] ret_type)
264
265        shared_ptr[CDataType] ret_type() const
266
267        const c_string& base_name() const
268
269        vector[shared_ptr[CDataType]] param_types() const
270
271        c_string ToString() const
272
273cdef extern from "gandiva/expression_registry.h" namespace "gandiva" nogil:
274
275    cdef vector[shared_ptr[CFunctionSignature]] \
276        GetRegisteredFunctionSignatures()
277
278cdef extern from "gandiva/configuration.h" namespace "gandiva" nogil:
279
280    cdef cppclass CConfiguration" gandiva::Configuration":
281        pass
282
283    cdef cppclass CConfigurationBuilder \
284            " gandiva::ConfigurationBuilder":
285        @staticmethod
286        shared_ptr[CConfiguration] DefaultConfiguration()
287