Visitar URL original
fix: deserialize Source and SourceWithKeys correctly so SampleRowKeys… · googleapis/java-bigtable-hbase@de7bdac · GitHub
Skip to content

Commit de7bdac

Browse files
mutianfmutianf
authored
fix: deserialize Source and SourceWithKeys correctly so SampleRowKeys only get called once (#4078)
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly: - [ ] Make sure to open an issue as a [bug/issue](https://togithub.com/googleapis/java-bigtable-hbase/issues/new/choose) before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea - [ ] Ensure the tests and linter pass - [ ] Code coverage does not decrease (if any source code was changed) - [ ] Appropriate docs were updated (if necessary) Fixes #<issue_number_goes_here> ☕️ If you write sample code, please follow the [samples format]( https://togithub.com/GoogleCloudPlatform/java-docs-samples/blob/main/SAMPLE_FORMAT.md).
1 parent c7831e5 commit de7bdac

2 files changed

Lines changed: 116 additions & 17 deletions

File tree

‎bigtable-dataflow-parent/bigtable-hbase-beam/src/main/java/com/google/cloud/bigtable/beam/CloudBigtableIO.java‎

Lines changed: 42 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -409,25 +409,15 @@ public void populateDisplayData(DisplayData.Builder builder) {
409409
configuration.populateDisplayData(builder);
410410
}
411411

412-
/**
413-
* The writeReplace method allows the developer to provide a replacement object that will be
414-
* serialized instead of the original one. We use this to keep the enclosed class immutable. For
415-
* more details on the technique see <a
416-
* href="https://lingpipe-blog.com/2009/08/10/serializing-immutable-singletons-serialization-proxy/">this
417-
* article</a>.
418-
*/
419-
private Object writeReplace() {
420-
return new SerializationProxy(getConfiguration());
421-
}
412+
abstract static class AbstractSerializationProxy implements Serializable {
422413

423-
static class SerializationProxy implements Serializable {
424414
private ValueProvider<String> projectId;
425415
private ValueProvider<String> instanceId;
426416
private ValueProvider<String> tableId;
427417
private Map<String, ValueProvider<String>> additionalConfiguration;
428418
private transient ValueProvider<Scan> scan;
429419

430-
public SerializationProxy(CloudBigtableScanConfiguration configuration) {
420+
public AbstractSerializationProxy(CloudBigtableScanConfiguration configuration) {
431421
this.projectId = configuration.getProjectIdValueProvider();
432422
this.instanceId = configuration.getInstanceIdValueProvider();
433423
this.tableId = configuration.getTableIdValueProvider();
@@ -472,11 +462,13 @@ private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundE
472462
}
473463
}
474464

475-
Object readResolve() {
465+
abstract Object readResolve();
466+
467+
protected CloudBigtableScanConfiguration buildScanConfig() {
476468
CloudBigtableScanConfiguration conf =
477469
CloudBigtableScanConfiguration.createConfig(
478470
projectId, instanceId, tableId, scan, additionalConfiguration);
479-
return CloudBigtableIO.read(conf);
471+
return conf;
480472
}
481473
}
482474
}
@@ -532,8 +524,26 @@ public Coder<Result> getOutputCoder() {
532524
return getResultCoder();
533525
}
534526

527+
/**
528+
* The writeReplace method allows the developer to provide a replacement object that will be
529+
* serialized instead of the original one. We use this to keep the enclosed class immutable. For
530+
* more details on the technique see <a
531+
* href="https://lingpipe-blog.com/2009/08/10/serializing-immutable-singletons-serialization-proxy/">this
532+
* article</a>.
533+
*/
535534
private Object writeReplace() {
536-
return new SerializationProxy(getConfiguration());
535+
return new SourceSerializationProxy(getConfiguration());
536+
}
537+
538+
static class SourceSerializationProxy extends AbstractSerializationProxy {
539+
public SourceSerializationProxy(CloudBigtableScanConfiguration configuration) {
540+
super(configuration);
541+
}
542+
543+
@Override
544+
Object readResolve() {
545+
return new Source(buildScanConfig());
546+
}
537547
}
538548
}
539549

@@ -554,7 +564,6 @@ protected static class SourceWithKeys extends AbstractSource {
554564

555565
protected SourceWithKeys(CloudBigtableScanConfiguration configuration, long estimatedSize) {
556566
super(configuration);
557-
558567
byte[] stopRow = configuration.getStopRow();
559568
if (stopRow.length > 0) {
560569
byte[] startRow = configuration.getStartRow();
@@ -618,7 +627,23 @@ public String toString() {
618627
}
619628

620629
private Object writeReplace() {
621-
return new SerializationProxy(getConfiguration());
630+
return new SourceWithKeysSerializationProxy(getConfiguration(), estimatedSize);
631+
}
632+
633+
static class SourceWithKeysSerializationProxy extends AbstractSerializationProxy {
634+
635+
private long estimatedSize;
636+
637+
public SourceWithKeysSerializationProxy(
638+
CloudBigtableScanConfiguration configuration, long estimatedSize) {
639+
super(configuration);
640+
this.estimatedSize = estimatedSize;
641+
}
642+
643+
Object readResolve() {
644+
CloudBigtableScanConfiguration conf = buildScanConfig();
645+
return new SourceWithKeys(conf, estimatedSize);
646+
}
622647
}
623648
}
624649
/** Reads rows for a specific {@link Table}, usually filtered by a {@link Scan}. */

‎bigtable-dataflow-parent/bigtable-hbase-beam/src/test/java/com/google/cloud/bigtable/beam/CloudBigtableIOTest.java‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package com.google.cloud.bigtable.beam;
1717

18+
import static org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
1819
import static org.junit.Assert.assertEquals;
1920
import static org.junit.Assert.assertNotNull;
2021
import static org.mockito.Mockito.when;
@@ -28,18 +29,22 @@
2829
import java.util.ArrayList;
2930
import java.util.Collections;
3031
import java.util.Comparator;
32+
import java.util.HashMap;
3133
import java.util.Iterator;
3234
import java.util.List;
3335
import org.apache.beam.sdk.Pipeline;
3436
import org.apache.beam.sdk.coders.CannotProvideCoderException;
3537
import org.apache.beam.sdk.coders.Coder;
3638
import org.apache.beam.sdk.coders.CoderRegistry;
3739
import org.apache.beam.sdk.io.BoundedSource;
40+
import org.apache.beam.sdk.util.SerializableUtils;
3841
import org.apache.beam.sdk.values.TypeDescriptor;
3942
import org.apache.hadoop.hbase.client.Delete;
4043
import org.apache.hadoop.hbase.client.Mutation;
4144
import org.apache.hadoop.hbase.client.Put;
4245
import org.apache.hadoop.hbase.client.Result;
46+
import org.apache.hadoop.hbase.client.Scan;
47+
import org.apache.hadoop.hbase.filter.KeyOnlyFilter;
4348
import org.apache.hadoop.hbase.util.Bytes;
4449
import org.junit.Assert;
4550
import org.junit.Before;
@@ -217,4 +222,73 @@ public void testWriteToMultipleTablesValidateConfig() throws Exception {
217222
Assert.assertTrue(e.getMessage().contains("A instanceId must be set"));
218223
}
219224
}
225+
226+
@Test
227+
public void testSerializeSource() {
228+
String projectId = "my-project";
229+
String instanceId = "my-instance";
230+
String tableId = "my-table";
231+
Scan scan = new Scan().setFilter(new KeyOnlyFilter()).setMaxVersions(2);
232+
CloudBigtableScanConfiguration configuration =
233+
CloudBigtableScanConfiguration.createConfig(
234+
StaticValueProvider.of(projectId),
235+
StaticValueProvider.of(instanceId),
236+
StaticValueProvider.of(tableId),
237+
StaticValueProvider.of(scan),
238+
new HashMap<>());
239+
240+
Source source = new Source(configuration);
241+
242+
Source deserialized = SerializableUtils.ensureSerializable(source);
243+
244+
assertEquals(projectId, deserialized.getConfiguration().getProjectId());
245+
assertEquals(instanceId, deserialized.getConfiguration().getInstanceId());
246+
assertEquals(tableId, deserialized.getConfiguration().getTableId());
247+
assertEquals(
248+
scan.getFilter(), deserialized.getConfiguration().getScanValueProvider().get().getFilter());
249+
assertEquals(2, scan.getMaxVersions());
250+
}
251+
252+
@Test
253+
public void testSerializeSourceWithKeys() {
254+
String projectId = "my-project";
255+
String instanceId = "my-instance";
256+
String tableId = "my-table";
257+
String startKey = "aaa1";
258+
String endKey = "bbb3";
259+
260+
Scan scan =
261+
new Scan()
262+
.withStartRow(ByteString.copyFromUtf8(startKey).toByteArray())
263+
.withStopRow(ByteString.copyFromUtf8(endKey).toByteArray())
264+
.setFilter(new KeyOnlyFilter());
265+
266+
CloudBigtableScanConfiguration sourceWithKeysConfiguration =
267+
CloudBigtableScanConfiguration.createConfig(
268+
StaticValueProvider.of(projectId),
269+
StaticValueProvider.of(instanceId),
270+
StaticValueProvider.of(tableId),
271+
StaticValueProvider.of(scan),
272+
new HashMap<>());
273+
274+
long estimatedSize = 123456;
275+
SourceWithKeys sourceWithKeys = new SourceWithKeys(sourceWithKeysConfiguration, estimatedSize);
276+
277+
SourceWithKeys deserialized = SerializableUtils.ensureSerializable(sourceWithKeys);
278+
279+
assertEquals(projectId, deserialized.getConfiguration().getProjectId());
280+
assertEquals(instanceId, deserialized.getConfiguration().getInstanceId());
281+
assertEquals(tableId, deserialized.getConfiguration().getTableId());
282+
assertEquals(
283+
startKey,
284+
ByteString.copyFrom(
285+
deserialized.getConfiguration().getScanValueProvider().get().getStartRow())
286+
.toStringUtf8());
287+
assertEquals(
288+
endKey,
289+
ByteString.copyFrom(
290+
deserialized.getConfiguration().getScanValueProvider().get().getStopRow())
291+
.toStringUtf8());
292+
assertEquals(estimatedSize, deserialized.getEstimatedSize());
293+
}
220294
}

0 commit comments

Comments
 (0)