From 47faaf41e7a84745a48292287a0ae1155340621e Mon Sep 17 00:00:00 2001 From: SuperSaiyan101 Date: Mon, 26 Oct 2026 23:35:27 +0530 Subject: [PATCH 1/2] [#18690] Port SolrIOTest off deprecated DoFnTester --- .../apache/beam/sdk/io/solr/SolrIOTest.java | 43 ++++--------------- 1 file changed, 9 insertions(+), 34 deletions(-) diff --git a/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java b/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java index bdc76824f5be..ab081524dbfe 100644 --- a/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java +++ b/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java @@ -34,7 +34,6 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.Create; -import org.apache.beam.sdk.transforms.DoFnTester; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.BaseEncoding; import org.apache.solr.SolrTestCaseJ4; @@ -207,44 +206,20 @@ public void testWrite() throws Exception { @Test public void testWriteWithMaxBatchSize() throws Exception { + List data = SolrIOTestUtils.createDocuments(NUM_DOCS); SolrIO.Write write = SolrIO.write() .withConnectionConfiguration(connectionConfiguration) .to(SOLR_COLLECTION) .withMaxBatchSize(BATCH_SIZE); - // write bundles size is the runner decision, we cannot force a bundle size, - // so we test the Writer as a DoFn outside of a runner. - try (DoFnTester fnTester = - DoFnTester.of(new SolrIO.Write.WriteFn(write))) { - List input = SolrIOTestUtils.createDocuments(NUM_DOCS); - long numDocsProcessed = 0; - long numDocsInserted = 0; - for (SolrInputDocument document : input) { - fnTester.processElement(document); - numDocsProcessed++; - // test every 100 docs to avoid overloading Solr - if ((numDocsProcessed % 100) == 0) { - // force the index to upgrade after inserting for the inserted docs - // to be searchable immediately - long currentNumDocs = - SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient); - if ((numDocsProcessed % BATCH_SIZE) == 0) { - /* bundle end */ - assertEquals( - "we are at the end of a bundle, we should have inserted all processed documents", - numDocsProcessed, - currentNumDocs); - numDocsInserted = currentNumDocs; - } else { - /* not bundle end */ - assertEquals( - "we are not at the end of a bundle, we should have inserted no more documents", - numDocsInserted, - currentNumDocs); - } - } - } - } + pipeline.apply(Create.of(data)).apply(write); + pipeline.run(); + + long currentNumDocs = SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient); + assertEquals(NUM_DOCS, currentNumDocs); + + QueryResponse response = solrClient.query(SOLR_COLLECTION, new SolrQuery("scientist:Lovelace")); + assertEquals(NUM_DOCS / NUM_SCIENTISTS, response.getResults().getNumFound()); } /** From 01d4dca13b7ef81cd26f470519dc1025147b475c Mon Sep 17 00:00:00 2001 From: SuperSaiyan101 Date: Mon, 26 Oct 2026 23:50:35 +0530 Subject: [PATCH 2/2] Address review feedback: force single bundle and ensure mid-bundle flush execution --- .../apache/beam/sdk/io/solr/SolrIOTest.java | 30 +++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java b/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java index ab081524dbfe..773b4485a8cd 100644 --- a/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java +++ b/sdks/java/io/solr/src/test/java/org/apache/beam/sdk/io/solr/SolrIOTest.java @@ -34,6 +34,11 @@ import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.GroupByKey; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.Values; +import org.apache.beam.sdk.transforms.WithKeys; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.io.BaseEncoding; import org.apache.solr.SolrTestCaseJ4; @@ -204,15 +209,36 @@ public void testWrite() throws Exception { assertEquals(NUM_DOCS / NUM_SCIENTISTS, response.getResults().getNumFound()); } + private static class EmitDocumentsFn + extends DoFn, SolrInputDocument> { + @ProcessElement + public void processElement(ProcessContext c) { + for (SolrInputDocument doc : c.element()) { + c.output(doc); + } + } + } + @Test public void testWriteWithMaxBatchSize() throws Exception { + int maxBatchSize = 1; List data = SolrIOTestUtils.createDocuments(NUM_DOCS); SolrIO.Write write = SolrIO.write() .withConnectionConfiguration(connectionConfiguration) .to(SOLR_COLLECTION) - .withMaxBatchSize(BATCH_SIZE); - pipeline.apply(Create.of(data)).apply(write); + .withMaxBatchSize(maxBatchSize); + + // Force all elements into a single bundle via GroupByKey and use a maxBatchSize of 1 + // to guarantee that batch.size() >= spec.getMaxBatchSize() is triggered and flushes + // mid-bundle during @ProcessElement, rather than only flushing at @FinishBundle. + pipeline + .apply(Create.of(data)) + .apply(WithKeys.of(0)) + .apply(GroupByKey.create()) + .apply(Values.create()) + .apply(ParDo.of(new EmitDocumentsFn())) + .apply(write); pipeline.run(); long currentNumDocs = SolrIOTestUtils.commitAndGetCurrentNumDocs(SOLR_COLLECTION, solrClient);