This is an automated email from the ASF dual-hosted git repository.
kennknowles pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 876596e1afd Merge pull request #37731: RedisIO support PCollection of
<byte[], byte[]>
876596e1afd is described below
commit 876596e1afdca33d4462aee581507d56a3726b4a
Author: Michael Gruschke <[email protected]>
AuthorDate: Wed Sep 16 20:00:09 2026 +0200
Merge pull request #37731: RedisIO support PCollection of <byte[], byte[]>
---
.../java/org/apache/beam/sdk/io/redis/RedisIO.java | 547 +++++++++++++++++----
.../org/apache/beam/sdk/io/redis/RedisIOTest.java | 292 ++++++++++-
2 files changed, 736 insertions(+), 103 deletions(-)
diff --git
a/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
b/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
index 62109a6a666..6b5cd99c69e 100644
--- a/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
+++ b/sdks/java/io/redis/src/main/java/org/apache/beam/sdk/io/redis/RedisIO.java
@@ -20,19 +20,23 @@ package org.apache.beam.sdk.io.redis;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument;
import com.google.auto.value.AutoValue;
+import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.KvCoder;
-import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.range.ByteKey;
import org.apache.beam.sdk.io.range.ByteKeyRange;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Filter;
import org.apache.beam.sdk.transforms.Latest;
+import org.apache.beam.sdk.transforms.MapElements;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.SerializableFunctions;
+import org.apache.beam.sdk.transforms.SimpleFunction;
import org.apache.beam.sdk.transforms.View;
import org.apache.beam.sdk.transforms.display.DisplayData;
import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
@@ -41,6 +45,7 @@ import org.apache.beam.sdk.values.PBegin;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionView;
import org.apache.beam.sdk.values.PDone;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Longs;
import org.checkerframework.checker.nullness.qual.Nullable;
import redis.clients.jedis.Jedis;
import redis.clients.jedis.StreamEntryID;
@@ -106,6 +111,16 @@ import redis.clients.jedis.resps.ScanResult;
*
* }</pre>
*
+ * <p>{@link #writeBytes()} is similar to {@link #write()} but works with
key/value pairs
+ * represented as byte arrays:
+ *
+ * <pre>{@code
+ * pipeline.apply(...)
+ * // here we a have a PCollection<byte[], byte[]> with key/value pairs
+ * .apply(RedisIO.<byte[], byte[]>writeBytes().withEndpoint("::1", 6379))
+ *
+ * }</pre>
+ *
* <h3>Writing Redis Streams</h3>
*
* <p>{@link #writeStreams()} appends the entries of a {@link PCollection} of
key/value pairs
@@ -159,12 +174,32 @@ public class RedisIO {
.build();
}
+ /** Read binary data from a Redis server. */
+ public static ReadBytes readBytes() {
+ return new AutoValue_RedisIO_ReadBytes.Builder()
+ .setConnectionConfiguration(RedisConnectionConfiguration.create())
+ .setKeyPattern("*")
+ .setBatchSize(1000)
+ .setOutputParallelization(true)
+ .build();
+ }
+
/**
* Like {@link #read()} but executes multiple instances of the Redis query
substituting each
* element of a {@link PCollection} as key pattern.
*/
public static ReadKeyPatterns readKeyPatterns() {
return new AutoValue_RedisIO_ReadKeyPatterns.Builder()
+ .setReadKeyPatternsBytes(readKeyPatternsBytes())
+ .build();
+ }
+
+ /**
+ * Like {@link #read()} but executes multiple instances of the Redis query
substituting each
+ * element of a {@link PCollection} as key pattern.
+ */
+ public static ReadKeyPatternsBytes readKeyPatternsBytes() {
+ return new AutoValue_RedisIO_ReadKeyPatternsBytes.Builder()
.setConnectionConfiguration(RedisConnectionConfiguration.create())
.setBatchSize(1000)
.setOutputParallelization(true)
@@ -173,7 +208,12 @@ public class RedisIO {
/** Write data to a Redis server. */
public static Write write() {
- return new AutoValue_RedisIO_Write.Builder()
+ return new
AutoValue_RedisIO_Write.Builder().setWriteBytes(writeBytes()).build();
+ }
+
+ /** Write data to a Redis server. */
+ public static WriteBytes writeBytes() {
+ return new AutoValue_RedisIO_WriteBytes.Builder()
.setConnectionConfiguration(RedisConnectionConfiguration.create())
.setMethod(Write.Method.APPEND)
.build();
@@ -182,6 +222,12 @@ public class RedisIO {
/** Write stream data to a Redis server. */
public static WriteStreams writeStreams() {
return new AutoValue_RedisIO_WriteStreams.Builder()
+ .setWriteStreamsBytes(writeStreamsBytes())
+ .build();
+ }
+
+ public static WriteStreamsBytes writeStreamsBytes() {
+ return new AutoValue_RedisIO_WriteStreamsBytes.Builder()
.setConnectionConfiguration(RedisConnectionConfiguration.create())
.setMaxLen(0L)
.setApproximateTrim(true)
@@ -282,33 +328,187 @@ public class RedisIO {
}
}
- /** Implementation of {@link #readKeyPatterns()}. */
+ /** Implementation of {@link #readBytes()}. */
@AutoValue
- public abstract static class ReadKeyPatterns
- extends PTransform<PCollection<String>, PCollection<KV<String, String>>>
{
+ public abstract static class ReadBytes
+ extends PTransform<PBegin, PCollection<KV<byte[], byte[]>>> {
abstract @Nullable RedisConnectionConfiguration connectionConfiguration();
+ abstract @Nullable String keyPattern();
+
abstract int batchSize();
abstract boolean outputParallelization();
abstract Builder toBuilder();
+ public ReadBytes withEndpoint(String host, int port) {
+ checkArgument(host != null, "host can not be null");
+ checkArgument(0 < port && port < 65536, "port must be a positive integer
less than 65536");
+ return toBuilder()
+
.setConnectionConfiguration(connectionConfiguration().withHost(host).withPort(port))
+ .build();
+ }
+
+ public ReadBytes withAuth(String auth) {
+ checkArgument(auth != null, "auth can not be null");
+ return toBuilder()
+ .setConnectionConfiguration(connectionConfiguration().withAuth(auth))
+ .build();
+ }
+
+ public ReadBytes withTimeout(int timeout) {
+ checkArgument(timeout >= 0, "timeout can not be negative");
+ return toBuilder()
+
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
+ .build();
+ }
+
+ public ReadBytes withKeyPattern(String keyPattern) {
+ checkArgument(keyPattern != null, "keyPattern can not be null");
+ return toBuilder().setKeyPattern(keyPattern).build();
+ }
+
+ public ReadBytes withConnectionConfiguration(RedisConnectionConfiguration
connection) {
+ checkArgument(connection != null, "connection can not be null");
+ return toBuilder().setConnectionConfiguration(connection).build();
+ }
+
+ public ReadBytes withBatchSize(int batchSize) {
+ return toBuilder().setBatchSize(batchSize).build();
+ }
+
+ /**
+ * Whether to reshuffle the resulting PCollection so results are
distributed to all workers. The
+ * default is to parallelize and should only be changed if this is known
to be unnecessary.
+ */
+ public ReadBytes withOutputParallelization(boolean outputParallelization) {
+ return
toBuilder().setOutputParallelization(outputParallelization).build();
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ connectionConfiguration().populateDisplayData(builder);
+ }
+
+ @Override
+ public PCollection<KV<byte[], byte[]>> expand(PBegin input) {
+ checkArgument(connectionConfiguration() != null,
"withConnectionConfiguration() is required");
+
+ return input
+ .apply(Create.of(keyPattern()))
+ .apply(
+ RedisIO.readKeyPatternsBytes()
+ .withConnectionConfiguration(connectionConfiguration())
+ .withBatchSize(batchSize())
+ .withOutputParallelization(outputParallelization()));
+ }
+
@AutoValue.Builder
abstract static class Builder {
abstract @Nullable Builder setConnectionConfiguration(
RedisConnectionConfiguration connection);
+ abstract @Nullable Builder setKeyPattern(String keyPattern);
+
abstract Builder setBatchSize(int batchSize);
abstract Builder setOutputParallelization(boolean outputParallelization);
- abstract ReadKeyPatterns build();
+ abstract ReadBytes build();
}
+ }
+
+ /** Implementation of {@link #readKeyPatterns()}. */
+ @AutoValue
+ public abstract static class ReadKeyPatterns
+ extends PTransform<PCollection<String>, PCollection<KV<String, String>>>
{
+
+ abstract ReadKeyPatternsBytes readKeyPatternsBytes();
+
+ abstract Builder toBuilder();
public ReadKeyPatterns withEndpoint(String host, int port) {
+ return toBuilder()
+ .setReadKeyPatternsBytes(readKeyPatternsBytes().withEndpoint(host,
port))
+ .build();
+ }
+
+ public ReadKeyPatterns withAuth(String auth) {
+ return
toBuilder().setReadKeyPatternsBytes(readKeyPatternsBytes().withAuth(auth)).build();
+ }
+
+ public ReadKeyPatterns withTimeout(int timeout) {
+ return toBuilder()
+ .setReadKeyPatternsBytes(readKeyPatternsBytes().withTimeout(timeout))
+ .build();
+ }
+
+ public ReadKeyPatterns
withConnectionConfiguration(RedisConnectionConfiguration connection) {
+ return toBuilder()
+
.setReadKeyPatternsBytes(readKeyPatternsBytes().withConnectionConfiguration(connection))
+ .build();
+ }
+
+ public ReadKeyPatterns withBatchSize(int batchSize) {
+ return toBuilder()
+
.setReadKeyPatternsBytes(readKeyPatternsBytes().withBatchSize(batchSize))
+ .build();
+ }
+
+ /**
+ * Whether to reshuffle the resulting PCollection so results are
distributed to all workers. The
+ * default is to parallelize and should only be changed if this is known
to be unnecessary.
+ */
+ public ReadKeyPatterns withOutputParallelization(boolean
outputParallelization) {
+ return toBuilder()
+ .setReadKeyPatternsBytes(
+
readKeyPatternsBytes().withOutputParallelization(outputParallelization))
+ .build();
+ }
+
+ @Override
+ public PCollection<KV<String, String>> expand(PCollection<String> input) {
+
+ return input
+ .apply(readKeyPatternsBytes())
+ .apply(
+ MapElements.via(
+ new SimpleFunction<KV<byte[], byte[]>, KV<String, String>>()
{
+ @Override
+ public KV<String, String> apply(KV<byte[], byte[]> input) {
+ return KV.of(
+ new String(input.getKey(), StandardCharsets.UTF_8),
+ new String(input.getValue(),
StandardCharsets.UTF_8));
+ }
+ }));
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ abstract Builder setReadKeyPatternsBytes(ReadKeyPatternsBytes
readKeyPatternsBytes);
+
+ abstract ReadKeyPatterns build();
+ }
+ }
+
+ /** Implementation of {@link #readKeyPatternsBytes()}. */
+ @AutoValue
+ public abstract static class ReadKeyPatternsBytes
+ extends PTransform<PCollection<String>, PCollection<KV<byte[], byte[]>>>
{
+
+ abstract @Nullable RedisConnectionConfiguration connectionConfiguration();
+
+ abstract int batchSize();
+
+ abstract boolean outputParallelization();
+
+ abstract Builder toBuilder();
+
+ public ReadKeyPatternsBytes withEndpoint(String host, int port) {
checkArgument(host != null, "host can not be null");
checkArgument(port > 0, "port can not be negative or 0");
return toBuilder()
@@ -316,26 +516,27 @@ public class RedisIO {
.build();
}
- public ReadKeyPatterns withAuth(String auth) {
+ public ReadKeyPatternsBytes withAuth(String auth) {
checkArgument(auth != null, "auth can not be null");
return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withAuth(auth))
.build();
}
- public ReadKeyPatterns withTimeout(int timeout) {
+ public ReadKeyPatternsBytes withTimeout(int timeout) {
checkArgument(timeout >= 0, "timeout can not be negative");
return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
.build();
}
- public ReadKeyPatterns
withConnectionConfiguration(RedisConnectionConfiguration connection) {
+ public ReadKeyPatternsBytes withConnectionConfiguration(
+ RedisConnectionConfiguration connection) {
checkArgument(connection != null, "connection can not be null");
return toBuilder().setConnectionConfiguration(connection).build();
}
- public ReadKeyPatterns withBatchSize(int batchSize) {
+ public ReadKeyPatternsBytes withBatchSize(int batchSize) {
return toBuilder().setBatchSize(batchSize).build();
}
@@ -343,26 +544,40 @@ public class RedisIO {
* Whether to reshuffle the resulting PCollection so results are
distributed to all workers. The
* default is to parallelize and should only be changed if this is known
to be unnecessary.
*/
- public ReadKeyPatterns withOutputParallelization(boolean
outputParallelization) {
+ public ReadKeyPatternsBytes withOutputParallelization(boolean
outputParallelization) {
return
toBuilder().setOutputParallelization(outputParallelization).build();
}
@Override
- public PCollection<KV<String, String>> expand(PCollection<String> input) {
+ public PCollection<KV<byte[], byte[]>> expand(PCollection<String> input) {
checkArgument(connectionConfiguration() != null,
"withConnectionConfiguration() is required");
- PCollection<KV<String, String>> output =
+
+ PCollection<KV<byte[], byte[]>> output =
input
.apply(ParDo.of(new ReadFn(connectionConfiguration())))
- .setCoder(KvCoder.of(StringUtf8Coder.of(),
StringUtf8Coder.of()));
+ .setCoder(KvCoder.of(ByteArrayCoder.of(), ByteArrayCoder.of()));
if (outputParallelization()) {
output = output.apply(new Reparallelize());
}
return output;
}
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ abstract @Nullable Builder setConnectionConfiguration(
+ RedisConnectionConfiguration connection);
+
+ abstract Builder setBatchSize(int batchSize);
+
+ abstract Builder setOutputParallelization(boolean outputParallelization);
+
+ abstract ReadKeyPatternsBytes build();
+ }
}
@DoFn.BoundedPerElement
- private static class ReadFn extends DoFn<String, KV<String, String>> {
+ private static class ReadFn extends DoFn<String, KV<byte[], byte[]>> {
protected final RedisConnectionConfiguration connectionConfiguration;
transient Jedis jedis;
@@ -394,10 +609,11 @@ public class RedisIO {
ScanParams scanParams = new ScanParams();
scanParams.match(c.element());
while (tracker.tryClaim(cursor)) {
- ScanResult<String> scanResult = jedis.scan(redisCursor.getCursor(),
scanParams);
+ byte[] cursorBytes =
redisCursor.getCursor().getBytes(StandardCharsets.UTF_8);
+ ScanResult<byte[]> scanResult = jedis.scan(cursorBytes, scanParams);
if (scanResult.getResult().size() > 0) {
- String[] keys = scanResult.getResult().toArray(new
String[scanResult.getResult().size()]);
- List<String> results = jedis.mget(keys);
+ byte[][] keys = scanResult.getResult().toArray(new
byte[scanResult.getResult().size()][]);
+ List<byte[]> results = jedis.mget(keys);
for (int i = 0; i < results.size(); i++) {
if (results.get(i) != null) {
c.output(KV.of(keys[i], results.get(i)));
@@ -411,20 +627,20 @@ public class RedisIO {
}
private static class Reparallelize
- extends PTransform<PCollection<KV<String, String>>,
PCollection<KV<String, String>>> {
+ extends PTransform<PCollection<KV<byte[], byte[]>>,
PCollection<KV<byte[], byte[]>>> {
@Override
- public PCollection<KV<String, String>> expand(PCollection<KV<String,
String>> input) {
+ public PCollection<KV<byte[], byte[]>> expand(PCollection<KV<byte[],
byte[]>> input) {
// reparallelize mimics the same behavior as in JdbcIO, used to break
fusion
- PCollectionView<Iterable<KV<String, String>>> empty =
+ PCollectionView<Iterable<KV<byte[], byte[]>>> empty =
input
.apply("Consume",
Filter.by(SerializableFunctions.constant(false)))
.apply(View.asIterable());
- PCollection<KV<String, String>> materialized =
+ PCollection<KV<byte[], byte[]>> materialized =
input.apply(
"Identity",
ParDo.of(
- new DoFn<KV<String, String>, KV<String, String>>() {
+ new DoFn<KV<byte[], byte[]>, KV<byte[], byte[]>>() {
@ProcessElement
public void processElement(ProcessContext c) {
c.output(c.element());
@@ -441,6 +657,66 @@ public class RedisIO {
@AutoValue
public abstract static class Write extends PTransform<PCollection<KV<String,
String>>, PDone> {
+ abstract WriteBytes writeBytes();
+
+ public abstract Builder toBuilder();
+
+ public Write withEndpoint(String host, int port) {
+ return toBuilder().setWriteBytes(writeBytes().withEndpoint(host,
port)).build();
+ }
+
+ public Write withAuth(String auth) {
+ return toBuilder().setWriteBytes(writeBytes().withAuth(auth)).build();
+ }
+
+ public Write withTimeout(int timeout) {
+ return
toBuilder().setWriteBytes(writeBytes().withTimeout(timeout)).build();
+ }
+
+ public Write withConnectionConfiguration(RedisConnectionConfiguration
connection) {
+ return toBuilder()
+ .setWriteBytes(writeBytes().withConnectionConfiguration(connection))
+ .build();
+ }
+
+ public Write withMethod(Method method) {
+ return
toBuilder().setWriteBytes(writeBytes().withMethod(method)).build();
+ }
+
+ public Write withExpireTime(Long expireTimeMillis) {
+ return
toBuilder().setWriteBytes(writeBytes().withExpireTime(expireTimeMillis)).build();
+ }
+
+ @Override
+ public PDone expand(PCollection<KV<String, String>> input) {
+ MapElements<KV<String, String>, KV<byte[], byte[]>> toBytes;
+ Method method = writeBytes().method();
+ if (Method.INCRBY == method || Method.DECRBY == method) {
+ toBytes =
+ MapElements.via(
+ new SimpleFunction<KV<String, String>, KV<byte[], byte[]>>() {
+ @Override
+ public KV<byte[], byte[]> apply(KV<String, String> input) {
+ long value = Long.parseLong(input.getValue().trim());
+ return KV.of(
+ input.getKey().getBytes(StandardCharsets.UTF_8),
Longs.toByteArray(value));
+ }
+ });
+ } else {
+ toBytes =
+ MapElements.via(
+ new SimpleFunction<KV<String, String>, KV<byte[], byte[]>>() {
+ @Override
+ public KV<byte[], byte[]> apply(KV<String, String> input) {
+ return KV.of(
+ input.getKey().getBytes(StandardCharsets.UTF_8),
+ input.getValue().getBytes(StandardCharsets.UTF_8));
+ }
+ });
+ }
+ return input.apply(toBytes).apply(writeBytes());
+ }
+
/** Determines the method used to insert data in Redis. */
public enum Method {
@@ -480,9 +756,23 @@ public class RedisIO {
DECRBY,
}
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ abstract Builder setWriteBytes(WriteBytes writeBytes);
+
+ abstract Write build();
+ }
+ }
+
+ /** A {@link PTransform} to write to a Redis server. */
+ @AutoValue
+ public abstract static class WriteBytes
+ extends PTransform<PCollection<KV<byte[], byte[]>>, PDone> {
+
abstract @Nullable RedisConnectionConfiguration connectionConfiguration();
- abstract @Nullable Method method();
+ abstract Write.@Nullable Method method();
abstract @Nullable Long expireTime();
@@ -494,14 +784,14 @@ public class RedisIO {
abstract Builder setConnectionConfiguration(
RedisConnectionConfiguration connectionConfiguration);
- abstract Builder setMethod(Method method);
+ abstract Builder setMethod(Write.Method method);
abstract Builder setExpireTime(Long expireTimeMillis);
- abstract Write build();
+ abstract WriteBytes build();
}
- public Write withEndpoint(String host, int port) {
+ public WriteBytes withEndpoint(String host, int port) {
checkArgument(host != null, "host can not be null");
checkArgument(port > 0, "port can not be negative or 0");
return toBuilder()
@@ -509,56 +799,56 @@ public class RedisIO {
.build();
}
- public Write withAuth(String auth) {
+ public WriteBytes withAuth(String auth) {
checkArgument(auth != null, "auth can not be null");
return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withAuth(auth))
.build();
}
- public Write withTimeout(int timeout) {
+ public WriteBytes withTimeout(int timeout) {
checkArgument(timeout >= 0, "timeout can not be negative");
return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
.build();
}
- public Write withConnectionConfiguration(RedisConnectionConfiguration
connection) {
+ public WriteBytes withConnectionConfiguration(RedisConnectionConfiguration
connection) {
checkArgument(connection != null, "connection can not be null");
return toBuilder().setConnectionConfiguration(connection).build();
}
- public Write withMethod(Method method) {
+ public WriteBytes withMethod(Write.Method method) {
checkArgument(method != null, "method can not be null");
return toBuilder().setMethod(method).build();
}
- public Write withExpireTime(Long expireTimeMillis) {
+ public WriteBytes withExpireTime(Long expireTimeMillis) {
checkArgument(expireTimeMillis != null, "expireTimeMillis can not be
null");
checkArgument(expireTimeMillis > 0, "expireTimeMillis can not be
negative or 0");
return toBuilder().setExpireTime(expireTimeMillis).build();
}
@Override
- public PDone expand(PCollection<KV<String, String>> input) {
+ public PDone expand(PCollection<KV<byte[], byte[]>> input) {
checkArgument(connectionConfiguration() != null,
"withConnectionConfiguration() is required");
input.apply(ParDo.of(new WriteFn(this)));
return PDone.in(input.getPipeline());
}
- private static class WriteFn extends DoFn<KV<String, String>, Void> {
+ private static class WriteFn extends DoFn<KV<byte[], byte[]>, Void> {
private static final int DEFAULT_BATCH_SIZE = 1000;
- private final Write spec;
+ private final WriteBytes spec;
private transient Jedis jedis;
private transient @Nullable Transaction transaction;
private int batchCount;
- public WriteFn(Write spec) {
+ public WriteFn(WriteBytes spec) {
this.spec = spec;
}
@@ -575,7 +865,7 @@ public class RedisIO {
@ProcessElement
public void processElement(ProcessContext c) {
- KV<String, String> record = c.element();
+ KV<byte[], byte[]> record = c.element();
writeRecord(record);
@@ -588,39 +878,39 @@ public class RedisIO {
}
}
- private void writeRecord(KV<String, String> record) {
- Method method = spec.method();
+ private void writeRecord(KV<byte[], byte[]> record) {
+ Write.Method method = spec.method();
Long expireTime = spec.expireTime();
- if (Method.APPEND == method) {
+ if (Write.Method.APPEND == method) {
writeUsingAppendCommand(record, expireTime);
- } else if (Method.SET == method) {
+ } else if (Write.Method.SET == method) {
writeUsingSetCommand(record, expireTime);
- } else if (Method.LPUSH == method || Method.RPUSH == method) {
+ } else if (Write.Method.LPUSH == method || Write.Method.RPUSH ==
method) {
writeUsingListCommand(record, method, expireTime);
- } else if (Method.SADD == method) {
+ } else if (Write.Method.SADD == method) {
writeUsingSaddCommand(record, expireTime);
- } else if (Method.PFADD == method) {
+ } else if (Write.Method.PFADD == method) {
writeUsingHLLCommand(record, expireTime);
- } else if (Method.INCRBY == method) {
+ } else if (Write.Method.INCRBY == method) {
writeUsingIncrBy(record, expireTime);
- } else if (Method.DECRBY == method) {
+ } else if (Write.Method.DECRBY == method) {
writeUsingDecrBy(record, expireTime);
}
}
- private void writeUsingAppendCommand(KV<String, String> record, Long
expireTime) {
- String key = record.getKey();
- String value = record.getValue();
+ private void writeUsingAppendCommand(KV<byte[], byte[]> record, Long
expireTime) {
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
transaction.append(key, value);
setExpireTimeWhenRequired(key, expireTime);
}
- private void writeUsingSetCommand(KV<String, String> record, Long
expireTime) {
- String key = record.getKey();
- String value = record.getValue();
+ private void writeUsingSetCommand(KV<byte[], byte[]> record, Long
expireTime) {
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
if (expireTime != null) {
transaction.psetex(key, expireTime, value);
@@ -630,57 +920,57 @@ public class RedisIO {
}
private void writeUsingListCommand(
- KV<String, String> record, Method method, Long expireTime) {
+ KV<byte[], byte[]> record, Write.Method method, Long expireTime) {
- String key = record.getKey();
- String value = record.getValue();
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
- if (Method.LPUSH == method) {
+ if (Write.Method.LPUSH == method) {
transaction.lpush(key, value);
- } else if (Method.RPUSH == method) {
+ } else if (Write.Method.RPUSH == method) {
transaction.rpush(key, value);
}
setExpireTimeWhenRequired(key, expireTime);
}
- private void writeUsingSaddCommand(KV<String, String> record, Long
expireTime) {
- String key = record.getKey();
- String value = record.getValue();
+ private void writeUsingSaddCommand(KV<byte[], byte[]> record, Long
expireTime) {
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
transaction.sadd(key, value);
setExpireTimeWhenRequired(key, expireTime);
}
- private void writeUsingHLLCommand(KV<String, String> record, Long
expireTime) {
- String key = record.getKey();
- String value = record.getValue();
+ private void writeUsingHLLCommand(KV<byte[], byte[]> record, Long
expireTime) {
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
transaction.pfadd(key, value);
setExpireTimeWhenRequired(key, expireTime);
}
- private void writeUsingIncrBy(KV<String, String> record, Long
expireTime) {
- String key = record.getKey();
- String value = record.getValue();
- long inc = Long.parseLong(value);
+ private void writeUsingIncrBy(KV<byte[], byte[]> record, Long
expireTime) {
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
+ long inc = Longs.fromByteArray(value);
transaction.incrBy(key, inc);
setExpireTimeWhenRequired(key, expireTime);
}
- private void writeUsingDecrBy(KV<String, String> record, Long
expireTime) {
- String key = record.getKey();
- String value = record.getValue();
- long decr = Long.parseLong(value);
+ private void writeUsingDecrBy(KV<byte[], byte[]> record, Long
expireTime) {
+ byte[] key = record.getKey();
+ byte[] value = record.getValue();
+ long decr = Longs.fromByteArray(value);
transaction.decrBy(key, decr);
setExpireTimeWhenRequired(key, expireTime);
}
- private void setExpireTimeWhenRequired(String key, Long expireTime) {
+ private void setExpireTimeWhenRequired(byte[] key, Long expireTime) {
if (expireTime != null) {
transaction.pexpire(key, expireTime);
}
@@ -705,13 +995,77 @@ public class RedisIO {
}
}
+ @AutoValue
+ public abstract static class WriteStreams
+ extends PTransform<PCollection<KV<String, Map<String, String>>>, PDone> {
+
+ abstract WriteStreamsBytes writeStreamsBytes();
+
+ public abstract Builder toBuilder();
+
+ public WriteStreams withMaxLen(long maxLen) {
+ return
toBuilder().setWriteStreamsBytes(writeStreamsBytes().withMaxLen(maxLen)).build();
+ }
+
+ public WriteStreams withEndpoint(String host, int port) {
+ return
toBuilder().setWriteStreamsBytes(writeStreamsBytes().withEndpoint(host,
port)).build();
+ }
+
+ public WriteStreams
withConnectionConfiguration(RedisConnectionConfiguration connection) {
+ return toBuilder()
+
.setWriteStreamsBytes(writeStreamsBytes().withConnectionConfiguration(connection))
+ .build();
+ }
+
+ /**
+ * If {@link #withMaxLen(long)} is used, set the "~" prefix to the MAXLEN
value, indicating to
+ * the server that it should use "close enough" trimming.
+ */
+ public WriteStreams withApproximateTrim(boolean approximateTrim) {
+ return toBuilder()
+
.setWriteStreamsBytes(writeStreamsBytes().withApproximateTrim(approximateTrim))
+ .build();
+ }
+
+ @Override
+ public PDone expand(PCollection<KV<String, Map<String, String>>> input) {
+ return input.apply(ParDo.of(new
StringToByteFn())).apply(writeStreamsBytes());
+ }
+
+ @AutoValue.Builder
+ abstract static class Builder {
+
+ abstract Builder setWriteStreamsBytes(WriteStreamsBytes
writeStreamsBytes);
+
+ abstract WriteStreams build();
+ }
+
+ private static class StringToByteFn
+ extends DoFn<KV<String, Map<String, String>>, KV<byte[], Map<byte[],
byte[]>>> {
+
+ @ProcessElement
+ public void processElement(
+ @Element KV<String, Map<String, String>> element,
+ OutputReceiver<KV<byte[], Map<byte[], byte[]>>> out) {
+ byte[] key = element.getKey().getBytes(StandardCharsets.UTF_8);
+ Map<byte[], byte[]> byteMap =
+ element.getValue().entrySet().stream()
+ .collect(
+ Collectors.toMap(
+ e -> e.getKey().getBytes(StandardCharsets.UTF_8),
+ e -> e.getValue().getBytes(StandardCharsets.UTF_8)));
+ out.output(KV.of(key, byteMap));
+ }
+ }
+ }
+
/**
* A {@link PTransform} to write stream key pairs
(https://redis.io/topics/streams-intro) to a
* Redis server.
*/
@AutoValue
- public abstract static class WriteStreams
- extends PTransform<PCollection<KV<String, Map<String, String>>>, PDone> {
+ public abstract static class WriteStreamsBytes
+ extends PTransform<PCollection<KV<byte[], Map<byte[], byte[]>>>, PDone> {
abstract RedisConnectionConfiguration connectionConfiguration();
@@ -722,7 +1076,7 @@ public class RedisIO {
abstract Builder toBuilder();
/** Set the hostname and port of the Redis server to connect to. */
- public WriteStreams withEndpoint(String host, int port) {
+ public WriteStreamsBytes withEndpoint(String host, int port) {
checkArgument(host != null, "host can not be null");
checkArgument(port > 0, "port can not be negative or 0");
return toBuilder()
@@ -735,7 +1089,7 @@ public class RedisIO {
* either just a password or a username and password separated by a space.
See
* https://redis.io/commands/auth for details
*/
- public WriteStreams withAuth(String auth) {
+ public WriteStreamsBytes withAuth(String auth) {
checkArgument(auth != null, "auth can not be null");
return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withAuth(auth))
@@ -743,7 +1097,7 @@ public class RedisIO {
}
/** Set the connection timeout for the Redis server connection. */
- public WriteStreams withTimeout(int timeout) {
+ public WriteStreamsBytes withTimeout(int timeout) {
checkArgument(timeout >= 0, "timeout can not be negative");
return toBuilder()
.setConnectionConfiguration(connectionConfiguration().withTimeout(timeout))
@@ -751,13 +1105,13 @@ public class RedisIO {
}
/** Predefine a {@link RedisConnectionConfiguration} and pass it to the
builder. */
- public WriteStreams
withConnectionConfiguration(RedisConnectionConfiguration connection) {
+ public WriteStreamsBytes
withConnectionConfiguration(RedisConnectionConfiguration connection) {
checkArgument(connection != null, "connection can not be null");
return toBuilder().setConnectionConfiguration(connection).build();
}
/** When appending (XADD) to a stream, set a MAXLEN option. */
- public WriteStreams withMaxLen(long maxLen) {
+ public WriteStreamsBytes withMaxLen(long maxLen) {
checkArgument(maxLen >= 0L, "maxLen must be positive if set");
return toBuilder().setMaxLen(maxLen).build();
}
@@ -766,10 +1120,18 @@ public class RedisIO {
* If {@link #withMaxLen(long)} is used, set the "~" prefix to the MAXLEN
value, indicating to
* the server that it should use "close enough" trimming.
*/
- public WriteStreams withApproximateTrim(boolean approximateTrim) {
+ public WriteStreamsBytes withApproximateTrim(boolean approximateTrim) {
return toBuilder().setApproximateTrim(approximateTrim).build();
}
+ @Override
+ public PDone expand(PCollection<KV<byte[], Map<byte[], byte[]>>> input) {
+ checkArgument(connectionConfiguration() != null,
"withConnectionConfiguration() is required");
+
+ input.apply(ParDo.of(new WriteStreamFn(this)));
+ return PDone.in(input.getPipeline());
+ }
+
@AutoValue.Builder
abstract static class Builder {
@@ -780,29 +1142,20 @@ public class RedisIO {
abstract Builder setApproximateTrim(boolean approximateTrim);
- abstract WriteStreams build();
+ abstract WriteStreamsBytes build();
}
- @Override
- public PDone expand(PCollection<KV<String, Map<String, String>>> input) {
- checkArgument(connectionConfiguration() != null,
"withConnectionConfiguration() is required");
-
- input.apply(ParDo.of(new WriteStreamFn(this)));
- return PDone.in(input.getPipeline());
- }
-
- private static class WriteStreamFn extends DoFn<KV<String, Map<String,
String>>, Void> {
+ private static class WriteStreamFn extends DoFn<KV<byte[], Map<byte[],
byte[]>>, Void> {
private static final int DEFAULT_BATCH_SIZE = 1000;
- private final WriteStreams spec;
-
+ private final WriteStreamsBytes spec;
private transient Jedis jedis;
private transient @Nullable Transaction transaction;
private int batchCount;
- public WriteStreamFn(WriteStreams spec) {
+ public WriteStreamFn(WriteStreamsBytes spec) {
this.spec = spec;
}
@@ -819,7 +1172,7 @@ public class RedisIO {
@ProcessElement
public void processElement(ProcessContext c) {
- KV<String, Map<String, String>> record = c.element();
+ KV<byte[], Map<byte[], byte[]>> record = c.element();
writeRecord(record);
@@ -832,9 +1185,7 @@ public class RedisIO {
}
}
- private void writeRecord(KV<String, Map<String, String>> record) {
- String key = record.getKey();
- Map<String, String> value = record.getValue();
+ private void writeRecord(KV<byte[], Map<byte[], byte[]>> record) {
final XAddParams params = new XAddParams().id(StreamEntryID.NEW_ENTRY);
if (spec.maxLen() > 0L) {
params.maxLen(spec.maxLen());
@@ -842,7 +1193,7 @@ public class RedisIO {
params.approximateTrimming();
}
}
- transaction.xadd(key, params, value);
+ transaction.xadd(record.getKey(), params, record.getValue());
}
@FinishBundle
diff --git
a/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
b/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
index 7b305d846c7..88141fcc33f 100644
---
a/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
+++
b/sdks/java/io/redis/src/test/java/org/apache/beam/sdk/io/redis/RedisIOTest.java
@@ -21,9 +21,11 @@ import static java.util.stream.Collectors.toList;
import static
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists.transform;
import static org.hamcrest.CoreMatchers.hasItems;
import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;
+import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
@@ -31,8 +33,10 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
+import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
import org.apache.beam.sdk.coders.KvCoder;
import org.apache.beam.sdk.coders.MapCoder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
@@ -45,6 +49,8 @@ import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.PCollection;
import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Ints;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.primitives.Longs;
import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.ClassRule;
@@ -157,6 +163,36 @@ public class RedisIOTest {
p.run();
}
+ @Test
+ public void testReadBytes() {
+ for (int i = 0; i < 10; i++) {
+ byte[] key = ("binaryread" + i).getBytes(StandardCharsets.UTF_8);
+ byte[] value = new byte[] {(byte) i};
+ client.set(key, value);
+ }
+
+ PCollection<KV<byte[], byte[]>> read =
+ p.apply(
+ "Read",
+ RedisIO.readBytes().withEndpoint(redisHost,
port).withKeyPattern("binaryread*"));
+
+ // Verify actual byte content
+ PAssert.that(read)
+ .satisfies(
+ results -> {
+ Set<String> seen = new HashSet<>();
+ for (KV<byte[], byte[]> kv : results) {
+ String key = new String(kv.getKey(), StandardCharsets.UTF_8);
+ int idx =
Integer.parseInt(key.substring("binaryread".length()));
+ assertArrayEquals(new byte[] {(byte) idx}, kv.getValue());
+ seen.add(key);
+ }
+ assertEquals(10, seen.size());
+ return null;
+ });
+ p.run();
+ }
+
@Test
public void testWriteWithMethodSet() {
String key = "testWriteWithMethodSet";
@@ -187,7 +223,7 @@ public class RedisIOTest {
assertEquals(newValue, client.get(key));
Long expireTime = client.pttl(key);
- assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <=
10_0000);
+ assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <=
10_000);
client.del(key);
}
@@ -268,7 +304,7 @@ public class RedisIOTest {
long count = client.pfcount(key);
assertEquals(6, count);
Long expireTime = client.pttl(key);
- assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <=
10_0000);
+ assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <=
10_000);
client.del(key);
}
@@ -303,6 +339,196 @@ public class RedisIOTest {
assertEquals(-1, count);
}
+ @Test
+ public void testWriteBytesWithMethodSet() {
+ byte[] key = {1};
+ client.set(key, new byte[] {1});
+
+ byte[] newValue = {2};
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(KV.of(key,
newValue)));
+ write.apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.SET));
+ p.run();
+
+ assertArrayEquals(newValue, client.get(key));
+ assertEquals(NO_EXPIRATION, Long.valueOf(client.ttl(key)));
+ }
+
+ @Test
+ public void testWriteBytesWithMethodSetWithExpiration() {
+ byte[] key = {2};
+ client.set(key, new byte[] {1});
+
+ byte[] newValue = {2};
+
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(KV.of(key,
newValue)));
+ write.apply(
+ RedisIO.writeBytes()
+ .withEndpoint(redisHost, port)
+ .withMethod(Method.SET)
+ .withExpireTime(10_000L));
+ p.run();
+
+ assertArrayEquals(newValue, client.get(key));
+ Long expireTime = client.pttl(key);
+ assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <=
10_000);
+ client.del(key);
+ }
+
+ @Test
+ public void testWriteBytesWithMethodLPush() {
+ byte[] key = {3};
+ byte[] value = {1};
+ client.lpush(key, value);
+
+ byte[] newValue = {2};
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(KV.of(key,
newValue)));
+ write.apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.LPUSH));
+ p.run();
+
+ List<byte[]> values = client.lrange(key, 0, -1);
+ List<byte[]> expected = Arrays.asList(new byte[] {2}, new byte[] {1});
+ assertEquals(expected.size(), values.size());
+ for (int i = 0; i < expected.size(); i++) {
+ assertArrayEquals(expected.get(i), values.get(i));
+ }
+ }
+
+ @Test
+ public void testWriteBytesWithMethodRPush() {
+ byte[] key = {4};
+ byte[] value = {1};
+ client.lpush(key, value);
+
+ byte[] newValue = {2};
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(KV.of(key,
newValue)));
+ write.apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.RPUSH));
+ p.run();
+
+ List<byte[]> values = client.lrange(key, 0, -1);
+ List<byte[]> expected = Arrays.asList(new byte[] {1}, new byte[] {2});
+ assertEquals(expected.size(), values.size());
+ for (int i = 0; i < expected.size(); i++) {
+ assertArrayEquals(expected.get(i), values.get(i));
+ }
+ }
+
+ @Test
+ public void testWriteBytesWithMethodSAdd() {
+ byte[] key = {5};
+ List<byte[]> values =
+ Arrays.asList(
+ new byte[] {0},
+ new byte[] {1},
+ new byte[] {2},
+ new byte[] {3},
+ new byte[] {2},
+ new byte[] {4},
+ new byte[] {0},
+ new byte[] {5});
+ List<KV<byte[], byte[]>> data = buildConstantKeyList(key, values);
+
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(data));
+ write.apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.SADD));
+ p.run();
+
+ Set<byte[]> members = client.smembers(key);
+ assertEquals(6, members.size());
+ assertThat(members, hasItems(values.toArray(new byte[0][])));
+ }
+
+ @Test
+ public void testWriteBytesWithMethodPFAdd() {
+ byte[] key = {6};
+ List<byte[]> values =
+ Arrays.asList(
+ new byte[] {0},
+ new byte[] {1},
+ new byte[] {2},
+ new byte[] {3},
+ new byte[] {2},
+ new byte[] {4},
+ new byte[] {0},
+ new byte[] {5});
+ List<KV<byte[], byte[]>> data = buildConstantKeyList(key, values);
+
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(data));
+ write.apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.PFADD));
+ p.run();
+
+ long count = client.pfcount(key);
+ assertEquals(6, count);
+ assertEquals(NO_EXPIRATION, Long.valueOf(client.ttl(key)));
+ }
+
+ @Test
+ public void testWriteBytesWithMethodPFAddWithExpireTime() {
+ byte[] key = {7};
+ List<byte[]> values =
+ Arrays.asList(
+ new byte[] {0},
+ new byte[] {1},
+ new byte[] {2},
+ new byte[] {3},
+ new byte[] {2},
+ new byte[] {4},
+ new byte[] {0},
+ new byte[] {5});
+ List<KV<byte[], byte[]>> data = buildConstantKeyList(key, values);
+
+ PCollection<KV<byte[], byte[]>> write = p.apply(Create.of(data));
+ write.apply(
+ RedisIO.writeBytes()
+ .withEndpoint(redisHost, port)
+ .withMethod(Method.PFADD)
+ .withExpireTime(10_000L));
+ p.run();
+
+ long count = client.pfcount(key);
+ assertEquals(6, count);
+ Long expireTime = client.pttl(key);
+ assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <=
10_000);
+ client.del(key);
+ }
+
+ @Test
+ public void testWriteBytesUsingINCRBY() {
+ byte[] key = "key_incr_bytes".getBytes(StandardCharsets.UTF_8);
+ List<byte[]> values =
+ Arrays.asList(0L, 1L, 2L, -3L, 2L, 4L, 0L, 5L).stream()
+ .map(Longs::toByteArray)
+ .collect(Collectors.toList());
+ List<KV<byte[], byte[]>> data = buildConstantKeyList(key, values);
+
+ p.apply(Create.of(data))
+ .apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.INCRBY));
+
+ p.run();
+
+ byte[] response = client.get(key);
+ long count = Long.parseLong(new String(response, StandardCharsets.UTF_8));
+ assertEquals(11, count);
+ }
+
+ @Test
+ public void testWriteBytesUsingDECRBY() {
+ byte[] key = "key_decr_bytes".getBytes(StandardCharsets.UTF_8);
+
+ List<byte[]> values =
+ Arrays.asList(-10L, 1L, 2L, -3L, 2L, 4L, 0L, 5L).stream()
+ .map(Longs::toByteArray)
+ .collect(Collectors.toList());
+ List<KV<byte[], byte[]>> data = buildConstantKeyList(key, values);
+
+ p.apply(Create.of(data))
+ .apply(RedisIO.writeBytes().withEndpoint(redisHost,
port).withMethod(Method.DECRBY));
+
+ p.run();
+
+ byte[] response = client.get(key);
+ long count = Long.parseLong(new String(response, StandardCharsets.UTF_8));
+ assertEquals(-1, count);
+ }
+
@Test
public void testWriteStreams() {
@@ -336,6 +562,62 @@ public class RedisIOTest {
}
}
+ @Test
+ public void testWriteStreamsBytes() {
+
+ /* test data is 10 keys (stream IDs), each with two entries, each entry
having one k/v pair of data */
+ List<byte[]> redisKeys =
+ IntStream.range(0,
10).boxed().map(Ints::toByteArray).collect(Collectors.toList());
+
+ Map<byte[], byte[]> fooValues =
+ ImmutableMap.of(
+ "sensor-id".getBytes(StandardCharsets.UTF_8),
+ new byte[] {1, 2, 3, 4},
+ "temperature".getBytes(StandardCharsets.UTF_8),
+ new byte[] {19, 8});
+ Map<byte[], byte[]> barValues =
+ ImmutableMap.of(
+ "sensor-id".getBytes(StandardCharsets.UTF_8),
+ new byte[] {9, 9, 9, 9},
+ "temperature".getBytes(StandardCharsets.UTF_8),
+ new byte[] {18, 2});
+
+ List<KV<byte[], Map<byte[], byte[]>>> allData =
+ redisKeys.stream()
+ .flatMap(id -> Stream.of(KV.of(id, fooValues), KV.of(id,
barValues)))
+ .collect(toList());
+
+ PCollection<KV<byte[], Map<byte[], byte[]>>> write =
+ p.apply(
+ Create.of(allData)
+ .withCoder(
+ KvCoder.of(
+ ByteArrayCoder.of(),
+ MapCoder.of(ByteArrayCoder.of(),
ByteArrayCoder.of()))));
+ write.apply(RedisIO.writeStreamsBytes().withEndpoint(redisHost, port));
+ p.run();
+
+ for (byte[] key : redisKeys) {
+ String redisKey = new String(key, StandardCharsets.UTF_8);
+ List<StreamEntry> streamEntries =
+ client.xrange(redisKey, (StreamEntryID) null, (StreamEntryID) null,
Integer.MAX_VALUE);
+ assertEquals(2, streamEntries.size());
+ assertThat(
+ transform(streamEntries, StreamEntry::getFields),
+ hasItems(
+ ImmutableMap.of(
+ "sensor-id",
+ new String(new byte[] {9, 9, 9, 9}, StandardCharsets.UTF_8),
+ "temperature",
+ new String(new byte[] {18, 2}, StandardCharsets.UTF_8)),
+ ImmutableMap.of(
+ "sensor-id",
+ new String(new byte[] {1, 2, 3, 4}, StandardCharsets.UTF_8),
+ "temperature",
+ new String(new byte[] {19, 8}, StandardCharsets.UTF_8))));
+ }
+ }
+
@Test
public void testWriteStreamsWithTruncation() {
/* test data is 10 keys (stream IDs), each with two entries, each entry
having one k/v pair of data */
@@ -406,9 +688,9 @@ public class RedisIOTest {
assertEquals("1885267", redisCursor.getCursor());
}
- private static List<KV<String, String>> buildConstantKeyList(String key,
List<String> values) {
- List<KV<String, String>> data = new ArrayList<>();
- for (String value : values) {
+ private static <K, V> List<KV<K, V>> buildConstantKeyList(K key, List<V>
values) {
+ List<KV<K, V>> data = new ArrayList<>();
+ for (V value : values) {
data.add(KV.of(key, value));
}
return data;
