Merge pull request #6 from elliVM/set

aggregator returns set instead of list
teragrep · Jul 11, 2023 · 8f9a06b · 8f9a06b
2 parents 7bb0219 + 8f4f66d
commit 8f9a06b
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/src/main/java/com/teragrep/functions/dpf_03/TokenAggregator.java b/src/main/java/com/teragrep/functions/dpf_03/TokenAggregator.java
@@ -2,7 +2,7 @@
 
 /*
  * Teragrep Tokenizer DPF-03
- * Copyright (C) 2019, 2020, 2021, 2022  Suomen Kanuuna Oy
+ * Copyright (C) 2019, 2020, 2021, 2022, 2023  Suomen Kanuuna Oy
  *
  * This program is free software: you can redistribute it and/or modify
  * it under the terms of the GNU Affero General Public License as published by
@@ -47,8 +47,8 @@
  */
 
 import java.io.Serializable;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashSet;
+import java.util.Set;
 
 import com.teragrep.blf_01.tokenizer.Tokenizer;
 import org.apache.spark.sql.Encoder;
@@ -57,7 +57,7 @@
 import org.apache.spark.sql.expressions.Aggregator;
 
 
-public class TokenAggregator extends Aggregator<Row, TokenBuffer, List<String>> implements Serializable {
+public class TokenAggregator extends Aggregator<Row, TokenBuffer, Set<String>> implements Serializable {
 
     private final String column;
 
@@ -89,8 +89,8 @@ public TokenBuffer merge(TokenBuffer b1, TokenBuffer b2) {
     }
 
     @Override
-    public List<String> finish(TokenBuffer reduction) {
-	    return new ArrayList<>(reduction.getMap().keySet());
+    public Set<String> finish(TokenBuffer reduction) {
+	    return new HashSet<>(reduction.getMap().keySet());
     }
 
     @Override
@@ -100,7 +100,7 @@ public Encoder<TokenBuffer> bufferEncoder() {
 
     @Override
     public Encoder outputEncoder() {
-	    List<String> stringList = new ArrayList<>();
-        return Encoders.kryo(stringList.getClass());
+	    Set<String> stringSet = new HashSet<>();
+        return Encoders.kryo(stringSet.getClass());
     }
 }