Merge pull request #10 from kortemik/partial-token

Partial token
teragrep · Oct 11, 2023 · 4a2db47 · 4a2db47
2 parents 058b6a8 + f4f15f4
commit 4a2db47
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 2 deletions.
diff --git a/src/main/java/com/teragrep/blf_01/Tokenizer.java b/src/main/java/com/teragrep/blf_01/Tokenizer.java
@@ -59,8 +59,12 @@ public class Tokenizer {
     final Entanglement entanglement;
     final TokenScan majorTokenScan;
     final TokenScan minorTokenScan;
+    final long maxTokenCount;
 
     public Tokenizer() {
+        this(Long.MAX_VALUE);
+    }
+    public Tokenizer(long maxTokenCount) {
 
         final MajorDelimiters majorDelimiters = new MajorDelimiters();
         final MinorDelimiters minorDelimiters = new MinorDelimiters();
@@ -69,7 +73,7 @@ public Tokenizer() {
         this.entanglement = new Entanglement();
         this.majorTokenScan = new TokenScan(majorDelimiters);
         this.minorTokenScan = new TokenScan(minorDelimiters);
-
+        this.maxTokenCount = maxTokenCount;
     }
 
     /**
@@ -94,8 +98,14 @@ public List<Token> tokenize(InputStream is) {
 
             ArrayList<Token> minorTokens = minorTokenScan.findBy(stream);
 
+            ArrayList<Token> tokens;
+            if (minorTokens.size() > maxTokenCount) {
+                tokens = minorTokens;
+            } else {
+                tokens = entanglement.entangle(minorTokens);
+            }
 
-            allTokens.addAll(entanglement.entangle(minorTokens));
+            allTokens.addAll(tokens);
         }
 
         return allTokens;

diff --git a/src/test/java/com/teragrep/blf_01/TokenizerTest.java b/src/test/java/com/teragrep/blf_01/TokenizerTest.java
@@ -90,6 +90,26 @@ public void testTokenization() {
 
     }
 
+
+    @Test
+    public void testTokenizerSizeLimit() {
+        Tokenizer tokenizer = new Tokenizer(3);
+        String input = "Abc#####Xyz";
+        ByteArrayInputStream bais = new ByteArrayInputStream(input.getBytes(StandardCharsets.UTF_8));
+        List<Token> result = tokenizer.tokenize(bais);
+
+        List<String> expected =
+                Arrays.asList(
+                        "Abc#####Xyz", "Abc", "#", "#", "#", "#", "#", "Xyz"
+                );
+
+        assertTrue(result.stream()
+                .map(Token::toString)
+                .collect(Collectors.toList())
+                .containsAll(expected));
+
+    }
+
     @Test
     @Benchmark
     public void tokenizeFileInput() throws FileNotFoundException {