ARROW-15089: [C++][Compute] Implement kernel to lookup a MapArray item for a given key

cpp/src/arrow/compute/api_scalar.cc

@@ -254,6 +254,26 @@ struct EnumTraits<compute::RandomOptions::Initializer>
  }
 };
 
+template <>
+struct EnumTraits<compute::MapLookupOptions::Occurrence>
+ : BasicEnumTraits<compute::MapLookupOptions::Occurrence,
+ compute::MapLookupOptions::Occurrence::FIRST,
+ compute::MapLookupOptions::Occurrence::LAST,
+ compute::MapLookupOptions::Occurrence::ALL> {
+ static std::string name() { return "MapLookupOptions::Occurrence"; }
+ static std::string value_name(compute::MapLookupOptions::Occurrence value) {
+ switch (value) {
+ case compute::MapLookupOptions::Occurrence::FIRST:
+ return "FIRST";
+ case compute::MapLookupOptions::Occurrence::LAST:
+ return "LAST";
+ case compute::MapLookupOptions::Occurrence::ALL:
+ return "ALL";
+ }
+ return "<INVALID>";
+ }
+};
+
 } // namespace internal
 
 namespace compute {
@@ -287,6 +307,9 @@ static auto kMakeStructOptionsType = GetFunctionOptionsType<MakeStructOptions>(
  DataMember("field_names", &MakeStructOptions::field_names),
  DataMember("field_nullability", &MakeStructOptions::field_nullability),
  DataMember("field_metadata", &MakeStructOptions::field_metadata));
+static auto kMapLookupOptionsType = GetFunctionOptionsType<MapLookupOptions>(
+ DataMember("occurrence", &MapLookupOptions::occurrence),
+ DataMember("query_key", &MapLookupOptions::query_key));
 static auto kMatchSubstringOptionsType = GetFunctionOptionsType<MatchSubstringOptions>(
  DataMember("pattern", &MatchSubstringOptions::pattern),
  DataMember("ignore_case", &MatchSubstringOptions::ignore_case));
@@ -344,6 +367,7 @@ static auto kRandomOptionsType = GetFunctionOptionsType<RandomOptions>(
  DataMember("length", &RandomOptions::length),
  DataMember("initializer", &RandomOptions::initializer),
  DataMember("seed", &RandomOptions::seed));
+
 } // namespace
 } // namespace internal
 
@@ -399,6 +423,15 @@ MakeStructOptions::MakeStructOptions(std::vector<std::string> n)
 MakeStructOptions::MakeStructOptions() : MakeStructOptions(std::vector<std::string>()) {}
 constexpr char MakeStructOptions::kTypeName[];
 
+MapLookupOptions::MapLookupOptions(std::shared_ptr<Scalar> query_key,
+ Occurrence occurrence)
+ : FunctionOptions(internal::kMapLookupOptionsType),
+ query_key(std::move(query_key)),
+ occurrence(occurrence) {}
+MapLookupOptions::MapLookupOptions()
+ : MapLookupOptions(std::make_shared<NullScalar>(), Occurrence::FIRST) {}
+constexpr char MapLookupOptions::kTypeName[];
+
 MatchSubstringOptions::MatchSubstringOptions(std::string pattern, bool ignore_case)
  : FunctionOptions(internal::kMatchSubstringOptionsType),
  pattern(std::move(pattern)),
@@ -554,6 +587,7 @@ void RegisterScalarOptions(FunctionRegistry* registry) {
  DCHECK_OK(registry->AddFunctionOptionsType(kExtractRegexOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kJoinOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kMakeStructOptionsType));
+ DCHECK_OK(registry->AddFunctionOptionsType(kMapLookupOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kMatchSubstringOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kNullOptionsType));
  DCHECK_OK(registry->AddFunctionOptionsType(kPadOptionsType));
@@ -785,5 +819,13 @@ Result<Datum> Week(const Datum& arg, WeekOptions options, ExecContext* ctx) {
  return CallFunction("week", {arg}, &options, ctx);
 }
 
+// ----------------------------------------------------------------------
+// Structural transforms
+Result<Datum> MapLookup(const Datum& arg, MapLookupOptions options, ExecContext* ctx) {
+ return CallFunction("map_lookup", {arg}, &options, ctx);
+}
+
+// ----------------------------------------------------------------------
+
 } // namespace compute
 } // namespace arrow

cpp/src/arrow/compute/api_scalar.h

@@ -470,6 +470,30 @@ class ARROW_EXPORT RandomOptions : public FunctionOptions {
  uint64_t seed;
 };
 
+/// Options for map_lookup function
+class ARROW_EXPORT MapLookupOptions : public FunctionOptions {
+ public:
+ enum Occurrence {
+ /// Return the first matching value
+ FIRST,
+ /// Return the last matching value
+ LAST,
+ /// Return all matching values
+ ALL
+ };
+
+ explicit MapLookupOptions(std::shared_ptr<Scalar> query_key, Occurrence occurrence);
+ MapLookupOptions();
+
+ constexpr static char const kTypeName[] = "MapLookupOptions";
+
+ /// The key to lookup in the map
+ std::shared_ptr<Scalar> query_key;
+
+ /// Whether to return the first, last, or all matching values
+ Occurrence occurrence;
+};
+
 /// @}
 
 /// \brief Get the absolute value of a value.
@@ -1350,5 +1374,20 @@ ARROW_EXPORT Result<Datum> AssumeTimezone(const Datum& values,
  AssumeTimezoneOptions options,
  ExecContext* ctx = NULLPTR);
 
+/// \brief Finds either the FIRST, LAST, or ALL items with a key that matches the given
+/// query key in a map.
+///
+/// Returns an array of items for FIRST and LAST, and an array of list of items for ALL.
+///
+/// \param[in] map to look in
+/// \param[in] options to pass a query key and choose which matching keys to return
+/// (FIRST, LAST or ALL)
+/// \param[in] ctx the function execution context, optional
+/// \return the resulting datum
+///
+/// \since 8.0.0
+/// \note API not yet finalized
+ARROW_EXPORT Result<Datum> MapLookup(const Datum& map, MapLookupOptions options,
+ ExecContext* ctx = NULLPTR);
 } // namespace compute
 } // namespace arrow