diff --git a/.gitignore b/.gitignore
index 9d1acbb..d883a6d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,12 +2,3 @@
 **/bazel-java
 **/bazel-out
 **/bazel-testlogs
-
-**/bazel-differential-privacy
-**/.ijwb/
-**/pipelinedp4j/.ijwb/
-**/pipelinedp4j/bazel-pipelinedp4j
-**/pipelinedp4j/MODULE**
-**/examples/.idea/
-**/examples/pipelinedp4j/bazel-pipelinedp4j
-**/examples/pipelinedp4j/MODULE**
diff --git a/cc/algorithms/BUILD b/cc/algorithms/BUILD
index 2ab7d58..764e780 100644
--- a/cc/algorithms/BUILD
+++ b/cc/algorithms/BUILD
@@ -139,7 +139,6 @@
     deps = [
         ":algorithm",
         ":approx-bounds",
-        ":bounded-algorithm",
         ":numerical-mechanisms",
         ":util",
         "//proto:util-lib",
@@ -245,7 +244,6 @@
     deps = [
         ":algorithm",
         ":approx-bounds",
-        ":bounded-algorithm",
         ":numerical-mechanisms",
         ":util",
         "//proto:util-lib",
@@ -291,10 +289,8 @@
     deps = [
         ":algorithm",
         ":approx-bounds",
-        ":bounded-algorithm",
         ":bounded-variance",
         ":numerical-mechanisms",
-        ":util",
         "//proto:util-lib",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
@@ -534,6 +530,7 @@
 cc_library(
     name = "bounded-algorithm",
     hdrs = ["bounded-algorithm.h"],
+    visibility = ["//visibility:private"],
     deps = [
         ":algorithm",
         ":approx-bounds",
@@ -668,10 +665,14 @@
     visibility = ["//visibility:public"],
     deps = [
         ":algorithm",
-        ":bounded-algorithm",
+        ":numerical-mechanisms",
         ":quantile-tree",
+        ":util",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_cc_differential_privacy//base:status_macros",
     ],
 )
 
diff --git a/cc/algorithms/bounded-standard-deviation.h b/cc/algorithms/bounded-standard-deviation.h
index 31676dd..96fad76 100644
--- a/cc/algorithms/bounded-standard-deviation.h
+++ b/cc/algorithms/bounded-standard-deviation.h
@@ -28,10 +28,8 @@
 #include "absl/status/statusor.h"
 #include "algorithms/algorithm.h"
 #include "algorithms/approx-bounds.h"
-#include "algorithms/bounded-algorithm.h"
 #include "algorithms/bounded-variance.h"
 #include "algorithms/numerical-mechanisms.h"
-#include "algorithms/util.h"
 #include "proto/util.h"
 #include "proto/data.pb.h"
 #include "proto/summary.pb.h"
diff --git a/cc/algorithms/bounded-sum.h b/cc/algorithms/bounded-sum.h
index 784ad2c..5de00b2 100644
--- a/cc/algorithms/bounded-sum.h
+++ b/cc/algorithms/bounded-sum.h
@@ -31,13 +31,11 @@
 
 #include "google/protobuf/any.pb.h"
 #include "absl/log/log.h"
-#include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "algorithms/algorithm.h"
 #include "algorithms/approx-bounds.h"
-#include "algorithms/bounded-algorithm.h"
 #include "algorithms/numerical-mechanisms.h"
 #include "algorithms/util.h"
 #include "proto/util.h"
diff --git a/cc/algorithms/bounded-variance.h b/cc/algorithms/bounded-variance.h
index ad4080e..b298977 100644
--- a/cc/algorithms/bounded-variance.h
+++ b/cc/algorithms/bounded-variance.h
@@ -37,7 +37,6 @@
 #include "absl/strings/str_cat.h"
 #include "algorithms/algorithm.h"
 #include "algorithms/approx-bounds.h"
-#include "algorithms/bounded-algorithm.h"
 #include "algorithms/numerical-mechanisms.h"
 #include "algorithms/util.h"
 #include "proto/util.h"
diff --git a/cc/algorithms/quantiles.h b/cc/algorithms/quantiles.h
index 119bb73..fbc2bf2 100644
--- a/cc/algorithms/quantiles.h
+++ b/cc/algorithms/quantiles.h
@@ -18,13 +18,21 @@
 #define DIFFERENTIAL_PRIVACY_CPP_ALGORITHMS_QUANTILES_H_
 
 #include <cstdint>
+#include <memory>
 #include <optional>
+#include <type_traits>
+#include <utility>
+#include <vector>
 
+#include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "algorithms/algorithm.h"
-#include "algorithms/bounded-algorithm.h"
+#include "algorithms/numerical-mechanisms.h"
 #include "algorithms/quantile-tree.h"
+#include "algorithms/util.h"
+#include "base/status_macros.h"
 
 namespace differential_privacy {
 
@@ -239,7 +247,7 @@
   int max_partitions_contributed_ = 1;
   int max_contributions_per_partition_ = 1;
   std::unique_ptr<NumericalMechanismBuilder> mechanism_builder_ =
-      absl::make_unique<LaplaceMechanism::Builder>();
+      std::make_unique<LaplaceMechanism::Builder>();
   std::vector<double> quantiles_;
 
   static absl::Status ValidateQuantiles(std::vector<double>& quantiles) {
diff --git a/cc/algorithms/util.cc b/cc/algorithms/util.cc
index aadcdff..f63b0b9 100644
--- a/cc/algorithms/util.cc
+++ b/cc/algorithms/util.cc
@@ -303,6 +303,10 @@
                             "contributed to (i.e., L0 sensitivity)");
 }
 
+absl::Status ValidateMaxWindows(std::optional<int> max_windows) {
+  return ValidateIsPositive(max_windows, "Maximum number of windows");
+}
+
 absl::Status ValidateMaxContributionsPerPartition(
     std::optional<double> max_contributions_per_partition) {
   return ValidateIsPositive(max_contributions_per_partition,
diff --git a/cc/algorithms/util.h b/cc/algorithms/util.h
index 39001ca..389016f 100644
--- a/cc/algorithms/util.h
+++ b/cc/algorithms/util.h
@@ -522,6 +522,7 @@
 absl::Status ValidateDelta(std::optional<double> delta);
 absl::Status ValidateMaxPartitionsContributed(
     std::optional<double> max_partitions_contributed);
+absl::Status ValidateMaxWindows(std::optional<int> max_windows);
 absl::Status ValidateMaxContributionsPerPartition(
     std::optional<double> max_contributions_per_partition);
 absl::Status ValidateMaxContributions(std::optional<int> max_contributions);
diff --git a/examples/pipelinedp4j/README.md b/examples/pipelinedp4j/README.md
index b95245f..df1f718 100644
--- a/examples/pipelinedp4j/README.md
+++ b/examples/pipelinedp4j/README.md
@@ -90,7 +90,7 @@
 1.  Run the program:
 
     ```shell
-    bazel-bin/BeamExample --inputFilePath=netflix_data.csv --outputFilePath=output.txt
+    bazel-bin/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BeamExample --inputFilePath=netflix_data.csv --outputFilePath=output.txt
     ```
 
 1.  View the results:
diff --git a/examples/pipelinedp4j/pom.xml b/examples/pipelinedp4j/pom.xml
index 4aaed76..af5066c 100644
--- a/examples/pipelinedp4j/pom.xml
+++ b/examples/pipelinedp4j/pom.xml
@@ -46,6 +46,10 @@
                 <configuration>
                     <source>11</source>
                     <target>11</target>
+                    <excludes>
+                        <!-- SparkExample is not supoprted yet. -->
+                        <exclude>com/google/privacy/differentialprivacy/pipelinedp4j/examples/SparkExample.java</exclude>
+                    </excludes>
                 </configuration>
             </plugin>
         </plugins>
diff --git a/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BUILD.bazel b/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BUILD.bazel
index a731aae..26f721e 100644
--- a/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BUILD.bazel
+++ b/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BUILD.bazel
@@ -37,23 +37,22 @@
 java_binary(
     name = "SparkExample",
     srcs = [
-        "SparkExample.java",
         "MovieMetrics.java",
         "MovieView.java",
+        "SparkExample.java",
     ],
     main_class = "com.google.privacy.differentialprivacy.pipelinedp4j.examples.SparkExample",
     deps = [
         "@com_google_privacy_differentialprivacy_pipielinedp4j//main/com/google/privacy/differentialprivacy/pipelinedp4j/api",
-        "@maven//:com_google_guava_guava",
-        "@maven//:info_picocli_picocli",
-        "@maven//:org_jetbrains_kotlin_kotlin_stdlib",
-
-        "@maven//:org_apache_spark_spark_core_2_13",
-        "@maven//:org_apache_spark_spark_sql_2_13",
-        "@maven//:org_apache_spark_spark_mllib_2_13",
-        "@maven//:org_apache_spark_spark_catalyst_2_13",
         "@maven//:com_fasterxml_jackson_core_jackson_databind",
         "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
+        "@maven//:com_google_guava_guava",
+        "@maven//:info_picocli_picocli",
+        "@maven//:org_apache_spark_spark_catalyst_2_13",
+        "@maven//:org_apache_spark_spark_core_2_13",
+        "@maven//:org_apache_spark_spark_mllib_2_13",
+        "@maven//:org_apache_spark_spark_sql_2_13",
+        "@maven//:org_jetbrains_kotlin_kotlin_stdlib",
         "@maven//:org_scala_lang_scala_library",
     ],
 )
diff --git a/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BeamExample.java b/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BeamExample.java
index 5c3aa42..6909f9f 100644
--- a/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BeamExample.java
+++ b/examples/pipelinedp4j/src/main/java/com/google/privacy/differentialprivacy/pipelinedp4j/examples/BeamExample.java
@@ -185,7 +185,7 @@
 
   /**
    * Movie ids (which are group keys for this dataset) are integers from 1 to ~17000. Set public
-   * groups 4500-4509.
+   * groups to a subset of them.
    */
   private static PCollection<String> publiclyKnownMovieIds(Pipeline pipeline) {
     var publicGroupsAsJavaList =
diff --git a/pipelinedp4j/BUILD.bazel b/pipelinedp4j/BUILD.bazel
index b6661f4..8872ad3 100644
--- a/pipelinedp4j/BUILD.bazel
+++ b/pipelinedp4j/BUILD.bazel
@@ -63,7 +63,6 @@
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/spark:spark_collections",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/spark:spark_dp_engine_factory",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/spark:spark_encoders",
-
     ],
     template_file = "pom.template",
 )
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel
index 4ae979f..e177b32 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel
@@ -41,6 +41,5 @@
         "@maven//:com_google_guava_guava",
         "@maven//:org_apache_beam_beam_sdks_java_core",
         "@maven//:org_apache_beam_beam_sdks_java_extensions_avro",
-
     ],
 )
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/PipelineDpCollection.kt b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/PipelineDpCollection.kt
index b52e2e1..0aa60df 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/PipelineDpCollection.kt
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/PipelineDpCollection.kt
@@ -16,16 +16,16 @@
 
 package com.google.privacy.differentialprivacy.pipelinedp4j.api
 
+import com.google.privacy.differentialprivacy.pipelinedp4j.spark.SparkCollection
+import com.google.privacy.differentialprivacy.pipelinedp4j.spark.SparkEncoderFactory
+import org.apache.spark.sql.Dataset
 import com.google.privacy.differentialprivacy.pipelinedp4j.beam.BeamCollection
 import com.google.privacy.differentialprivacy.pipelinedp4j.beam.BeamEncoderFactory
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.EncoderFactory
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.FrameworkCollection
 import com.google.privacy.differentialprivacy.pipelinedp4j.local.LocalCollection
 import com.google.privacy.differentialprivacy.pipelinedp4j.local.LocalEncoderFactory
-import com.google.privacy.differentialprivacy.pipelinedp4j.spark.SparkCollection
-import com.google.privacy.differentialprivacy.pipelinedp4j.spark.SparkEncoderFactory
 import org.apache.beam.sdk.values.PCollection as BeamPCollection
-import org.apache.spark.sql.Dataset
 
 /**
  * An internal interface to represent an arbitrary collection that is supported by PipelineDP4j.
@@ -57,7 +57,8 @@
 }
 
 /** Spark Collection represented as a Spark Dataset. */
-internal data class SparkPipelineDpCollection<T>(val data: Dataset<T>) : PipelineDpCollection<T> {
+internal data class SparkPipelineDpCollection<T>(val data: Dataset<T>) : PipelineDpCollection<T>
+{
   override val encoderFactory = SparkEncoderFactory()
 
   override fun toFrameworkCollection() = SparkCollection<T>(data)
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/QueryBuilder.kt b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/QueryBuilder.kt
index 15de370..80df376 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/QueryBuilder.kt
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/api/QueryBuilder.kt
@@ -17,6 +17,7 @@
 package com.google.privacy.differentialprivacy.pipelinedp4j.api
 
 import org.apache.beam.sdk.values.PCollection as BeamPCollection
+
 import org.apache.spark.sql.Dataset
 
 /**
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncoders.kt b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncoders.kt
index ff6559f..a6e05cf 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncoders.kt
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncoders.kt
@@ -21,13 +21,13 @@
 import com.google.protobuf.Message
 import java.io.InputStream
 import java.io.OutputStream
-import kotlin.reflect.KClass
 import org.apache.beam.sdk.coders.Coder
 import org.apache.beam.sdk.coders.CustomCoder
 import org.apache.beam.sdk.coders.DoubleCoder
 import org.apache.beam.sdk.coders.StringUtf8Coder
 import org.apache.beam.sdk.coders.VarIntCoder
 import org.apache.beam.sdk.extensions.avro.coders.AvroCoder
+
 import org.apache.beam.sdk.extensions.protobuf.ProtoCoder
 
 class BeamEncoder<T>(val coder: Coder<T>) : Encoder<T>
@@ -39,11 +39,10 @@
 
   override fun ints() = BeamEncoder<Int>(VarIntCoder.of())
 
-  override fun <T : Any> records(recordClass: KClass<T>) =
-    BeamEncoder<T>(AvroCoder.of(recordClass.java))
+  override fun <T : Any> records(recordClass: Class<T>) = BeamEncoder<T>(AvroCoder.of(recordClass))
 
-  override fun <T : Message> protos(protoClass: KClass<T>) =
-    BeamEncoder<T>(ProtoCoder.of(protoClass.java))
+  override fun <T : Message> protos(protoClass: Class<T>) =
+    BeamEncoder<T>(ProtoCoder.of(protoClass))
 
   override fun <T1 : Any, T2 : Any> tuple2sOf(first: Encoder<T1>, second: Encoder<T2>) =
     BeamEncoder(
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/Encoders.kt b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/Encoders.kt
index d37234a..84ebd90 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/Encoders.kt
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/core/Encoders.kt
@@ -40,16 +40,56 @@
   /** Returns an [Encoder] for an integer value, which can be stored in a [FrameworkCollection]. */
   fun ints(): Encoder<Int>
 
-  /** Encoder for data classes. */
-  fun <T : Any> records(recordClass: KClass<T>): Encoder<T>
+  /** Encoder for classes. */
+  fun <T : Any> records(recordClass: Class<T>): Encoder<T>
+
+  /** Same as [records(Class)] but accepts Kotlin class. */
+  fun <T : Any> records(recordClass: KClass<T>) = records(recordClass.java)
 
   /** Returns an [Encoder] for a protobuf value, which can be stored in a [FrameworkCollection]. */
-  fun <T : Message> protos(protoClass: KClass<T>): Encoder<T>
+  fun <T : Message> protos(protoClass: Class<T>): Encoder<T>
+
+  /** Same as [protos(Class)] but accepts Kotlin class. */
+  fun <T : Message> protos(protoClass: KClass<T>) = protos(protoClass.java)
 
   /** Returns an [Encoder] for a pair of tuples, which can be stored in a [FrameworkCollection]. */
   fun <T1 : Any, T2 : Any> tuple2sOf(first: Encoder<T1>, second: Encoder<T2>): Encoder<Pair<T1, T2>>
+
+  /**
+   * Returns the most specific [Encoder] for any record type given its [KClass], including primitive
+   * types and proto but except pairs.
+   *
+   * Use it when the record type is not known at compile time and it can be a primitive type or a
+   * proto. This method will return the most appropriate (and efficient) [Encoder] for the given
+   * type.
+   *
+   * Note that this method does not work for pairs ([tuple2sOf]) and for any other classes that are
+   * parameterized by generic types.
+   */
+  fun <T : Any> recordsOfUnknownClass(recordClass: Class<T>) =
+    when {
+      recordClass == String::class.java -> strings()
+      recordClass == Double::class.java -> doubles()
+      recordClass == Int::class.java -> ints()
+      Message::class.java.isAssignableFrom(recordClass) -> {
+        @Suppress("UNCHECKED_CAST") protos(recordClass as Class<out Message>)
+      }
+      else -> records(recordClass)
+    }
+
+  /** Same as [recordsOfUnknownClass(Class)] but accepts Kotlin class. */
+  fun <T : Any> recordsOfUnknownClass(recordClass: KClass<T>) =
+    recordsOfUnknownClass(recordClass.java)
 }
 
+/**
+ * Inlines the function and the type parameter which allows to use [EncoderFactory.records] without
+ * specifying the class.
+ */
 inline fun <reified T : Any> EncoderFactory.records() = this.records(T::class)
 
+/**
+ * Inlines the function and the type parameter which allows to use [EncoderFactory.protos] without
+ * specifying the class.
+ */
 inline fun <reified T : Message> EncoderFactory.protos() = this.protos(T::class)
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/local/LocalEncoderFactory.kt b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/local/LocalEncoderFactory.kt
index 76489cd..3c8098d 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/local/LocalEncoderFactory.kt
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/local/LocalEncoderFactory.kt
@@ -19,7 +19,6 @@
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.Encoder
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.EncoderFactory
 import com.google.protobuf.Message
-import kotlin.reflect.KClass
 
 class LocalEncoderFactory() : EncoderFactory {
   // The implementation of local encoders is empty because when the data is being processed
@@ -36,9 +35,9 @@
     return object : Encoder<Int> {}
   }
 
-  override fun <T : Any> records(recordClass: KClass<T>): Encoder<T> = object : Encoder<T> {}
+  override fun <T : Any> records(recordClass: Class<T>): Encoder<T> = object : Encoder<T> {}
 
-  override fun <T : Message> protos(protoClass: KClass<T>): Encoder<T> = object : Encoder<T> {}
+  override fun <T : Message> protos(protoClass: Class<T>): Encoder<T> = object : Encoder<T> {}
 
   override fun <T1 : Any, T2 : Any> tuple2sOf(first: Encoder<T1>, second: Encoder<T2>) =
     object : Encoder<Pair<T1, T2>> {}
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel
index 8092662..2efe020 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel
@@ -11,32 +11,34 @@
     srcs = ["SparkEncoders.kt"],
     deps = [
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core:encoders",
-        "@maven//:com_google_protobuf_protobuf_java",
-        "@maven//:org_apache_spark_spark_core_2_13",
-        "@maven//:org_apache_spark_spark_sql_2_13",
-        "@maven//:org_apache_spark_spark_mllib_2_13",
-        "@maven//:org_apache_spark_spark_catalyst_2_13",
         "@maven//:com_fasterxml_jackson_core_jackson_databind",
         "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
+        "@maven//:com_google_protobuf_protobuf_java",
+        "@maven//:org_apache_spark_spark_catalyst_2_13",
+        "@maven//:org_apache_spark_spark_core_2_13",
+        "@maven//:org_apache_spark_spark_mllib_2_13",
+        "@maven//:org_apache_spark_spark_sql_2_13",
         "@maven//:org_scala_lang_scala_library",
     ],
 )
 
 kt_jvm_library(
     name = "spark_collections",
-    srcs = ["SparkCollection.kt",
-            "SparkTable.kt",],
+    srcs = [
+        "SparkCollection.kt",
+        "SparkTable.kt",
+    ],
     deps = [
         ":spark_encoders",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core:encoders",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core:framework_collections",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/local:local_collections",
-         "@maven//:org_apache_spark_spark_core_2_13",
-         "@maven//:org_apache_spark_spark_sql_2_13",
-         "@maven//:org_apache_spark_spark_mllib_2_13",
-         "@maven//:org_apache_spark_spark_catalyst_2_13",
-         "@maven//:com_fasterxml_jackson_core_jackson_databind",
-         "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
+        "@maven//:com_fasterxml_jackson_core_jackson_databind",
+        "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
+        "@maven//:org_apache_spark_spark_catalyst_2_13",
+        "@maven//:org_apache_spark_spark_core_2_13",
+        "@maven//:org_apache_spark_spark_mllib_2_13",
+        "@maven//:org_apache_spark_spark_sql_2_13",
     ],
 )
 
diff --git a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncoders.kt b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncoders.kt
index 226cb3c..5d95010 100644
--- a/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncoders.kt
+++ b/pipelinedp4j/main/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncoders.kt
@@ -19,7 +19,6 @@
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.Encoder
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.EncoderFactory
 import com.google.protobuf.Message
-import kotlin.reflect.KClass
 import org.apache.spark.sql.Encoders
 
 /** A serializer and a deserializer for the data types to convert into Spark internal data types. */
@@ -39,12 +38,12 @@
     return SparkEncoder<Int>(Encoders.INT())
   }
 
-  override fun <T : Any> records(recordClass: KClass<T>): SparkEncoder<T> {
-    return SparkEncoder(Encoders.bean(recordClass.java))
+  override fun <T : Any> records(recordClass: Class<T>): SparkEncoder<T> {
+    return SparkEncoder(Encoders.bean(recordClass))
   }
 
-  override fun <T : Message> protos(protoClass: KClass<T>): SparkEncoder<T> {
-    return SparkEncoder<T>(Encoders.kryo(protoClass.java))
+  override fun <T : Message> protos(protoClass: Class<T>): SparkEncoder<T> {
+    return SparkEncoder<T>(Encoders.kryo(protoClass))
   }
 
   override fun <T1 : Any, T2 : Any> tuple2sOf(
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel
index 692eea9..edba452 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/BUILD.bazel
@@ -21,7 +21,7 @@
         "BeamQueryBuilderTest.kt",
         "BeamQueryTest.kt",
         "SparkQueryBuilderTest.kt",
-        "SparkQueryTest.kt"
+        "SparkQueryTest.kt",
     ],
     test_class = "com.google.privacy.differentialprivacy.pipelinedp4j.api.ApiTests",
     runtime_deps = [
@@ -32,15 +32,15 @@
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core:dp_functions_params",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/core/budget:budget_spec",
         "//tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark:spark_tests",
+        "@maven//:com_fasterxml_jackson_core_jackson_databind",
+        "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
         "@maven//:org_apache_beam_beam_sdks_java_core",
-        "@maven//:org_jetbrains_kotlin_kotlin_test",
-        "@maven//:org_apache_spark_spark_core_2_13",
-        "@maven//:org_apache_spark_spark_sql_2_13",
-        "@maven//:org_apache_spark_spark_mllib_2_13",
         "@maven//:org_apache_spark_spark_catalyst_2_13",
-        "@maven//:com_fasterxml_jackson_core_jackson_databind",
-        "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
+        "@maven//:org_apache_spark_spark_core_2_13",
+        "@maven//:org_apache_spark_spark_mllib_2_13",
+        "@maven//:org_apache_spark_spark_sql_2_13",
+        "@maven//:org_jetbrains_kotlin_kotlin_test",
     ],
 )
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryBuilderTest.kt b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryBuilderTest.kt
index e3e046b..5902b8e 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryBuilderTest.kt
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryBuilderTest.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.google.privacy.differentialprivacy.pipelinedp4j.api
 
 import com.google.common.truth.Truth.assertThat
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryTest.kt b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryTest.kt
index 2d5b5c0..bffdb88 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryTest.kt
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/api/SparkQueryTest.kt
@@ -1,3 +1,19 @@
+/*
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.google.privacy.differentialprivacy.pipelinedp4j.api
 
 import com.google.common.truth.Truth.assertThat
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BUILD.bazel b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BUILD.bazel
index b3734a9..28c0e27 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BUILD.bazel
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BUILD.bazel
@@ -37,6 +37,8 @@
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
         "@maven//:org_apache_beam_beam_sdks_java_core",
+        "@maven//:org_apache_beam_beam_sdks_java_extensions_avro",
+        "@maven//:org_apache_beam_beam_sdks_java_extensions_protobuf",
         "@maven//:org_hamcrest_hamcrest",
     ],
 )
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncodersTest.kt b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncodersTest.kt
index c76201d..7a1b9b2 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncodersTest.kt
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/beam/BeamEncodersTest.kt
@@ -16,6 +16,8 @@
 
 package com.google.privacy.differentialprivacy.pipelinedp4j.beam
 
+import org.apache.beam.sdk.extensions.protobuf.ProtoCoder
+import com.google.common.truth.Truth.assertThat
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.ContributionWithPrivacyId
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.contributionWithPrivacyId
 import com.google.privacy.differentialprivacy.pipelinedp4j.core.encoderOfContributionWithPrivacyId
@@ -25,6 +27,10 @@
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.quantilesAccumulator
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.sumAccumulator
 import com.google.protobuf.ByteString
+import org.apache.beam.sdk.coders.DoubleCoder
+import org.apache.beam.sdk.coders.StringUtf8Coder
+import org.apache.beam.sdk.coders.VarIntCoder
+import org.apache.beam.sdk.extensions.avro.coders.AvroCoder
 import org.apache.beam.sdk.testing.PAssert
 import org.apache.beam.sdk.testing.TestPipeline
 import org.apache.beam.sdk.transforms.Create
@@ -75,21 +81,9 @@
 
   @Test
   fun records_isPossibleToCreateBeamPCollectionOfThatType() {
-    val input =
-      listOf(
-        contributionWithPrivacyId("privacyId1", "partitionKey1", -1.0),
-        contributionWithPrivacyId("privacyId2", "partitionKey1", 0.0),
-        contributionWithPrivacyId("privacyId1", "partitionKey2", 1.0),
-        contributionWithPrivacyId("privacyId3", "partitionKey3", 1.2345),
-      )
+    val input = listOf(TestRecord("privacyId1", 1.0, -1), TestRecord("privacyId2", 2.0, 2))
     val inputCoder =
-      (encoderOfContributionWithPrivacyId(
-          beamEncoderFactory.strings(),
-          beamEncoderFactory.strings(),
-          beamEncoderFactory,
-        )
-          as BeamEncoder<ContributionWithPrivacyId<String, String>>)
-        .coder
+      (beamEncoderFactory.records(TestRecord::class) as BeamEncoder<TestRecord>).coder
 
     val pCollection = testPipeline.apply(Create.of(input).withCoder(inputCoder))
 
@@ -115,7 +109,9 @@
         },
         compoundAccumulator {},
       )
-    val inputCoder = beamEncoderFactory.protos(CompoundAccumulator::class).coder
+    val inputCoder =
+      (beamEncoderFactory.protos(CompoundAccumulator::class) as BeamEncoder<CompoundAccumulator>)
+        .coder
 
     val pCollection = testPipeline.apply(Create.of(input).withCoder(inputCoder))
 
@@ -137,6 +133,79 @@
     testPipeline.run().waitUntilFinish()
   }
 
+  @Test
+  fun contributionWithPrivacyIdOf_isPossibleToCreateBeamPCollectionOfThatType() {
+    val input =
+      listOf(
+        contributionWithPrivacyId("privacyId1", "partitionKey1", -1.0),
+        contributionWithPrivacyId("privacyId2", "partitionKey1", 0.0),
+        contributionWithPrivacyId("privacyId1", "partitionKey2", 1.0),
+        contributionWithPrivacyId("privacyId3", "partitionKey3", 1.2345),
+      )
+    val inputCoder =
+      (encoderOfContributionWithPrivacyId(
+          beamEncoderFactory.strings(),
+          beamEncoderFactory.strings(),
+          beamEncoderFactory,
+        )
+          as BeamEncoder<ContributionWithPrivacyId<String, String>>)
+        .coder
+
+    val pCollection = testPipeline.apply(Create.of(input).withCoder(inputCoder))
+
+    PAssert.that(pCollection).containsInAnyOrder(input)
+
+    testPipeline.run().waitUntilFinish()
+  }
+
+  @Test
+  fun recordsOfUnknownClass_string_createsEncoderWithStringCoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder = beamEncoderFactory.recordsOfUnknownClass(String::class) as BeamEncoder<String>
+
+    assertThat(encoder.coder).isInstanceOf(StringUtf8Coder::class.java)
+  }
+
+  @Test
+  fun recordsOfUnknownClass_double_createsEncoderWithDoubleCoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder = beamEncoderFactory.recordsOfUnknownClass(Double::class) as BeamEncoder<Double>
+
+    assertThat(encoder.coder).isInstanceOf(DoubleCoder::class.java)
+  }
+
+  @Test
+  fun recordsOfUnknownClass_int_createsEncoderWithIntCoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder = beamEncoderFactory.recordsOfUnknownClass(Int::class) as BeamEncoder<Int>
+
+    assertThat(encoder.coder).isInstanceOf(VarIntCoder::class.java)
+  }
+
+  @Test
+  fun recordsOfUnknownClass_kotlinClass_createsEncoderWithAvroCoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder =
+      beamEncoderFactory.recordsOfUnknownClass(TestRecord::class) as BeamEncoder<TestRecord>
+
+    assertThat(encoder.coder).isInstanceOf(AvroCoder::class.java)
+  }
+
+  @Test
+  fun recordsOfUnknownClass_proto_createsEncoderWithProtoCoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder =
+      beamEncoderFactory.recordsOfUnknownClass(CompoundAccumulator::class)
+        as BeamEncoder<CompoundAccumulator>
+
+    assertThat(encoder.coder).isInstanceOf(ProtoCoder::class.java)
+  }
+
+  private data class TestRecord(val string: String, val double: Double, val int: Int) {
+    // Required for Beam serialization.
+    private constructor() : this("", 0.0, 0)
+  }
+
   companion object {
     private val beamEncoderFactory = BeamEncoderFactory()
   }
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel
index 804cd08..c1cef1f 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/BUILD.bazel
@@ -31,15 +31,15 @@
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/proto:accumulators_kt_proto",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/spark:spark_collections",
         "//main/com/google/privacy/differentialprivacy/pipelinedp4j/spark:spark_encoders",
+        "@maven//:com_fasterxml_jackson_core_jackson_databind",
+        "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
         "@maven//:com_google_protobuf_protobuf_java",
         "@maven//:com_google_testparameterinjector_test_parameter_injector",
         "@maven//:com_google_truth_truth",
         "@maven//:junit_junit",
-        "@maven//:org_apache_spark_spark_core_2_13",
-        "@maven//:org_apache_spark_spark_sql_2_13",
-        "@maven//:org_apache_spark_spark_mllib_2_13",
         "@maven//:org_apache_spark_spark_catalyst_2_13",
-        "@maven//:com_fasterxml_jackson_core_jackson_databind",
-        "@maven//:com_fasterxml_jackson_module_jackson_module_paranamer",
+        "@maven//:org_apache_spark_spark_core_2_13",
+        "@maven//:org_apache_spark_spark_mllib_2_13",
+        "@maven//:org_apache_spark_spark_sql_2_13",
     ],
 )
diff --git a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncodersTest.kt b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncodersTest.kt
index 420008e..1267e62 100644
--- a/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncodersTest.kt
+++ b/pipelinedp4j/tests/com/google/privacy/differentialprivacy/pipelinedp4j/spark/SparkEncodersTest.kt
@@ -26,6 +26,7 @@
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.quantilesAccumulator
 import com.google.privacy.differentialprivacy.pipelinedp4j.proto.sumAccumulator
 import com.google.protobuf.ByteString
+import org.apache.spark.sql.Encoders
 import org.junit.ClassRule
 import org.junit.Test
 import org.junit.runner.RunWith
@@ -38,7 +39,9 @@
   fun strings_isPossibleToCreateSparkCollectionOfThatType() {
     val input = listOf("a", "b", "c")
     val inputCoder = sparkEncoderFactory.strings().encoder
+
     val dataset = sparkSession.spark.createDataset(input, inputCoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
@@ -46,7 +49,9 @@
   fun doubles_isPossibleToCreateSparkCollectionOfThatType() {
     val input = listOf(-1.2, 0.0, 2.1)
     val inputCoder = sparkEncoderFactory.doubles().encoder
+
     val dataset = sparkSession.spark.createDataset(input, inputCoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
@@ -54,28 +59,20 @@
   fun ints_isPossibleToCreateSparkCollectionOfThatType() {
     val input = listOf(-1, 0, 1)
     val inputCoder = sparkEncoderFactory.ints().encoder
+
     val dataset = sparkSession.spark.createDataset(input, inputCoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
   @Test
   fun records_isPossibleToCreateSparkCollectionOfThatType() {
-    val input =
-      listOf(
-        contributionWithPrivacyId("privacyId1", "partitionKey1", -1.0),
-        contributionWithPrivacyId("privacyId2", "partitionKey1", 0.0),
-        contributionWithPrivacyId("privacyId1", "partitionKey2", 1.0),
-        contributionWithPrivacyId("privacyId3", "partitionKey3", 1.2345),
-      )
+    val input = listOf(TestRecord("privacyId1", 1.0, -1), TestRecord("privacyId2", 2.0, 2))
     val inputCoder =
-      (encoderOfContributionWithPrivacyId(
-          sparkEncoderFactory.strings(),
-          sparkEncoderFactory.strings(),
-          sparkEncoderFactory,
-        )
-          as SparkEncoder<ContributionWithPrivacyId<String, String>>)
-        .encoder
+      (sparkEncoderFactory.records(TestRecord::class) as SparkEncoder<TestRecord>).encoder
+
     val dataset = sparkSession.spark.createDataset(input, inputCoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
@@ -96,8 +93,12 @@
         },
         compoundAccumulator {},
       )
-    val inputCoder = sparkEncoderFactory.protos(CompoundAccumulator::class).encoder
+    val inputCoder =
+      (sparkEncoderFactory.protos(CompoundAccumulator::class) as SparkEncoder<CompoundAccumulator>)
+        .encoder
+
     val dataset = sparkSession.spark.createDataset(input, inputCoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
@@ -110,6 +111,7 @@
         .encoder
 
     val dataset = sparkSession.spark.createDataset(input, inputEncoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
@@ -129,10 +131,80 @@
           sparkEncoderFactory.strings(),
         )
         .encoder
+
     val dataset = sparkSession.spark.createDataset(input, inputEncoder)
+
     assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
   }
 
+  @Test
+  fun contributionWithPrivacyIdOf_isPossibleToCreateSparkCollectionOfThatType() {
+    val input =
+      listOf(
+        contributionWithPrivacyId("privacyId1", "partitionKey1", -1.0),
+        contributionWithPrivacyId("privacyId2", "partitionKey1", 0.0),
+        contributionWithPrivacyId("privacyId1", "partitionKey2", 1.0),
+        contributionWithPrivacyId("privacyId3", "partitionKey3", 1.2345),
+      )
+    val inputCoder =
+      (encoderOfContributionWithPrivacyId(
+          sparkEncoderFactory.strings(),
+          sparkEncoderFactory.strings(),
+          sparkEncoderFactory,
+        )
+          as SparkEncoder<ContributionWithPrivacyId<String, String>>)
+        .encoder
+
+    val dataset = sparkSession.spark.createDataset(input, inputCoder)
+
+    assertThat(dataset.collectAsList()).containsExactlyElementsIn(input)
+  }
+
+  @Test
+  fun recordsOfUnknownClass_string_createsEncoderWithSparkStringEncoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder = sparkEncoderFactory.recordsOfUnknownClass(String::class) as SparkEncoder<String>
+
+    assertThat(encoder.encoder).isEqualTo(Encoders.STRING())
+  }
+
+  @Test
+  fun recordsOfUnknownClass_double_createsEncoderWithSparkDoubleEncoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder = sparkEncoderFactory.recordsOfUnknownClass(Double::class) as SparkEncoder<Double>
+
+    assertThat(encoder.encoder).isEqualTo(Encoders.DOUBLE())
+  }
+
+  @Test
+  fun recordsOfUnknownClass_int_createsEncoderWithSparkIntEncoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder = sparkEncoderFactory.recordsOfUnknownClass(Int::class) as SparkEncoder<Int>
+
+    assertThat(encoder.encoder).isEqualTo(Encoders.INT())
+  }
+
+  @Test
+  fun recordsOfUnknownClass_kotlinClass_createsEncoderWithSparkBeanEncoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder =
+      sparkEncoderFactory.recordsOfUnknownClass(TestRecord::class) as SparkEncoder<TestRecord>
+
+    assertThat(encoder.encoder).isEqualTo(Encoders.bean(TestRecord::class.java))
+  }
+
+  @Test
+  fun recordsOfUnknownClass_proto_createsEncoderWithSparkKryoEncoder() {
+    @Suppress("UNCHECKED_CAST")
+    val encoder =
+      sparkEncoderFactory.recordsOfUnknownClass(CompoundAccumulator::class)
+        as SparkEncoder<CompoundAccumulator>
+
+    assertThat(encoder.encoder).isEqualTo(Encoders.kryo(CompoundAccumulator::class.java))
+  }
+
+  data class TestRecord(var string: String = "", var double: Double = 0.0, var int: Int = 0)
+
   companion object {
     @JvmField @ClassRule val sparkSession = SparkSessionRule()
     private val sparkEncoderFactory = SparkEncoderFactory()
diff --git a/privacy-on-beam/pbeam/public_partitions.go b/privacy-on-beam/pbeam/public_partitions.go
index f95f6ae..145d0ac 100644
--- a/privacy-on-beam/pbeam/public_partitions.go
+++ b/privacy-on-beam/pbeam/public_partitions.go
@@ -22,12 +22,19 @@
 	"bytes"
 	"encoding/base64"
 	"fmt"
+	"math/rand"
 	"reflect"
 
+	"flag"
 	"github.com/google/differential-privacy/privacy-on-beam/v3/internal/kv"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/typex"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/transforms/filter"
+)
+
+var (
+	enableShardedPublicPartitions = flag.Bool("enable_sharded_public_partitions", false, "Enable sharded public partitions. This is a temporary flag to allow us to test the new sharded implementation of public partition filtering.")
 )
 
 func init() {
@@ -51,6 +58,11 @@
 	register.Function4x0[beam.W, func(*beam.V) bool, func(*beam.V) bool, func(beam.W, beam.V)](mergeResultWithEmptyPublicPartitionsFn)
 	register.Iter1[beam.V]()
 	register.Emitter2[beam.W, beam.V]()
+
+	register.Function1x2[kv.Pair, ShardedKey, []byte](addRandomShardIDFn)
+	register.Function4x0[[]byte, func(*[]byte) bool, func(*int) bool, func(ShardedKey, int)](extractEmittableKeysWithShardIDFn)
+	register.DoFn4x1[ShardedKey, func(*int) bool, func(*[]byte) bool, func(beam.T, beam.W), error](&filterKeysWithShardIDFn{})
+	register.Function1x2[kv.Pair, []byte, []byte](unwrapPairFn)
 }
 
 // newAddZeroValuesToPublicPartitionsFn turns a PCollection<V> into PCollection<V,0>.
@@ -138,6 +150,116 @@
 	}
 }
 
+// ShardedKey is an key encoded as bytes with a int shardID.
+type ShardedKey struct {
+	K       []byte
+	ShardID int
+}
+
+func addRandomShardIDFn(encoded kv.Pair) (ShardedKey, []byte) {
+	return ShardedKey{K: encoded.K, ShardID: rand.Intn(2048)}, encoded.V
+}
+
+func extractEmittableKeysWithShardIDFn(k []byte, isAllowedKeyIter func(*[]byte) bool, shardIDIter func(*int) bool, emit func(ShardedKey, int)) {
+	var isAllowedKey []byte
+	if !isAllowedKeyIter(&isAllowedKey) {
+		// k is not an allow listed key, filter it out.
+		return
+	}
+
+	var subkey int
+	for shardIDIter(&subkey) {
+		emit(ShardedKey{K: k, ShardID: subkey}, 0)
+	}
+}
+
+type filterKeysWithShardIDFn struct {
+	KType     beam.EncodedType
+	VType     beam.EncodedType
+	PairCodec *kv.Codec
+}
+
+func newFilterKeysWithShardIDFn(kType, vType reflect.Type) *filterKeysWithShardIDFn {
+	return &filterKeysWithShardIDFn{
+		KType: beam.EncodedType{T: kType},
+		VType: beam.EncodedType{T: vType},
+	}
+}
+
+func (fn *filterKeysWithShardIDFn) Setup() error {
+	fn.PairCodec = kv.NewCodec(fn.KType.T, fn.VType.T)
+	return fn.PairCodec.Setup()
+}
+
+func (fn *filterKeysWithShardIDFn) ProcessElement(k ShardedKey, isEmittableShardIDIter func(*int) bool, pcolValueIter func(*[]byte) bool, emit func(beam.T, beam.W)) error {
+	var isEmittableShardID int
+	if !isEmittableShardIDIter(&isEmittableShardID) {
+		// k isn't a key from the public partitions collection.
+		return nil
+	}
+
+	var pcolValue []byte
+	for pcolValueIter(&pcolValue) {
+		k, v, err := fn.PairCodec.Decode(kv.Pair{K: k.K, V: pcolValue})
+		if err != nil {
+			return err
+		}
+		emit(k, v)
+	}
+
+	return nil
+}
+
+func unwrapPairFn(encoded kv.Pair) ([]byte, []byte) {
+	return encoded.K, encoded.V
+}
+
+func unwrapShardedKeyFn(shardedKey ShardedKey) ([]byte, int) {
+	return shardedKey.K, shardedKey.ShardID
+}
+
+// Filters out KV-s from col that have a key in 'keys'.
+//
+// A single key in col may have a huge number of values. This function handles that
+// case by sharding col randomly into 2048 collections, and then joining the keys
+// within each of those shards. This reduces stragglers when there is a hot key in col,
+// since its processing is parallelized 2048 ways.
+//
+// Each value in col is randomly selected to be in one of the 2048 shards. Then the
+// 'keys' collection is joined with the sharded col collection to find the shardIds
+// that need to be present per key. Then the sharded col collection is joined with
+// the sharded 'keys' collection to compute the final filtered result.
+func filterKeysImbalanced(s beam.Scope, col beam.PCollection, keys beam.PCollection) beam.PCollection {
+	s = s.Scope("filterKeysImbalanced")
+
+	kT, vT := beam.ValidateKVType(col)
+
+	// Add a random shardId (one of 2048) to each element in col.
+	// PCollection<KV<ShardedKey[key, randInt[0;2048)], []byte>>
+	pcolAsBytesWithShardID := beam.ParDo(s, addRandomShardIDFn, beam.ParDo(s, kv.NewEncodeFn(kT, vT), col))
+	// Drop values and remove duplicates.
+	// PCollection<ShardedKey[key, randInt[0;2048)]>
+	uniqueKeysWithShardID := filter.Distinct(s, beam.DropValue(s, pcolAsBytesWithShardID))
+
+	// Prepare the keys for a CoGroupBy with uniqueSubkeys.
+	// PCollection<KV<key, 0>>
+	keysWithZero := beam.ParDo(s, addZeroValuesToPublicPartitionsInt64, keys)
+	// PCollection<kv.Pair<key, 0>>
+	keysEncodedWithZero := beam.ParDo(
+		s,
+		kv.NewEncodeFn(keysWithZero.Type().Components()[0],
+			keysWithZero.Type().Components()[1]), keysWithZero)
+
+	// Find the shardIds per key in the keys collection.
+	// PCollection<KV<key, *>>
+	groupedByKey := beam.CoGroupByKey(s, beam.ParDo(s, unwrapPairFn, keysEncodedWithZero), beam.ParDo(s, unwrapShardedKeyFn, uniqueKeysWithShardID))
+	emittableKeysWithShardID := beam.ParDo(s, extractEmittableKeysWithShardIDFn, groupedByKey)
+
+	// Finally perform the sharded filter.
+	groupedByKeyAndShardID := beam.CoGroupByKey(s, emittableKeysWithShardID, pcolAsBytesWithShardID)
+	return beam.ParDo(s, newFilterKeysWithShardIDFn(kT.Type(), vT.Type()), groupedByKeyAndShardID, beam.TypeDefinition{Var: beam.TType, T: kT.Type()}, beam.TypeDefinition{Var: beam.WType, T: vT.Type()})
+}
+
 // dropNonPublicPartitionsVFn drops partitions not specified in
 // PublicPartitions from pcol. It can be used for aggregations on V values,
 // e.g. Count and DistinctPrivacyID.
@@ -152,9 +274,13 @@
 // Returns a PCollection<PrivacyKey, Value> only for values present in
 // publicPartitions.
 func dropNonPublicPartitionsVFn(s beam.Scope, publicPartitions beam.PCollection, pcol PrivatePCollection) beam.PCollection {
-	publicPartitionsWithZeros := beam.ParDo(s, addZeroValuesToPublicPartitionsInt64, publicPartitions)
-	groupedByValue := beam.CoGroupByKey(s, publicPartitionsWithZeros, beam.SwapKV(s, pcol.col))
-	return beam.ParDo(s, mergePublicValues, groupedByValue)
+	if *enableShardedPublicPartitions {
+		return beam.SwapKV(s, filterKeysImbalanced(s, beam.SwapKV(s, pcol.col), publicPartitions))
+	} else {
+		publicPartitionsWithZeros := beam.ParDo(s, addZeroValuesToPublicPartitionsInt64, publicPartitions)
+		groupedByValue := beam.CoGroupByKey(s, publicPartitionsWithZeros, beam.SwapKV(s, pcol.col))
+		return beam.ParDo(s, mergePublicValues, groupedByValue)
+	}
 }
 
 // dropNonPublicPartitionsKVFn drops partitions not specified in
@@ -175,12 +301,20 @@
 // Returns a PCollection<PrivacyKey, <PartitionKey, Value>> only for values present in
 // publicPartitions.
 func dropNonPublicPartitionsKVFn(s beam.Scope, publicPartitions beam.PCollection, pcol PrivatePCollection, idType typex.FullType) beam.PCollection {
-	publicPartitionsWithZeros := beam.ParDo(s, addZeroValuesToPublicPartitionsInt64, publicPartitions)
-	encodedIDV := beam.ParDo(s, newEncodeIDVFn(idType, pcol.codec), pcol.col, beam.TypeDefinition{Var: beam.WType, T: pcol.codec.KType.T})
-	groupedByValue := beam.CoGroupByKey(s, publicPartitionsWithZeros, encodedIDV)
-	merged := beam.SwapKV(s, beam.ParDo(s, mergePublicValues, groupedByValue))
-	decodeFn := newDecodeIDVFn(pcol.codec.KType, kv.NewCodec(idType.Type(), pcol.codec.VType.T))
-	return beam.ParDo(s, decodeFn, merged, beam.TypeDefinition{Var: beam.UType, T: idType.Type()})
+	if *enableShardedPublicPartitions {
+		encodedIDV := beam.ParDo(
+			s, newEncodeIDVFn(idType, pcol.codec), pcol.col, beam.TypeDefinition{Var: beam.WType, T: pcol.codec.KType.T})
+		filteredEncodedIDV := filterKeysImbalanced(s, encodedIDV, publicPartitions)
+		decodeFn := newDecodeIDVFn(pcol.codec.KType, kv.NewCodec(idType.Type(), pcol.codec.VType.T))
+		return beam.ParDo(s, decodeFn, filteredEncodedIDV, beam.TypeDefinition{Var: beam.UType, T: idType.Type()})
+	} else {
+		publicPartitionsWithZeros := beam.ParDo(s, addZeroValuesToPublicPartitionsInt64, publicPartitions)
+		encodedIDV := beam.ParDo(s, newEncodeIDVFn(idType, pcol.codec), pcol.col, beam.TypeDefinition{Var: beam.WType, T: pcol.codec.KType.T})
+		groupedByValue := beam.CoGroupByKey(s, publicPartitionsWithZeros, encodedIDV)
+		merged := beam.SwapKV(s, beam.ParDo(s, mergePublicValues, groupedByValue))
+		decodeFn := newDecodeIDVFn(pcol.codec.KType, kv.NewCodec(idType.Type(), pcol.codec.VType.T))
+		return beam.ParDo(s, decodeFn, merged, beam.TypeDefinition{Var: beam.UType, T: idType.Type()})
+	}
 }
 
 // encodeIDVFn takes a PCollection<ID,kv.Pair{K,V}> as input, and returns a
diff --git a/privacy-on-beam/pbeam/public_partitions_test.go b/privacy-on-beam/pbeam/public_partitions_test.go
index 54e83e1..1ed8a00 100644
--- a/privacy-on-beam/pbeam/public_partitions_test.go
+++ b/privacy-on-beam/pbeam/public_partitions_test.go
@@ -22,6 +22,7 @@
 	"reflect"
 	"testing"
 
+	"flag"
 	"github.com/google/differential-privacy/privacy-on-beam/v3/pbeam/testutils"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam"
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
@@ -104,6 +105,40 @@
 	}
 }
 
+// TODO: Remove once the enable_sharded_public_partitions flag is gone.
+func TestDropNonPublicPartitionsVFnShardedImpl(t *testing.T) {
+	flag.Set("enable_sharded_public_partitions", "true")
+
+	pairs := testutils.ConcatenatePairs(
+		testutils.MakePairsWithFixedV(7, 0),
+		testutils.MakePairsWithFixedVStartingFromKey(7, 10, 1),
+		testutils.MakePairsWithFixedVStartingFromKey(17, 83, 2),
+		testutils.MakePairsWithFixedVStartingFromKey(100, 10, 3),
+	)
+
+	// Keep partitions 0, 2;
+	// drop partitions 1, 3.
+	result := testutils.ConcatenatePairs(
+		testutils.MakePairsWithFixedV(7, 0),
+		testutils.MakePairsWithFixedVStartingFromKey(17, 83, 2),
+	)
+
+	p, s, col, want := ptest.CreateList2(pairs, result)
+	want = beam.ParDo(s, testutils.PairToKV, want)
+	col = beam.ParDo(s, testutils.PairToKV, col)
+	partitions := []int{0, 2}
+
+	partitionsCol := beam.CreateList(s, partitions)
+	epsilon := 50.0
+	pcol := MakePrivate(s, col, privacySpec(t, PrivacySpecParams{AggregationEpsilon: epsilon}))
+	got := dropNonPublicPartitionsVFn(s, partitionsCol, pcol)
+	testutils.EqualsKVInt(t, s, got, want)
+	if err := ptest.Run(p); err != nil {
+		t.Errorf("DropNonPublicPartitionsVFn did not drop non public partitions as expected: %v", err)
+	}
+	flag.Set("enable_sharded_public_partitions", "false")
+}
+
 // TestDropNonPublicPartitionsKVFn checks that int elements with non-public partitions
 // are dropped (tests function used for sum and mean).
 func TestDropNonPublicPartitionsKVFn(t *testing.T) {
@@ -147,6 +182,51 @@
 	}
 }
 
+// TODO: Remove once the enable_sharded_public_partitions flag is gone.
+func TestDropNonPublicPartitionsKVFnShardedImpl(t *testing.T) {
+	flag.Set("enable_sharded_public_partitions", "true")
+
+	triples := testutils.ConcatenateTriplesWithIntValue(
+		testutils.MakeTripleWithIntValueStartingFromKey(0, 7, 0, 0),
+		testutils.MakeTripleWithIntValueStartingFromKey(7, 3, 1, 0),
+		testutils.MakeTripleWithIntValueStartingFromKey(10, 90, 2, 0),
+		testutils.MakeTripleWithIntValueStartingFromKey(100, 100, 11, 0),
+		testutils.MakeTripleWithIntValueStartingFromKey(200, 5, 12, 0))
+	// Keep partitions 0, 2.
+	// Drop partitions 1, 33, 100.
+	result := testutils.ConcatenateTriplesWithIntValue(
+		testutils.MakeTripleWithIntValueStartingFromKey(0, 7, 0, 0),
+		testutils.MakeTripleWithIntValueStartingFromKey(10, 90, 2, 0))
+
+	p, s, col, col2 := ptest.CreateList2(triples, result)
+	// Doesn't matter that the values 3, 4, 5, 6, 9, 10
+	// are in the partitions PCollection because we are
+	// just dropping the values that are in our original PCollection
+	// that are not in public partitions.
+	partitionsCol := beam.CreateList(s, []int{0, 2, 3, 4, 5, 6, 9, 10})
+	col = beam.ParDo(s, testutils.ExtractIDFromTripleWithIntValue, col)
+	col2 = beam.ParDo(s, testutils.ExtractIDFromTripleWithIntValue, col2)
+	epsilon := 50.0
+
+	pcol := MakePrivate(s, col, privacySpec(t, PrivacySpecParams{AggregationEpsilon: epsilon}))
+	pcol = ParDo(s, testutils.TripleWithIntValueToKV, pcol)
+	idT, _ := beam.ValidateKVType(pcol.col)
+
+	got := dropNonPublicPartitionsKVFn(s, partitionsCol, pcol, idT)
+	got = beam.SwapKV(s, got)
+
+	pcol2 := MakePrivate(s, col2, privacySpec(t, PrivacySpecParams{AggregationEpsilon: epsilon}))
+	pcol2 = ParDo(s, testutils.TripleWithIntValueToKV, pcol2)
+	want := pcol2.col
+	want = beam.SwapKV(s, want)
+
+	testutils.EqualsKVInt(t, s, got, want)
+	if err := ptest.Run(p); err != nil {
+		t.Errorf("TestDropNonPublicPartitionsKVFn did not drop non public partitions as expected: %v", err)
+	}
+	flag.Set("enable_sharded_public_partitions", "false")
+}
+
 // Check that float elements with non-public partitions
 // are dropped (tests function used for sum and mean).
 func TestDropNonPublicPartitionsFloat(t *testing.T) {
@@ -190,3 +270,49 @@
 		t.Errorf("TestDropNonPublicPartitionsFloat did not drop non public partitions as expected: %v", err)
 	}
 }
+
+// TODO: Remove once the enable_sharded_public_partitions flag is gone.
+func TestDropNonPublicPartitionsFloatShardedImpl(t *testing.T) {
+	flag.Set("enable_sharded_public_partitions", "true")
+
+	// In this test, we check  that non-public partitions
+	// are dropped. This function is used for sum and mean.
+	// Used example values from the mean test.
+	triples := testutils.ConcatenateTriplesWithFloatValue(
+		testutils.MakeTripleWithFloatValue(7, 0, 2.0),
+		testutils.MakeTripleWithFloatValueStartingFromKey(7, 100, 1, 1.3),
+		testutils.MakeTripleWithFloatValueStartingFromKey(107, 150, 1, 2.5),
+	)
+	// Keep partition 0.
+	// drop partition 1.
+	result := testutils.ConcatenateTriplesWithFloatValue(
+		testutils.MakeTripleWithFloatValue(7, 0, 2.0))
+
+	p, s, col, col2 := ptest.CreateList2(triples, result)
+
+	// Doesn't matter that the values 2, 3, 4, 5, 6, 7 are in the partitions PCollection.
+	// We are just dropping the values that are in our original PCollection that are not in
+	// public partitions.
+	partitionsCol := beam.CreateList(s, []int{0, 2, 3, 4, 5, 6, 7})
+	col = beam.ParDo(s, testutils.ExtractIDFromTripleWithFloatValue, col)
+	col2 = beam.ParDo(s, testutils.ExtractIDFromTripleWithFloatValue, col2)
+	epsilon := 50.0
+
+	pcol := MakePrivate(s, col, privacySpec(t, PrivacySpecParams{AggregationEpsilon: epsilon}))
+	pcol = ParDo(s, testutils.TripleWithFloatValueToKV, pcol)
+	idT, _ := beam.ValidateKVType(pcol.col)
+
+	got := dropNonPublicPartitionsKVFn(s, partitionsCol, pcol, idT)
+	got = beam.SwapKV(s, got)
+
+	pcol2 := MakePrivate(s, col2, privacySpec(t, PrivacySpecParams{AggregationEpsilon: epsilon}))
+	pcol2 = ParDo(s, testutils.TripleWithFloatValueToKV, pcol2)
+	want := pcol2.col
+	want = beam.SwapKV(s, want)
+
+	testutils.EqualsKVInt(t, s, got, want)
+	if err := ptest.Run(p); err != nil {
+		t.Errorf("TestDropNonPublicPartitionsFloat did not drop non public partitions as expected: %v", err)
+	}
+	flag.Set("enable_sharded_public_partitions", "false")
+}
diff --git a/privacy-on-beam/pbeam/testutils/testutils.go b/privacy-on-beam/pbeam/testutils/testutils.go
index 04851e6..8151f2d 100644
--- a/privacy-on-beam/pbeam/testutils/testutils.go
+++ b/privacy-on-beam/pbeam/testutils/testutils.go
@@ -24,6 +24,7 @@
 	"math"
 	"math/big"
 	"reflect"
+	"sort"
 	"testing"
 
 	"github.com/google/differential-privacy/go/v3/dpagg"
@@ -510,6 +511,7 @@
 	for vIter(&v) {
 		vSlice = append(vSlice, float64(v))
 	}
+	sort.Float64s(vSlice)
 	return vSlice
 }