Prestissimo ApproxMostFrequent JSON (facebookincubator#12189)

Summary: X-link: prestodb/presto#24450 Pull Request resolved: facebookincubator#12189 Prestissimo ApproxMostFrequent is not implemented for JSON. This PR adds support for JSON type. Reviewed By: amitkdutta Differential Revision: D68287956 fbshipit-source-id: b1ebb9217b1148d52385f54a4864d67ddf85c773
yuandagits · Feb 5, 2025 · f6fe8a3 · f6fe8a3
1 parent 5d13c13
commit f6fe8a3
Show file tree

Hide file tree

Showing 2 changed files with 40 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/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