Skip to content

Add Drain mode to PaneInfo #34530

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kennknowles wants to merge 1 commit into from
Closed

Add Drain mode to PaneInfo #34530

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 92 additions & 13 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/windowing/PaneInfo.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.apache.beam.sdk.transforms.DoFn.WindowedContext;
import org.apache.beam.sdk.transforms.GroupByKey;
import org.apache.beam.sdk.util.VarInt;
import org.apache.beam.sdk.values.DrainMode;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
Expand Down Expand Up @@ -140,16 +141,32 @@ private static byte encodedByte(boolean isFirst, boolean isLast, Timing timing)
return result;
}

private static byte encodeExtendedMetadataByte(DrainMode drainMode) {
return (byte) (drainMode.ordinal() << 6);
}

private static DrainMode drainModeFromExtendedMetadata(byte extendedMetadataByte) {
return DrainMode.values()[(extendedMetadataByte >>> 6) & 0b11];
}

private static final ImmutableMap<Byte, PaneInfo> BYTE_TO_PANE_INFO;

static {
ImmutableMap.Builder<Byte, PaneInfo> decodingBuilder = ImmutableMap.builder();
for (Timing timing : Timing.values()) {
long onTimeIndex = timing == Timing.EARLY ? -1 : 0;
register(decodingBuilder, new PaneInfo(true, true, timing, 0, onTimeIndex));
register(decodingBuilder, new PaneInfo(true, false, timing, 0, onTimeIndex));
register(decodingBuilder, new PaneInfo(false, true, timing, -1, onTimeIndex));
register(decodingBuilder, new PaneInfo(false, false, timing, -1, onTimeIndex));
register(
decodingBuilder,
new PaneInfo(true, true, timing, DrainMode.NOT_DRAINING, 0, onTimeIndex));
register(
decodingBuilder,
new PaneInfo(true, false, timing, DrainMode.NOT_DRAINING, 0, onTimeIndex));
register(
decodingBuilder,
new PaneInfo(false, true, timing, DrainMode.NOT_DRAINING, -1, onTimeIndex));
register(
decodingBuilder,
new PaneInfo(false, false, timing, DrainMode.NOT_DRAINING, -1, onTimeIndex));
}
BYTE_TO_PANE_INFO = decodingBuilder.build();
}
Expand All @@ -158,11 +175,17 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
builder.put(info.encodedByte, info);
}

// Byte containing the encoding tag, timing, isFirst, and isLast.
// These bytes are cached for all known possibilities to avoid repeat processing
private final byte encodedByte;

// Extended metadata byte, containing drain mode
private final byte extendedMetadataByte;

private final boolean isFirst;
private final boolean isLast;
private final Timing timing;
private final DrainMode drainMode;
private final long index;
private final long nonSpeculativeIndex;

Expand All @@ -177,18 +200,28 @@ private static void register(ImmutableMap.Builder<Byte, PaneInfo> builder, PaneI
public static final PaneInfo ON_TIME_AND_ONLY_FIRING =
PaneInfo.createPane(true, true, Timing.ON_TIME, 0, 0);

private PaneInfo(boolean isFirst, boolean isLast, Timing timing, long index, long onTimeIndex) {
private PaneInfo(
boolean isFirst,
boolean isLast,
Timing timing,
DrainMode drainMode,
long index,
long onTimeIndex) {
this.encodedByte = encodedByte(isFirst, isLast, timing);
this.extendedMetadataByte =
drainMode == DrainMode.DRAINING ? encodeExtendedMetadataByte(drainMode) : 0x00;
this.isFirst = isFirst;
this.isLast = isLast;
this.timing = timing;
this.drainMode = drainMode;
this.index = index;
this.nonSpeculativeIndex = onTimeIndex;
}

public static PaneInfo createPane(boolean isFirst, boolean isLast, Timing timing) {
checkArgument(isFirst, "Indices must be provided for non-first pane info.");
return createPane(isFirst, isLast, timing, 0, timing == Timing.EARLY ? -1 : 0);
return createPane(
isFirst, isLast, timing, DrainMode.NOT_DRAINING, 0, timing == Timing.EARLY ? -1 : 0);
}

/** Factory method to create a {@link PaneInfo} with the specified parameters. */
Expand All @@ -197,7 +230,22 @@ public static PaneInfo createPane(
if (isFirst || timing == Timing.UNKNOWN) {
return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing)));
} else {
return new PaneInfo(isFirst, isLast, timing, index, onTimeIndex);
return new PaneInfo(isFirst, isLast, timing, DrainMode.NOT_DRAINING, index, onTimeIndex);
}
}

/** Factory method to create a {@link PaneInfo} with the specified parameters. */
public static PaneInfo createPane(
boolean isFirst,
boolean isLast,
Timing timing,
DrainMode drainMode,
long index,
long onTimeIndex) {
if (drainMode != DrainMode.DRAINING && (isFirst || timing == Timing.UNKNOWN)) {
return checkNotNull(BYTE_TO_PANE_INFO.get(encodedByte(isFirst, isLast, timing)));
} else {
return new PaneInfo(isFirst, isLast, timing, drainMode, index, onTimeIndex);
}
}

Expand Down Expand Up @@ -241,6 +289,14 @@ public long getIndex() {
return index;
}

/**
* Indicates whether this element resulted from an aggregation that fired during a drain
* operation.
*/
public DrainMode getDrainMode() {
return drainMode;
}

/**
* The zero-based index of this trigger firing among non-speculative panes.
*
Expand All @@ -253,13 +309,17 @@ public long getNonSpeculativeIndex() {
return nonSpeculativeIndex;
}

int getEncodedByte() {
byte getEncodedByte() {
return encodedByte;
}

byte getExtendedMetadataByte() {
return extendedMetadataByte;
}

@Override
public int hashCode() {
return Objects.hash(encodedByte, index, nonSpeculativeIndex);
return Objects.hash(encodedByte, extendedMetadataByte, index, nonSpeculativeIndex);
}

@Override
Expand All @@ -270,6 +330,7 @@ public boolean equals(@Nullable Object obj) {
} else if (obj instanceof PaneInfo) {
PaneInfo that = (PaneInfo) obj;
return this.encodedByte == that.encodedByte
&& this.extendedMetadataByte == that.extendedMetadataByte
&& this.index == that.index
&& this.nonSpeculativeIndex == that.nonSpeculativeIndex;
} else {
Expand All @@ -290,6 +351,7 @@ public String toString() {
.add("timing", timing)
.add("index", index)
.add("onTimeIndex", nonSpeculativeIndex != -1 ? nonSpeculativeIndex : null)
.add("drainMode", drainMode == DrainMode.DRAINING ? "DRAINING" : null)
.toString();
}

Expand All @@ -298,7 +360,8 @@ public static class PaneInfoCoder extends AtomicCoder<PaneInfo> {
private enum Encoding {
FIRST,
ONE_INDEX,
TWO_INDICES;
TWO_INDICES,
EXTENDED_METADATA;

// NOTE: Do not reorder fields. The ordinal is used as part of
// the encoding.
Expand All @@ -311,12 +374,15 @@ private enum Encoding {
}

public static Encoding fromTag(byte b) {
return Encoding.values()[b >> 4];
return Encoding.values()[(b >>> 4) & 0xF];
}
}

private Encoding chooseEncoding(PaneInfo value) {
if ((value.index == 0 && value.nonSpeculativeIndex == 0) || value.timing == Timing.UNKNOWN) {
if (value.drainMode == DrainMode.DRAINING) {
return Encoding.EXTENDED_METADATA;
} else if ((value.index == 0 && value.nonSpeculativeIndex == 0)
|| value.timing == Timing.UNKNOWN) {
return Encoding.FIRST;
} else if (value.index == value.nonSpeculativeIndex || value.timing == Timing.EARLY) {
return Encoding.ONE_INDEX;
Expand Down Expand Up @@ -350,6 +416,12 @@ public void encode(PaneInfo value, final OutputStream outStream)
VarInt.encode(value.index, outStream);
VarInt.encode(value.nonSpeculativeIndex, outStream);
break;
case EXTENDED_METADATA:
outStream.write(value.encodedByte | encoding.tag);
outStream.write(value.extendedMetadataByte);
VarInt.encode(value.index, outStream);
VarInt.encode(value.nonSpeculativeIndex, outStream);
break;
default:
throw new CoderException("Unknown encoding " + encoding);
}
Expand All @@ -360,6 +432,7 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
byte keyAndTag = (byte) inStream.read();
PaneInfo base = Preconditions.checkNotNull(BYTE_TO_PANE_INFO.get((byte) (keyAndTag & 0x0F)));
long index, onTimeIndex;
DrainMode drainMode = DrainMode.NOT_DRAINING;
switch (Encoding.fromTag(keyAndTag)) {
case FIRST:
return base;
Expand All @@ -371,10 +444,16 @@ public PaneInfo decode(final InputStream inStream) throws CoderException, IOExce
index = VarInt.decodeLong(inStream);
onTimeIndex = VarInt.decodeLong(inStream);
break;
case EXTENDED_METADATA:
byte extendedMetadata = (byte) inStream.read();
drainMode = drainModeFromExtendedMetadata(extendedMetadata);
index = VarInt.decodeLong(inStream);
onTimeIndex = VarInt.decodeLong(inStream);
break;
default:
throw new CoderException("Unknown encoding " + (keyAndTag & 0xF0));
}
return new PaneInfo(base.isFirst, base.isLast, base.timing, index, onTimeIndex);
return new PaneInfo(base.isFirst, base.isLast, base.timing, drainMode, index, onTimeIndex);
}

@Override
Expand Down
24 changes: 24 additions & 0 deletions sdks/java/core/src/main/java/org/apache/beam/sdk/values/DrainMode.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.beam.sdk.values;

/** {@link DrainMode} indicates whether a pipeline is known to be draining or not. */
public enum DrainMode {
NOT_DRAINING,
DRAINING;
}
36 changes: 27 additions & 9 deletions sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/windowing/PaneInfoTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.testing.CoderProperties;
import org.apache.beam.sdk.transforms.windowing.PaneInfo.Timing;
import org.apache.beam.sdk.values.DrainMode;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
Expand All @@ -41,19 +42,23 @@ public void testInterned() throws Exception {
@Test
public void testEncodingRoundTrip() throws Exception {
Coder<PaneInfo> coder = PaneInfo.PaneInfoCoder.INSTANCE;
for (Timing timing : Timing.values()) {
long onTimeIndex = timing == Timing.EARLY ? -1 : 37;
CoderProperties.coderDecodeEncodeEqual(
coder, PaneInfo.createPane(false, false, timing, 389, onTimeIndex));
CoderProperties.coderDecodeEncodeEqual(
coder, PaneInfo.createPane(false, true, timing, 5077, onTimeIndex));
CoderProperties.coderDecodeEncodeEqual(coder, PaneInfo.createPane(true, false, timing, 0, 0));
CoderProperties.coderDecodeEncodeEqual(coder, PaneInfo.createPane(true, true, timing, 0, 0));
for (DrainMode drainMode : DrainMode.values()) {
for (Timing timing : Timing.values()) {
long onTimeIndex = timing == Timing.EARLY ? -1 : 37;
CoderProperties.coderDecodeEncodeEqual(
coder, PaneInfo.createPane(false, false, timing, drainMode, 389, onTimeIndex));
CoderProperties.coderDecodeEncodeEqual(
coder, PaneInfo.createPane(false, true, timing, drainMode, 5077, onTimeIndex));
CoderProperties.coderDecodeEncodeEqual(
coder, PaneInfo.createPane(true, false, timing, drainMode, 0, 0));
CoderProperties.coderDecodeEncodeEqual(
coder, PaneInfo.createPane(true, true, timing, drainMode, 0, 0));
}
}
}

@Test
public void testEncodings() {
public void testHeaderByteEncoding() {
assertEquals(
"PaneInfo encoding assumes that there are only 4 Timing values.",
4,
Expand Down Expand Up @@ -83,4 +88,17 @@ public void testEncodings() {
0xF,
PaneInfo.createPane(true, true, Timing.UNKNOWN).getEncodedByte());
}

/**
* NOTE: this does not reproduce the logic in the main code by doing exhaustive testing
* programmatically. It is deliberately a smoke test of constant values.
*/
@Test
public void testExtendedMetadataByteEncoding() {
assertEquals(
"PaneInfo with drainMode should have known extended metadata byte",
(byte) 0b01000000,
PaneInfo.createPane(true, true, Timing.UNKNOWN, DrainMode.DRAINING, 1, 1)
.getExtendedMetadataByte());
}
}
Loading