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 62109a6a6660..6b5cd99c69e4 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 @@ 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.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 @@ * * } * + *

{@link #writeBytes()} is similar to {@link #write()} but works with key/value pairs + * represented as byte arrays: + * + * 

{@code
+ * pipeline.apply(...)
+ *   // here we a have a PCollection with key/value pairs
+ *   .apply(RedisIO.writeBytes().withEndpoint("::1", 6379))
+ *
+ * }
+ * *

Writing Redis Streams

* *

{@link #writeStreams()} appends the entries of a {@link PCollection} of key/value pairs @@ -159,12 +174,32 @@ public static Read read() { .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 static ReadKeyPatterns readKeyPatterns() { /** 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 static Write write() { /** 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 PCollection> expand(PBegin input) { } } - /** Implementation of {@link #readKeyPatterns()}. */ + /** Implementation of {@link #readBytes()}. */ @AutoValue - public abstract static class ReadKeyPatterns - extends PTransform, PCollection>> { + public abstract static class ReadBytes + extends PTransform>> { 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> 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>> { + + 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> expand(PCollection input) { + + return input + .apply(readKeyPatternsBytes()) + .apply( + MapElements.via( + new SimpleFunction, KV>() { + @Override + public KV apply(KV 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>> { + + 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 ReadKeyPatterns withEndpoint(String host, int port) { .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 ReadKeyPatterns withBatchSize(int batchSize) { * 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> expand(PCollection input) { + public PCollection> expand(PCollection input) { checkArgument(connectionConfiguration() != null, "withConnectionConfiguration() is required"); - PCollection> output = + + PCollection> 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> { + private static class ReadFn extends DoFn> { protected final RedisConnectionConfiguration connectionConfiguration; transient Jedis jedis; @@ -394,10 +609,11 @@ public void processElement( ScanParams scanParams = new ScanParams(); scanParams.match(c.element()); while (tracker.tryClaim(cursor)) { - ScanResult scanResult = jedis.scan(redisCursor.getCursor(), scanParams); + byte[] cursorBytes = redisCursor.getCursor().getBytes(StandardCharsets.UTF_8); + ScanResult scanResult = jedis.scan(cursorBytes, scanParams); if (scanResult.getResult().size() > 0) { - String[] keys = scanResult.getResult().toArray(new String[scanResult.getResult().size()]); - List results = jedis.mget(keys); + byte[][] keys = scanResult.getResult().toArray(new byte[scanResult.getResult().size()][]); + List 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 void processElement( } private static class Reparallelize - extends PTransform>, PCollection>> { + extends PTransform>, PCollection>> { @Override - public PCollection> expand(PCollection> input) { + public PCollection> expand(PCollection> input) { // reparallelize mimics the same behavior as in JdbcIO, used to break fusion - PCollectionView>> empty = + PCollectionView>> empty = input .apply("Consume", Filter.by(SerializableFunctions.constant(false))) .apply(View.asIterable()); - PCollection> materialized = + PCollection> materialized = input.apply( "Identity", ParDo.of( - new DoFn, KV>() { + new DoFn, KV>() { @ProcessElement public void processElement(ProcessContext c) { c.output(c.element()); @@ -441,6 +657,66 @@ public void processElement(ProcessContext c) { @AutoValue public abstract static class Write extends PTransform>, 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> input) { + MapElements, KV> toBytes; + Method method = writeBytes().method(); + if (Method.INCRBY == method || Method.DECRBY == method) { + toBytes = + MapElements.via( + new SimpleFunction, KV>() { + @Override + public KV apply(KV 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>() { + @Override + public KV apply(KV 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 enum Method { 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>, PDone> { + abstract @Nullable RedisConnectionConfiguration connectionConfiguration(); - abstract @Nullable Method method(); + abstract Write.@Nullable Method method(); abstract @Nullable Long expireTime(); @@ -494,14 +784,14 @@ abstract static class Builder { 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 Write withEndpoint(String host, int port) { .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> input) { + public PDone expand(PCollection> 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, Void> { + private static class WriteFn extends DoFn, 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 void startBundle() { @ProcessElement public void processElement(ProcessContext c) { - KV record = c.element(); + KV record = c.element(); writeRecord(record); @@ -588,39 +878,39 @@ public void processElement(ProcessContext c) { } } - private void writeRecord(KV record) { - Method method = spec.method(); + private void writeRecord(KV 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 record, Long expireTime) { - String key = record.getKey(); - String value = record.getValue(); + private void writeUsingAppendCommand(KV record, Long expireTime) { + byte[] key = record.getKey(); + byte[] value = record.getValue(); transaction.append(key, value); setExpireTimeWhenRequired(key, expireTime); } - private void writeUsingSetCommand(KV record, Long expireTime) { - String key = record.getKey(); - String value = record.getValue(); + private void writeUsingSetCommand(KV record, Long expireTime) { + byte[] key = record.getKey(); + byte[] value = record.getValue(); if (expireTime != null) { transaction.psetex(key, expireTime, value); @@ -630,57 +920,57 @@ private void writeUsingSetCommand(KV record, Long expireTime) { } private void writeUsingListCommand( - KV record, Method method, Long expireTime) { + KV 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 record, Long expireTime) { - String key = record.getKey(); - String value = record.getValue(); + private void writeUsingSaddCommand(KV record, Long expireTime) { + byte[] key = record.getKey(); + byte[] value = record.getValue(); transaction.sadd(key, value); setExpireTimeWhenRequired(key, expireTime); } - private void writeUsingHLLCommand(KV record, Long expireTime) { - String key = record.getKey(); - String value = record.getValue(); + private void writeUsingHLLCommand(KV record, Long expireTime) { + byte[] key = record.getKey(); + byte[] value = record.getValue(); transaction.pfadd(key, value); setExpireTimeWhenRequired(key, expireTime); } - private void writeUsingIncrBy(KV record, Long expireTime) { - String key = record.getKey(); - String value = record.getValue(); - long inc = Long.parseLong(value); + private void writeUsingIncrBy(KV 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 record, Long expireTime) { - String key = record.getKey(); - String value = record.getValue(); - long decr = Long.parseLong(value); + private void writeUsingDecrBy(KV 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 void teardown() { } } + @AutoValue + public abstract static class WriteStreams + extends PTransform>>, 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>> 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>> { + + @ProcessElement + public void processElement( + @Element KV> element, + OutputReceiver>> out) { + byte[] key = element.getKey().getBytes(StandardCharsets.UTF_8); + Map 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>>, PDone> { + public abstract static class WriteStreamsBytes + extends PTransform>>, PDone> { abstract RedisConnectionConfiguration connectionConfiguration(); @@ -722,7 +1076,7 @@ public abstract static class WriteStreams 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 WriteStreams withEndpoint(String host, int port) { * 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 WriteStreams withAuth(String auth) { } /** 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 WriteStreams withTimeout(int timeout) { } /** 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 WriteStreams withMaxLen(long maxLen) { * 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>> 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 @@ abstract Builder setConnectionConfiguration( abstract Builder setApproximateTrim(boolean approximateTrim); - abstract WriteStreams build(); + abstract WriteStreamsBytes build(); } - @Override - public PDone expand(PCollection>> 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>, Void> { + private static class WriteStreamFn extends DoFn>, 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 void startBundle() { @ProcessElement public void processElement(ProcessContext c) { - KV> record = c.element(); + KV> record = c.element(); writeRecord(record); @@ -832,9 +1185,7 @@ public void processElement(ProcessContext c) { } } - private void writeRecord(KV> record) { - String key = record.getKey(); - Map value = record.getValue(); + private void writeRecord(KV> record) { final XAddParams params = new XAddParams().id(StreamEntryID.NEW_ENTRY); if (spec.maxLen() > 0L) { params.maxLen(spec.maxLen()); @@ -842,7 +1193,7 @@ private void writeRecord(KV> record) { 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 5754de89ec3a..0037999d722c 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 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.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; @@ -46,6 +50,8 @@ 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.Rule; @@ -151,6 +157,36 @@ public void testReadWithKeyPattern() { 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> read = + p.apply( + "Read", + RedisIO.readBytes().withEndpoint(REDIS_HOST, port).withKeyPattern("binaryread*")); + + // Verify actual byte content + PAssert.that(read) + .satisfies( + results -> { + Set seen = new HashSet<>(); + for (KV 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"; @@ -181,7 +217,7 @@ public void testWriteWithMethodSetWithExpiration() { 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); } @@ -255,14 +291,14 @@ public void testWriteWithMethodPFAddWithExpireTime() { write.apply( RedisIO.write() .withEndpoint(REDIS_HOST, port) - .withMethod(Method.PFADD) + .withMethod(RedisIO.Write.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_0000); + assertTrue(expireTime.toString(), 9_000 <= expireTime && expireTime <= 10_000); client.del(key); } @@ -297,6 +333,196 @@ public void testWriteUsingDECRBY() { assertEquals(-1, count); } + @Test + public void testWriteBytesWithMethodSet() { + byte[] key = {1}; + client.set(key, new byte[] {1}); + + byte[] newValue = {2}; + PCollection> write = p.apply(Create.of(KV.of(key, newValue))); + write.apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, 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> write = p.apply(Create.of(KV.of(key, newValue))); + write.apply( + RedisIO.writeBytes() + .withEndpoint(REDIS_HOST, 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> write = p.apply(Create.of(KV.of(key, newValue))); + write.apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, port).withMethod(Method.LPUSH)); + p.run(); + + List values = client.lrange(key, 0, -1); + List 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> write = p.apply(Create.of(KV.of(key, newValue))); + write.apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, port).withMethod(Method.RPUSH)); + p.run(); + + List values = client.lrange(key, 0, -1); + List 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 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> data = buildConstantKeyList(key, values); + + PCollection> write = p.apply(Create.of(data)); + write.apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, port).withMethod(Method.SADD)); + p.run(); + + Set members = client.smembers(key); + assertEquals(6, members.size()); + assertThat(members, hasItems(values.toArray(new byte[0][]))); + } + + @Test + public void testWriteBytesWithMethodPFAdd() { + byte[] key = {6}; + List 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> data = buildConstantKeyList(key, values); + + PCollection> write = p.apply(Create.of(data)); + write.apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, 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 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> data = buildConstantKeyList(key, values); + + PCollection> write = p.apply(Create.of(data)); + write.apply( + RedisIO.writeBytes() + .withEndpoint(REDIS_HOST, 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 values = + Arrays.asList(0L, 1L, 2L, -3L, 2L, 4L, 0L, 5L).stream() + .map(Longs::toByteArray) + .collect(Collectors.toList()); + List> data = buildConstantKeyList(key, values); + + p.apply(Create.of(data)) + .apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, 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 values = + Arrays.asList(-10L, 1L, 2L, -3L, 2L, 4L, 0L, 5L).stream() + .map(Longs::toByteArray) + .collect(Collectors.toList()); + List> data = buildConstantKeyList(key, values); + + p.apply(Create.of(data)) + .apply(RedisIO.writeBytes().withEndpoint(REDIS_HOST, 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() { @@ -330,6 +556,62 @@ public void testWriteStreams() { } } + @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 redisKeys = + IntStream.range(0, 10).boxed().map(Ints::toByteArray).collect(Collectors.toList()); + + Map 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 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>> allData = + redisKeys.stream() + .flatMap(id -> Stream.of(KV.of(id, fooValues), KV.of(id, barValues))) + .collect(toList()); + + PCollection>> write = + p.apply( + Create.of(allData) + .withCoder( + KvCoder.of( + ByteArrayCoder.of(), + MapCoder.of(ByteArrayCoder.of(), ByteArrayCoder.of())))); + write.apply(RedisIO.writeStreamsBytes().withEndpoint(REDIS_HOST, port)); + p.run(); + + for (byte[] key : redisKeys) { + String redisKey = new String(key, StandardCharsets.UTF_8); + List 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 */ @@ -400,9 +682,9 @@ public void redisByteKeyToRedisCursor() { assertEquals("1885267", redisCursor.getCursor()); } - private static List> buildConstantKeyList(String key, List values) { - List> data = new ArrayList<>(); - for (String value : values) { + private static List> buildConstantKeyList(K key, List values) { + List> data = new ArrayList<>(); + for (V value : values) { data.add(KV.of(key, value)); } return data;