lidavidm commented on a change in pull request #12162:
URL: https://github.com/apache/arrow/pull/12162#discussion_r792954414
##########
File path: cpp/src/arrow/compute/kernels/scalar_nested.cc
##########
@@ -428,40 +428,82 @@ const FunctionDoc make_struct_doc{"Wrap Arrays into a
StructArray",
"specified through MakeStructOptions."),
{"*args"},
"MakeStructOptions"};
-
+template <typename KeyType>
struct MapArrayLookupFunctor {
- static Result<int64_t> FindOneMapValueIndex(const Array& keys, const Scalar&
query_key,
+ static Result<int64_t> FindOneMapValueIndex(const Array& keys,
+ const Scalar& query_key_scalar,
const int64_t start, const
int64_t end,
const bool from_back = false) {
- if (!from_back) {
- for (int64_t idx = start; idx < end; ++idx) {
- ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Scalar> key,
keys.GetScalar(idx));
-
- if (key->Equals(query_key)) return idx;
- }
- } else {
- for (int64_t idx = end - 1; idx >= start; --idx) {
- ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Scalar> key,
keys.GetScalar(idx));
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ int64_t match_idx = -1;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
+ if (!from_back) {
+ match_idx = index;
+ return Status::Cancelled("Found first matching key");
+ } else {
+ match_idx = index;
+ }
+ }
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ },
+ [&]() -> Status {
+ if (index < end) {
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ }));
- if (key->Equals(query_key)) return idx;
- }
- }
- return -1;
+ return match_idx;
}
static Result<std::unique_ptr<ArrayBuilder>> GetBuiltArray(
Review comment:
This could perhaps take in `ArrayBuilder* builder` as a parameter
instead of returning it. That way, for the Array case below, it could directly
append to the child items builder of the list builder, avoiding an extra copy
of the data/allocation of a scalar. And the Scalar case could just use it as
normal.
##########
File path: cpp/src/arrow/compute/kernels/scalar_nested.cc
##########
@@ -428,6 +428,290 @@ const FunctionDoc make_struct_doc{"Wrap Arrays into a
StructArray",
"specified through MakeStructOptions."),
{"*args"},
"MakeStructOptions"};
+template <typename KeyType>
+struct MapArrayLookupFunctor {
+ static Result<int64_t> FindOneMapValueIndex(const Array& keys,
+ const Scalar& query_key_scalar,
+ const int64_t start, const
int64_t end,
+ const bool from_back = false) {
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ int64_t match_idx = -1;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
Review comment:
The latter will allocate a new Array which will do some shared_ptr
copies. That's not too bad, though, so if it makes the code clearer, let's go
with that. (We can adjust if it turns out to be significant overhead but I'm
not too sure it will. If so, we probably need to add explicit offset/length
support to VisitArrayValuesInline)
##########
File path: cpp/src/arrow/compute/kernels/scalar_nested.cc
##########
@@ -428,6 +428,290 @@ const FunctionDoc make_struct_doc{"Wrap Arrays into a
StructArray",
"specified through MakeStructOptions."),
{"*args"},
"MakeStructOptions"};
+template <typename KeyType>
+struct MapArrayLookupFunctor {
+ static Result<int64_t> FindOneMapValueIndex(const Array& keys,
+ const Scalar& query_key_scalar,
+ const int64_t start, const
int64_t end,
+ const bool from_back = false) {
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ int64_t match_idx = -1;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
Review comment:
Hmm, I also see that there's some duplicate logic below for the array
case. Instead of having this return a single index, you could organize it as
follows:
```
template <typename FoundItem>
Status FindMatchingIndices(const Array& keys, const Scalar& query, const
bool from_back, FoundItem callback);
RETURN_NOT_OK(FindMatchingIndices(...field(0), query_key_scalar, from_back,
[&](int64_t index) {
// Scalar case: append to builder or call GetScalar()
// Array case: append to ItemType builder or List<ItemType> builder
});
```
that way all the logic gets centralized
##########
File path: cpp/src/arrow/compute/kernels/scalar_nested.cc
##########
@@ -428,6 +428,290 @@ const FunctionDoc make_struct_doc{"Wrap Arrays into a
StructArray",
"specified through MakeStructOptions."),
{"*args"},
"MakeStructOptions"};
+template <typename KeyType>
+struct MapArrayLookupFunctor {
+ static Result<int64_t> FindOneMapValueIndex(const Array& keys,
+ const Scalar& query_key_scalar,
+ const int64_t start, const
int64_t end,
+ const bool from_back = false) {
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ int64_t match_idx = -1;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
+ if (!from_back) {
+ match_idx = index;
+ return Status::Cancelled("Found first matching key");
+ } else {
+ match_idx = index;
+ }
+ }
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ },
+ [&]() -> Status {
+ if (index < end) {
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ }));
+
+ return match_idx;
+ }
+
+ static Result<std::unique_ptr<ArrayBuilder>> GetBuiltArray(
+ const Array& keys, const Array& items, const Scalar& query_key_scalar,
+ bool& found_at_least_one_key, const int64_t& start, const int64_t& end,
+ KernelContext* ctx) {
+ std::unique_ptr<ArrayBuilder> builder;
+ RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(), items.type(), &builder));
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
+ found_at_least_one_key = true;
+ RETURN_NOT_OK(builder->AppendArraySlice(*items.data(), index,
1));
+ }
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ },
+ [&]() -> Status {
+ if (index < end) {
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ }));
+
+ return std::move(builder);
+ }
+
+ static Status ExecMapArray(KernelContext* ctx, const ExecBatch& batch,
Datum* out) {
+ const auto& options = OptionsWrapper<MapArrayLookupOptions>::Get(ctx);
+ const auto& query_key = options.query_key;
+ const auto& occurrence = options.occurrence;
+ const MapArray map_array(batch[0].array());
+
+ std::shared_ptr<arrow::Array> keys = map_array.keys();
+ std::shared_ptr<arrow::Array> items = map_array.items();
+ auto offsets = std::dynamic_pointer_cast<Int32Array>(map_array.offsets());
+
+ std::unique_ptr<ArrayBuilder> builder;
+ if (occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(),
+ list(map_array.map_type()->item_type()),
&builder));
+
+ for (int64_t map_array_idx = 0; map_array_idx < map_array.length();
+ ++map_array_idx) {
+ if (!map_array.IsValid(map_array_idx)) {
+ RETURN_NOT_OK(builder->AppendNull());
+ continue;
+ }
+
+ int64_t start = offsets->Value(map_array_idx);
+ int64_t end = offsets->Value(map_array_idx + 1);
+ std::unique_ptr<ArrayBuilder> list_builder;
+ bool found_at_least_one_key = false;
+ ARROW_ASSIGN_OR_RAISE(
+ list_builder, GetBuiltArray(*keys, *items, *query_key,
found_at_least_one_key,
+ start, end, ctx));
+ if (!found_at_least_one_key) {
+ RETURN_NOT_OK(builder->AppendNull());
+ } else {
+ ARROW_ASSIGN_OR_RAISE(auto list_result, list_builder->Finish());
+ RETURN_NOT_OK(builder->AppendScalar(ListScalar(list_result)));
+ }
+ list_builder->Reset();
+ }
+ ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish());
+ out->value = result->data();
+ } else {
+ RETURN_NOT_OK(
+ MakeBuilder(ctx->memory_pool(), map_array.map_type()->item_type(),
&builder));
+
+ for (int64_t map_array_idx = 0; map_array_idx < map_array.length();
+ ++map_array_idx) {
+ if (!map_array.IsValid(map_array_idx)) {
+ RETURN_NOT_OK(builder->AppendNull());
+ continue;
+ }
+ int64_t start = offsets->Value(map_array_idx);
+ int64_t end = offsets->Value(map_array_idx + 1);
+ bool from_back = (occurrence == MapArrayLookupOptions::LAST);
+
+ ARROW_ASSIGN_OR_RAISE(
+ int64_t key_match_idx,
+ FindOneMapValueIndex(*keys, *query_key, start, end, from_back));
+ if (key_match_idx != -1) {
+ RETURN_NOT_OK(builder->AppendArraySlice(*items->data(),
key_match_idx, 1));
+ } else {
+ RETURN_NOT_OK(builder->AppendNull());
+ }
+ }
+ ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish());
+ out->value = result->data();
+ }
+
+ return Status::OK();
+ }
+
+ static Status ExecMapScalar(KernelContext* ctx, const ExecBatch& batch,
Datum* out) {
+ const auto& options = OptionsWrapper<MapArrayLookupOptions>::Get(ctx);
+ const auto& query_key = options.query_key;
+ const auto& occurrence = options.occurrence;
+
+ std::shared_ptr<DataType> item_type =
+ checked_cast<const MapType&>(*batch[0].type()).item_type();
+ const auto& map_scalar = batch[0].scalar_as<MapScalar>();
+
+ if (ARROW_PREDICT_FALSE(!map_scalar.is_valid)) {
+ if (options.occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ out->value = MakeNullScalar(list(item_type));
+ } else {
+ out->value = MakeNullScalar(item_type);
+ }
+ return Status::OK();
+ }
+
+ const auto& struct_array = checked_cast<const
StructArray&>(*map_scalar.value);
+ const std::shared_ptr<Array> keys = struct_array.field(0);
+ const std::shared_ptr<Array> items = struct_array.field(1);
+
+ if (occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ bool found_at_least_one_key = false;
+ std::unique_ptr<ArrayBuilder> builder;
+ ARROW_ASSIGN_OR_RAISE(
+ builder, GetBuiltArray(*keys, *items, *query_key,
found_at_least_one_key, 0,
+ struct_array.length(), ctx));
+
+ if (!found_at_least_one_key) {
+ out->value = MakeNullScalar(list(items->type()));
+ } else {
+ ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish());
+ ARROW_ASSIGN_OR_RAISE(out->value, MakeScalar(list(items->type()),
result));
+ }
+ } else { /* occurrence == FIRST || LAST */
+ bool from_back = (occurrence == MapArrayLookupOptions::LAST);
+
+ ARROW_ASSIGN_OR_RAISE(
+ int64_t key_match_idx,
+ FindOneMapValueIndex(*keys, *query_key, 0, struct_array.length(),
from_back));
+ if (key_match_idx != -1) {
+ ARROW_ASSIGN_OR_RAISE(out->value, items->GetScalar(key_match_idx));
+ } else {
+ out->value = MakeNullScalar(items->type());
+ }
+ }
+ return Status::OK();
+ }
+};
+
+Result<ValueDescr> ResolveMapArrayLookupType(KernelContext* ctx,
+ const std::vector<ValueDescr>&
descrs) {
+ const auto& options = OptionsWrapper<MapArrayLookupOptions>::Get(ctx);
+ std::shared_ptr<DataType> type = descrs.front().type;
+ std::shared_ptr<DataType> item_type = checked_cast<const
MapType&>(*type).item_type();
+ std::shared_ptr<DataType> key_type = checked_cast<const
MapType&>(*type).key_type();
+
+ if (!options.query_key || !options.query_key->type ||
+ !options.query_key->type->Equals(key_type)) {
+ return Status::TypeError(
+ "map_array_lookup: query_key type and MapArray key_type don't match.
Expected "
+ "type: ",
+ *item_type, ", but got type: ", *options.query_key->type);
+ }
+
+ if (options.occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ return ValueDescr(list(item_type), descrs.front().shape);
+ } else { /* occurrence == FIRST || LAST */
+ return ValueDescr(item_type, descrs.front().shape);
+ }
+}
+
+struct ResolveMapArrayLookup {
+ KernelContext* ctx;
+ const ExecBatch& batch;
+ Datum* out;
+
+ template <typename KeyType>
+ Status Execute() {
+ if (batch[0].kind() == Datum::SCALAR) {
+ return MapArrayLookupFunctor<KeyType>::ExecMapScalar(ctx, batch, out);
+ }
+ return MapArrayLookupFunctor<KeyType>::ExecMapArray(ctx, batch, out);
+ }
+
+ template <typename KeyType>
+ enable_if_physical_integer<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
+ template <typename KeyType>
+ enable_if_decimal<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
+ template <typename KeyType>
+ enable_if_base_binary<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
+ template <typename KeyType>
+ enable_if_boolean<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
Review comment:
For MonthDayNanoInterval, you can just overload it like normal (`Status
Visit(const MonthDayNanoInterval& type)`). I think `VisitArrayValuesInline`
should work (the type should be the `MonthDayNanos` struct:
https://github.com/apache/arrow/blob/8daa7a4ed5629c0020dadf7325a6b523bdfc62e9/cpp/src/arrow/type.h#L1382)
##########
File path: cpp/src/arrow/compute/kernels/scalar_nested.cc
##########
@@ -428,6 +428,290 @@ const FunctionDoc make_struct_doc{"Wrap Arrays into a
StructArray",
"specified through MakeStructOptions."),
{"*args"},
"MakeStructOptions"};
+template <typename KeyType>
+struct MapArrayLookupFunctor {
+ static Result<int64_t> FindOneMapValueIndex(const Array& keys,
+ const Scalar& query_key_scalar,
+ const int64_t start, const
int64_t end,
+ const bool from_back = false) {
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ int64_t match_idx = -1;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
+ if (!from_back) {
+ match_idx = index;
+ return Status::Cancelled("Found first matching key");
+ } else {
+ match_idx = index;
+ }
+ }
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ },
+ [&]() -> Status {
+ if (index < end) {
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ }));
+
+ return match_idx;
+ }
+
+ static Result<std::unique_ptr<ArrayBuilder>> GetBuiltArray(
+ const Array& keys, const Array& items, const Scalar& query_key_scalar,
+ bool& found_at_least_one_key, const int64_t& start, const int64_t& end,
+ KernelContext* ctx) {
+ std::unique_ptr<ArrayBuilder> builder;
+ RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(), items.type(), &builder));
+ const auto query_key = UnboxScalar<KeyType>::Unbox(query_key_scalar);
+ int64_t index = 0;
+ ARROW_UNUSED(VisitArrayValuesInline<KeyType>(
+ *keys.data(),
+ [&](decltype(query_key) key) -> Status {
+ if (index < start) {
+ ++index;
+ return Status::OK();
+ } else if (index < end) {
+ if (key == query_key) {
+ found_at_least_one_key = true;
+ RETURN_NOT_OK(builder->AppendArraySlice(*items.data(), index,
1));
+ }
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ },
+ [&]() -> Status {
+ if (index < end) {
+ ++index;
+ return Status::OK();
+ } else {
+ return Status::Cancelled("End reached");
+ }
+ }));
+
+ return std::move(builder);
+ }
+
+ static Status ExecMapArray(KernelContext* ctx, const ExecBatch& batch,
Datum* out) {
+ const auto& options = OptionsWrapper<MapArrayLookupOptions>::Get(ctx);
+ const auto& query_key = options.query_key;
+ const auto& occurrence = options.occurrence;
+ const MapArray map_array(batch[0].array());
+
+ std::shared_ptr<arrow::Array> keys = map_array.keys();
+ std::shared_ptr<arrow::Array> items = map_array.items();
+ auto offsets = std::dynamic_pointer_cast<Int32Array>(map_array.offsets());
+
+ std::unique_ptr<ArrayBuilder> builder;
+ if (occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ RETURN_NOT_OK(MakeBuilder(ctx->memory_pool(),
+ list(map_array.map_type()->item_type()),
&builder));
+
+ for (int64_t map_array_idx = 0; map_array_idx < map_array.length();
+ ++map_array_idx) {
+ if (!map_array.IsValid(map_array_idx)) {
+ RETURN_NOT_OK(builder->AppendNull());
+ continue;
+ }
+
+ int64_t start = offsets->Value(map_array_idx);
+ int64_t end = offsets->Value(map_array_idx + 1);
+ std::unique_ptr<ArrayBuilder> list_builder;
+ bool found_at_least_one_key = false;
+ ARROW_ASSIGN_OR_RAISE(
+ list_builder, GetBuiltArray(*keys, *items, *query_key,
found_at_least_one_key,
+ start, end, ctx));
+ if (!found_at_least_one_key) {
+ RETURN_NOT_OK(builder->AppendNull());
+ } else {
+ ARROW_ASSIGN_OR_RAISE(auto list_result, list_builder->Finish());
+ RETURN_NOT_OK(builder->AppendScalar(ListScalar(list_result)));
+ }
+ list_builder->Reset();
+ }
+ ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish());
+ out->value = result->data();
+ } else {
+ RETURN_NOT_OK(
+ MakeBuilder(ctx->memory_pool(), map_array.map_type()->item_type(),
&builder));
+
+ for (int64_t map_array_idx = 0; map_array_idx < map_array.length();
+ ++map_array_idx) {
+ if (!map_array.IsValid(map_array_idx)) {
+ RETURN_NOT_OK(builder->AppendNull());
+ continue;
+ }
+ int64_t start = offsets->Value(map_array_idx);
+ int64_t end = offsets->Value(map_array_idx + 1);
+ bool from_back = (occurrence == MapArrayLookupOptions::LAST);
+
+ ARROW_ASSIGN_OR_RAISE(
+ int64_t key_match_idx,
+ FindOneMapValueIndex(*keys, *query_key, start, end, from_back));
+ if (key_match_idx != -1) {
+ RETURN_NOT_OK(builder->AppendArraySlice(*items->data(),
key_match_idx, 1));
+ } else {
+ RETURN_NOT_OK(builder->AppendNull());
+ }
+ }
+ ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish());
+ out->value = result->data();
+ }
+
+ return Status::OK();
+ }
+
+ static Status ExecMapScalar(KernelContext* ctx, const ExecBatch& batch,
Datum* out) {
+ const auto& options = OptionsWrapper<MapArrayLookupOptions>::Get(ctx);
+ const auto& query_key = options.query_key;
+ const auto& occurrence = options.occurrence;
+
+ std::shared_ptr<DataType> item_type =
+ checked_cast<const MapType&>(*batch[0].type()).item_type();
+ const auto& map_scalar = batch[0].scalar_as<MapScalar>();
+
+ if (ARROW_PREDICT_FALSE(!map_scalar.is_valid)) {
+ if (options.occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ out->value = MakeNullScalar(list(item_type));
+ } else {
+ out->value = MakeNullScalar(item_type);
+ }
+ return Status::OK();
+ }
+
+ const auto& struct_array = checked_cast<const
StructArray&>(*map_scalar.value);
+ const std::shared_ptr<Array> keys = struct_array.field(0);
+ const std::shared_ptr<Array> items = struct_array.field(1);
+
+ if (occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ bool found_at_least_one_key = false;
+ std::unique_ptr<ArrayBuilder> builder;
+ ARROW_ASSIGN_OR_RAISE(
+ builder, GetBuiltArray(*keys, *items, *query_key,
found_at_least_one_key, 0,
+ struct_array.length(), ctx));
+
+ if (!found_at_least_one_key) {
+ out->value = MakeNullScalar(list(items->type()));
+ } else {
+ ARROW_ASSIGN_OR_RAISE(auto result, builder->Finish());
+ ARROW_ASSIGN_OR_RAISE(out->value, MakeScalar(list(items->type()),
result));
+ }
+ } else { /* occurrence == FIRST || LAST */
+ bool from_back = (occurrence == MapArrayLookupOptions::LAST);
+
+ ARROW_ASSIGN_OR_RAISE(
+ int64_t key_match_idx,
+ FindOneMapValueIndex(*keys, *query_key, 0, struct_array.length(),
from_back));
+ if (key_match_idx != -1) {
+ ARROW_ASSIGN_OR_RAISE(out->value, items->GetScalar(key_match_idx));
+ } else {
+ out->value = MakeNullScalar(items->type());
+ }
+ }
+ return Status::OK();
+ }
+};
+
+Result<ValueDescr> ResolveMapArrayLookupType(KernelContext* ctx,
+ const std::vector<ValueDescr>&
descrs) {
+ const auto& options = OptionsWrapper<MapArrayLookupOptions>::Get(ctx);
+ std::shared_ptr<DataType> type = descrs.front().type;
+ std::shared_ptr<DataType> item_type = checked_cast<const
MapType&>(*type).item_type();
+ std::shared_ptr<DataType> key_type = checked_cast<const
MapType&>(*type).key_type();
+
+ if (!options.query_key || !options.query_key->type ||
+ !options.query_key->type->Equals(key_type)) {
+ return Status::TypeError(
+ "map_array_lookup: query_key type and MapArray key_type don't match.
Expected "
+ "type: ",
+ *item_type, ", but got type: ", *options.query_key->type);
+ }
+
+ if (options.occurrence == MapArrayLookupOptions::Occurrence::ALL) {
+ return ValueDescr(list(item_type), descrs.front().shape);
+ } else { /* occurrence == FIRST || LAST */
+ return ValueDescr(item_type, descrs.front().shape);
+ }
+}
+
+struct ResolveMapArrayLookup {
+ KernelContext* ctx;
+ const ExecBatch& batch;
+ Datum* out;
+
+ template <typename KeyType>
+ Status Execute() {
+ if (batch[0].kind() == Datum::SCALAR) {
+ return MapArrayLookupFunctor<KeyType>::ExecMapScalar(ctx, batch, out);
+ }
+ return MapArrayLookupFunctor<KeyType>::ExecMapArray(ctx, batch, out);
+ }
+
+ template <typename KeyType>
+ enable_if_physical_integer<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
+ template <typename KeyType>
+ enable_if_decimal<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
+ template <typename KeyType>
+ enable_if_base_binary<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
+ template <typename KeyType>
+ enable_if_boolean<KeyType, Status> Visit(const KeyType& type) {
+ return Execute<KeyType>();
+ }
+
Review comment:
You can use `enable_if_same` to provide an overload specifically for
FixedSizeBinaryType
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
