Auto Bucket Partitioner

src/integrationTest/java/com/mongodb/spark/sql/connector/mongodb/MongoSparkConnectorHelper.java

@@ -58,7 +58,10 @@ public class MongoSparkConnectorHelper
   public static final String URI_SYSTEM_PROPERTY_NAME = "org.mongodb.test.uri";
   public static final String DEFAULT_DATABASE_NAME = "MongoSparkConnectorTest";
   public static final String DEFAULT_COLLECTION_NAME = "coll";
-  private static final String SAMPLE_DATA_TEMPLATE = "{_id: '%s', pk: '%s', dups: '%s', s: '%s'}";
+  private static final String SAMPLE_DATA_TEMPLATE =
+      "{_id: '%s', pk: '%s', dups: '%s', i: %d, s: '%s'}";
+  private static final String COMPLEX_SAMPLE_DATA_TEMPLATE =
+      "{_id: '%s', nested: {pk: '%s', dups: '%s', i: %d}, s: '%s'}";
 
   private static final Logger LOGGER = LoggerFactory.getLogger(MongoSparkConnectorHelper.class);
 
@@ -194,13 +197,33 @@ public boolean isOnline() {
    */
   public void loadSampleData(
       final int numberOfDocuments, final int sizeInMB, final MongoConfig config) {
+    loadSampleData(numberOfDocuments, sizeInMB, config, SAMPLE_DATA_TEMPLATE);
+  }
+
+  /**
+   * Creates complex sample data
+   *
+   * @param numberOfDocuments the total number of documents to create
+   * @param sizeInMB the total size of the documents
+   * @param config the config used for the database and collection
+   */
+  public void loadComplexSampleData(
+      final int numberOfDocuments, final int sizeInMB, final MongoConfig config) {
+    loadSampleData(numberOfDocuments, sizeInMB, config, COMPLEX_SAMPLE_DATA_TEMPLATE);
+  }
+
+  private void loadSampleData(
+      final int numberOfDocuments,
+      final int sizeInMB,
+      final MongoConfig config,
+      final String template) {
     if (!isOnline()) {
       return;
     }
     int sizeBytes = sizeInMB * 1000 * 1000;
     int totalDocumentSize = sizeBytes / numberOfDocuments;
     int sampleDataWithEmptySampleStringSize = RawBsonDocument.parse(
-            format(SAMPLE_DATA_TEMPLATE, "00000", "_10000", "00000", ""))
+            format(template, "00000", "_10000", "00000", 1, ""))
         .getByteBuffer()
         .limit();
     String sampleString =
@@ -213,7 +236,7 @@ public void loadSampleData(
           String pkString = format("_%s", i + 10000);
           String dupsString = StringUtils.leftPad(format("%s", i % 3 == 0 ? 0 : i), 5, "0");
           return RawBsonDocument.parse(
-              format(SAMPLE_DATA_TEMPLATE, idString, pkString, dupsString, sampleString));
+              format(template, idString, pkString, dupsString, i % 10, sampleString));
         })
         .collect(Collectors.toList());
     MongoCollection<BsonDocument> coll = getMongoClient()

src/integrationTest/java/com/mongodb/spark/sql/connector/mongodb/MongoSparkConnectorTestCase.java

@@ -155,6 +155,18 @@ public void loadSampleData(
     HELPER.loadSampleData(numberOfDocuments, sizeInMB, config);
   }
 
+  /**
+   * Creates complex sample data
+   *
+   * @param numberOfDocuments the total number of documents to create
+   * @param sizeInMB the total size of the documents
+   * @param config the config used for the database and collection
+   */
+  public void loadComplexSampleData(
+      final int numberOfDocuments, final int sizeInMB, final MongoConfig config) {
+    HELPER.loadComplexSampleData(numberOfDocuments, sizeInMB, config);
+  }
+
   /** Runs events with the profiler on. */
   public void assertCommentsInProfile(final Runnable runnable, final ReadConfig readConfig) {
     assumeFalse(isSharded());