From 416f407ccd168a397658f25cbd0d879d17909343 Mon Sep 17 00:00:00 2001
From: ejblanco <ericjoel.blancohermida@telefonica.com>
Date: Tue, 5 Nov 2024 09:10:39 +0100
Subject: [PATCH] feat: add option to hash elements and to repeate them X
 number of times

---
 README.md                                     | 14 +++++
 .../avro/random/generator/Generator.java      | 61 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/README.md b/README.md
index 228a79f..237af0b 100644
--- a/README.md
+++ b/README.md
@@ -86,7 +86,21 @@ if &lt;restart&gt; is equal to &lt;start&gt;. If provided with a boolean
 schema, only &lt;start&gt; may be specified; the resulting values will
 begin with &lt;start&gt; and alternate from `true` to `false` and from
 `false` to `true` from that point on.
+  + __hashed:__ If this options is set to `true`, the iteration will be hashed. 
+  This is useful when you want to generate IDs that can join with other datasets.
+  + __num_repetitions:__ Number of times you want each element to be repeated. 
+  This is useful when you want to generate a dataset with a lot of repeated IDs.
+```json
+              "arg.properties": {
+                "iteration": {
+                  "start": "0"
+                },
+                "hashed": true,
+                "num_repetitions": 4
+              }
+```
 > ITERATION_STEP environment var can be used as script argument
+
 + __range:__ A JSON object that conforms to the following formats:
     - `{"min": <min>, "max": <max>}` (at least one of "min" or "max" must be
     specified). If provided, ensures that the generated number will be
diff --git a/app/src/main/java/io/confluent/avro/random/generator/Generator.java b/app/src/main/java/io/confluent/avro/random/generator/Generator.java
index 71df76f..9160ffb 100644
--- a/app/src/main/java/io/confluent/avro/random/generator/Generator.java
+++ b/app/src/main/java/io/confluent/avro/random/generator/Generator.java
@@ -36,6 +36,8 @@
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.file.Files;
+import java.security.MessageDigest;
+import java.security.NoSuchAlgorithmException;
 import java.util.*;
 
 /**
@@ -93,6 +95,15 @@ public class Generator {
    */
   public static final String SUFFIX_PROP = "suffix";
 
+    /**
+     * If you want to apply a hash function to generated values in iterations. As a boolean.
+     */
+  public static final String HASHED_PROP = "hashed";
+
+    /**
+     * Number of times the same element would be repeated in iterations. As an integer.
+     */
+  public static final String REPEAT_PROP = "num_repetitions";
   /**
    * The name of the attribute for specifying specific values which should be randomly chosen from
    * when generating values for the schema. Can be given as either an array of values or an object
@@ -1042,6 +1053,9 @@ private Iterator<Object> parseIterations(Schema schema, Map propertiesProp) {
       case DOUBLE:
         return getDoubleIterator(iterationProps);
       case STRING:
+        if (propertiesProp.containsKey(HASHED_PROP) && (Boolean) propertiesProp.get(HASHED_PROP)) {
+          return createHashedStringIterator(getIntegerIterator(iterationProps), propertiesProp);
+        }
         return createStringIterator(getIntegerIterator(iterationProps), propertiesProp);
       default:
         throw new UnsupportedOperationException(String.format(
@@ -1157,6 +1171,53 @@ public Object next() {
     };
   }
 
+  private String hashNumber(BigInteger number) {
+    try {
+      MessageDigest md = MessageDigest.getInstance("SHA-256");
+      byte[] hash = md.digest(number.toByteArray());
+      StringBuilder hexString = new StringBuilder();
+      for (byte b : hash) {
+        String hex = Integer.toHexString(0xff & b);
+        if (hex.length() == 1) hexString.append('0');
+        hexString.append(hex);
+      }
+      return hexString.toString();
+    } catch (NoSuchAlgorithmException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private Iterator<Object> createHashedStringIterator(Iterator<Object> inner, Map propertiesProp) {
+    return new Iterator<Object>() {
+      private Integer count = 0;
+      private String currentValue = "";
+      private final Integer repeatValue = getIntegerNumberField(
+          HASHED_PROP,
+          REPEAT_PROP,
+          propertiesProp
+      );
+      private final Integer numRepeat = repeatValue == null ? 1 : repeatValue;
+
+      @Override
+      public boolean hasNext() {
+        return inner.hasNext();
+      }
+
+      @Override
+      public Object next() {
+        if (numRepeat == 1) {
+          return hashNumber(BigInteger.valueOf((Integer) inner.next()));
+        }
+        if (count % numRepeat == 0) {
+          currentValue = hashNumber(BigInteger.valueOf((Integer) inner.next()));
+          count = 0;
+        }
+        count++;
+        return currentValue;
+      }
+    };
+  }
+
   private Iterator<Object> getIntegerIterator(Map iterationProps) {
     Integer iterationStartField = getIntegerNumberField(
         ITERATION_PROP,