Skip to content

Commit b3286d3

Browse files
wesmwesm
committed
Port R C++ code, though there is a failing test
1 parent 6f71f77 commit b3286d3

File tree

1 file changed

+16
-59
lines changed

1 file changed

+16
-59
lines changed

‎r/src/compute.cpp‎

Lines changed: 16 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -34,21 +34,16 @@ std::shared_ptr<arrow::Array> Array__cast(
3434
const std::shared_ptr<arrow::Array>& array,
3535
const std::shared_ptr<arrow::DataType>& target_type,
3636
const std::shared_ptr<arrow::compute::CastOptions>& options) {
37-
std::shared_ptr<arrow::Array> out;
38-
arrow::compute::FunctionContext context;
39-
StopIfNotOk(arrow::compute::Cast(&context, *array, target_type, *options, &out));
40-
return out;
37+
return ValueOrStop(arrow::compute::Cast(*array, target_type, *options));
4138
}
4239

4340
// [[arrow::export]]
4441
std::shared_ptr<arrow::ChunkedArray> ChunkedArray__cast(
4542
const std::shared_ptr<arrow::ChunkedArray>& chunked_array,
4643
const std::shared_ptr<arrow::DataType>& target_type,
4744
const std::shared_ptr<arrow::compute::CastOptions>& options) {
48-
arrow::compute::Datum value(chunked_array);
49-
arrow::compute::Datum out;
50-
arrow::compute::FunctionContext context;
51-
StopIfNotOk(arrow::compute::Cast(&context, value, target_type, *options, &out));
45+
arrow::Datum value(chunked_array);
46+
arrow::Datum out = ValueOrStop(arrow::compute::Cast(value, target_type, *options));
5247
return out.chunked_array();
5348
}
5449

@@ -85,155 +80,119 @@ std::shared_ptr<arrow::Table> Table__cast(
8580
// [[arrow::export]]
8681
std::shared_ptr<arrow::Array> Array__Take(const std::shared_ptr<arrow::Array>& values,
8782
const std::shared_ptr<arrow::Array>& indices) {
88-
std::shared_ptr<arrow::Array> out;
89-
arrow::compute::FunctionContext context;
9083
arrow::compute::TakeOptions options;
91-
StopIfNotOk(arrow::compute::Take(&context, *values, *indices, options, &out));
92-
return out;
84+
return ValueOrStop(arrow::compute::Take(*values, *indices, options));
9385
}
9486

9587
// [[arrow::export]]
9688
std::shared_ptr<arrow::ChunkedArray> Array__TakeChunked(
9789
const std::shared_ptr<arrow::Array>& values,
9890
const std::shared_ptr<arrow::ChunkedArray>& indices) {
99-
std::shared_ptr<arrow::ChunkedArray> out;
100-
arrow::compute::FunctionContext context;
10191
arrow::compute::TakeOptions options;
102-
103-
StopIfNotOk(arrow::compute::Take(&context, *values, *indices, options, &out));
104-
return out;
92+
return ValueOrStop(arrow::compute::Take(*values, *indices, options));
10593
}
10694

10795
// [[arrow::export]]
10896
std::shared_ptr<arrow::RecordBatch> RecordBatch__Take(
10997
const std::shared_ptr<arrow::RecordBatch>& batch,
11098
const std::shared_ptr<arrow::Array>& indices) {
111-
std::shared_ptr<arrow::RecordBatch> out;
112-
arrow::compute::FunctionContext context;
11399
arrow::compute::TakeOptions options;
114-
StopIfNotOk(arrow::compute::Take(&context, *batch, *indices, options, &out));
115-
return out;
100+
return ValueOrStop(arrow::compute::Take(*batch, *indices, options));
116101
}
117102

118103
// [[arrow::export]]
119104
std::shared_ptr<arrow::ChunkedArray> ChunkedArray__Take(
120105
const std::shared_ptr<arrow::ChunkedArray>& values,
121106
const std::shared_ptr<arrow::Array>& indices) {
122-
std::shared_ptr<arrow::ChunkedArray> out;
123-
arrow::compute::FunctionContext context;
124107
arrow::compute::TakeOptions options;
125-
126-
StopIfNotOk(arrow::compute::Take(&context, *values, *indices, options, &out));
127-
return out;
108+
return ValueOrStop(arrow::compute::Take(*values, *indices, options));
128109
}
129110

130111
// [[arrow::export]]
131112
std::shared_ptr<arrow::ChunkedArray> ChunkedArray__TakeChunked(
132113
const std::shared_ptr<arrow::ChunkedArray>& values,
133114
const std::shared_ptr<arrow::ChunkedArray>& indices) {
134-
std::shared_ptr<arrow::ChunkedArray> out;
135-
arrow::compute::FunctionContext context;
136115
arrow::compute::TakeOptions options;
137-
138-
StopIfNotOk(arrow::compute::Take(&context, *values, *indices, options, &out));
139-
return out;
116+
return ValueOrStop(arrow::compute::Take(*values, *indices, options));
140117
}
141118

142119
// [[arrow::export]]
143120
std::shared_ptr<arrow::Table> Table__Take(const std::shared_ptr<arrow::Table>& table,
144121
const std::shared_ptr<arrow::Array>& indices) {
145-
std::shared_ptr<arrow::Table> out;
146-
arrow::compute::FunctionContext context;
147122
arrow::compute::TakeOptions options;
148-
149-
StopIfNotOk(arrow::compute::Take(&context, *table, *indices, options, &out));
150-
return out;
123+
return ValueOrStop(arrow::compute::Take(*table, *indices, options));
151124
}
152125

153126
// [[arrow::export]]
154127
std::shared_ptr<arrow::Table> Table__TakeChunked(
155128
const std::shared_ptr<arrow::Table>& table,
156129
const std::shared_ptr<arrow::ChunkedArray>& indices) {
157-
std::shared_ptr<arrow::Table> out;
158-
arrow::compute::FunctionContext context;
159130
arrow::compute::TakeOptions options;
160-
161-
StopIfNotOk(arrow::compute::Take(&context, *table, *indices, options, &out));
162-
return out;
131+
return ValueOrStop(arrow::compute::Take(*table, *indices, options));
163132
}
164133

165134
// [[arrow::export]]
166135
std::shared_ptr<arrow::Array> Array__Filter(const std::shared_ptr<arrow::Array>& values,
167136
const std::shared_ptr<arrow::Array>& filter,
168137
bool keep_na) {
169-
arrow::compute::FunctionContext context;
170-
arrow::compute::Datum out;
171138
// Use the EMIT_NULL filter option to match R's behavior in [
172139
arrow::compute::FilterOptions options;
173140
if (keep_na) {
174141
options.null_selection_behavior = arrow::compute::FilterOptions::EMIT_NULL;
175142
}
176-
StopIfNotOk(arrow::compute::Filter(&context, values, filter, {}, &out));
143+
arrow::Datum out = ValueOrStop(arrow::compute::Filter(values, filter, options));
177144
return out.make_array();
178145
}
179146

180147
// [[arrow::export]]
181148
std::shared_ptr<arrow::RecordBatch> RecordBatch__Filter(
182149
const std::shared_ptr<arrow::RecordBatch>& batch,
183150
const std::shared_ptr<arrow::Array>& filter, bool keep_na) {
184-
arrow::compute::FunctionContext context;
185-
arrow::compute::Datum out;
186151
// Use the EMIT_NULL filter option to match R's behavior in [
187152
arrow::compute::FilterOptions options;
188153
if (keep_na) {
189154
options.null_selection_behavior = arrow::compute::FilterOptions::EMIT_NULL;
190155
}
191-
StopIfNotOk(arrow::compute::Filter(&context, batch, filter, options, &out));
156+
arrow::Datum out = ValueOrStop(arrow::compute::Filter(batch, filter, options));
192157
return out.record_batch();
193158
}
194159

195160
// [[arrow::export]]
196161
std::shared_ptr<arrow::ChunkedArray> ChunkedArray__Filter(
197162
const std::shared_ptr<arrow::ChunkedArray>& values,
198163
const std::shared_ptr<arrow::Array>& filter, bool keep_na) {
199-
arrow::compute::FunctionContext context;
200-
arrow::compute::Datum out;
201164
// Use the EMIT_NULL filter option to match R's behavior in [
202165
arrow::compute::FilterOptions options;
203166
if (keep_na) {
204167
options.null_selection_behavior = arrow::compute::FilterOptions::EMIT_NULL;
205168
}
206-
StopIfNotOk(arrow::compute::Filter(&context, values, filter, options, &out));
169+
arrow::Datum out = ValueOrStop(arrow::compute::Filter(values, filter, options));
207170
return out.chunked_array();
208171
}
209172

210173
// [[arrow::export]]
211174
std::shared_ptr<arrow::ChunkedArray> ChunkedArray__FilterChunked(
212175
const std::shared_ptr<arrow::ChunkedArray>& values,
213176
const std::shared_ptr<arrow::ChunkedArray>& filter, bool keep_na) {
214-
arrow::compute::FunctionContext context;
215-
arrow::compute::Datum out;
216177
// Use the EMIT_NULL filter option to match R's behavior in [
217178
arrow::compute::FilterOptions options;
218179
if (keep_na) {
219180
options.null_selection_behavior = arrow::compute::FilterOptions::EMIT_NULL;
220181
}
221-
StopIfNotOk(arrow::compute::Filter(&context, values, filter, options, &out));
182+
arrow::Datum out = ValueOrStop(arrow::compute::Filter(values, filter, options));
222183
return out.chunked_array();
223184
}
224185

225186
// [[arrow::export]]
226187
std::shared_ptr<arrow::Table> Table__Filter(const std::shared_ptr<arrow::Table>& table,
227188
const std::shared_ptr<arrow::Array>& filter,
228189
bool keep_na) {
229-
arrow::compute::FunctionContext context;
230-
arrow::compute::Datum out;
231190
// Use the EMIT_NULL filter option to match R's behavior in [
232191
arrow::compute::FilterOptions options;
233192
if (keep_na) {
234193
options.null_selection_behavior = arrow::compute::FilterOptions::EMIT_NULL;
235194
}
236-
StopIfNotOk(arrow::compute::Filter(&context, table, filter, options, &out));
195+
arrow::Datum out = ValueOrStop(arrow::compute::Filter(table, filter, options));
237196
std::shared_ptr<arrow::Table> tab = out.table();
238197
if (tab->num_rows() == 0) {
239198
// Slight hack: if there are no rows in the result, instead do a 0-length
@@ -248,14 +207,12 @@ std::shared_ptr<arrow::Table> Table__Filter(const std::shared_ptr<arrow::Table>&
248207
std::shared_ptr<arrow::Table> Table__FilterChunked(
249208
const std::shared_ptr<arrow::Table>& table,
250209
const std::shared_ptr<arrow::ChunkedArray>& filter, bool keep_na) {
251-
arrow::compute::FunctionContext context;
252-
arrow::compute::Datum out;
253210
// Use the EMIT_NULL filter option to match R's behavior in [
254211
arrow::compute::FilterOptions options;
255212
if (keep_na) {
256213
options.null_selection_behavior = arrow::compute::FilterOptions::EMIT_NULL;
257214
}
258-
StopIfNotOk(arrow::compute::Filter(&context, table, filter, options, &out));
215+
arrow::Datum out = ValueOrStop(arrow::compute::Filter(table, filter, options));
259216
std::shared_ptr<arrow::Table> tab = out.table();
260217
if (tab->num_rows() == 0) {
261218
// Slight hack: if there are no rows in the result, instead do a 0-length

0 commit comments

Comments
 (0)