Prestissimo ApproxMostFrequent JSON (facebookincubator#12189)

Summary: X-link: prestodb/presto#24450 Prestissimo ApproxMostFrequent is not implemented for JSON. This PR adds support for JSON type. Differential Revision: D68287956
natashasehgal · Jan 31, 2025 · 9731553 · 9731553
1 parent 4bd1221
commit 9731553
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 1 deletion.
diff --git a/velox/functions/prestosql/aggregates/ApproxMostFrequentAggregate.cpp b/velox/functions/prestosql/aggregates/ApproxMostFrequentAggregate.cpp
@@ -477,7 +477,13 @@ void registerApproxMostFrequentAggregate(
     bool overwrite) {
   std::vector<std::shared_ptr<exec::AggregateFunctionSignature>> signatures;
   for (const auto& valueType :
-       {"boolean", "tinyint", "smallint", "integer", "bigint", "varchar"}) {
+       {"boolean",
+        "tinyint",
+        "smallint",
+        "integer",
+        "bigint",
+        "varchar",
+        "json"}) {
     signatures.push_back(
         exec::AggregateFunctionSignatureBuilder()
             .returnType(fmt::format("map({},bigint)", valueType))

diff --git a/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp b/velox/functions/prestosql/aggregates/RegisterAggregateFunctions.cpp
@@ -15,6 +15,7 @@
  */
 #include "velox/functions/prestosql/aggregates/RegisterAggregateFunctions.h"
 #include "velox/exec/Aggregate.h"
+#include "velox/functions/prestosql/types/JsonType.h"
 
 namespace facebook::velox::aggregate::prestosql {
 
@@ -155,6 +156,7 @@ void registerAllAggregateFunctions(
     bool withCompanionFunctions,
     bool onlyPrestoSignatures,
     bool overwrite) {
+  registerJsonType();
   registerApproxDistinctAggregates(prefix, withCompanionFunctions, overwrite);
   registerApproxMostFrequentAggregate(
       prefix, withCompanionFunctions, overwrite);

diff --git a/velox/functions/prestosql/aggregates/tests/ApproxMostFrequentTest.cpp b/velox/functions/prestosql/aggregates/tests/ApproxMostFrequentTest.cpp
@@ -281,5 +281,38 @@ TEST_F(ApproxMostFrequentTestBoolean, basic) {
       {input}, {"c0"}, {"approx_most_frequent(3, c5, 31)"}, {expected});
 }
 
+class ApproxMostFrequentTestJson : public AggregationTestBase {
+ protected:
+  void SetUp() override {
+    AggregationTestBase::SetUp();
+  }
+};
+
+TEST_F(ApproxMostFrequentTestJson, basic) {
+  // JSON strings as input
+  std::vector<std::string> jsonStrings = {
+      "{\"type\": \"store\"}",
+      "{\"type\": \"fruit\"}",
+      "{\"type\": \"fruit\"}",
+      "{\"type\": \"book\"}",
+      "{\"type\": \"store\"}",
+      "{\"type\": \"fruit\"}"};
+
+  auto inputVector = makeFlatVector<StringView>(
+      static_cast<vector_size_t>(jsonStrings.size()),
+      [&](auto row) { return StringView(jsonStrings[row]); });
+
+  MapVectorPtr expectedMap = makeMapVector<StringView, int64_t>(
+      {{{StringView("{\"type\": \"fruit\"}"), 3},
+        {StringView("{\"type\": \"store\"}"), 2}}});
+  auto expected = makeRowVector({{expectedMap}});
+
+  testAggregations(
+      {makeRowVector({inputVector})},
+      {},
+      {"approx_most_frequent(2, c0, 31)"},
+      {expected});
+}
+
 } // namespace
 } // namespace facebook::velox::aggregate::test