damccorm commented on code in PR #40609: URL: https://github.com/apache/beam/pull/40609#discussion_r4287464538
########## runners/flink/src/main/java/org/apache/beam/runners/flink/translation/types/LengthPrefixedCoderTypeSerializer.java: ########## @@ -0,0 +1,217 @@ +/* + * 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.runners.flink.translation.types; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.EOFException; +import java.io.IOException; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.CoderException; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.coders.VoidCoder; +import org.apache.beam.sdk.values.WindowedValues; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.apache.flink.core.memory.DataInputView; +import org.apache.flink.core.memory.DataOutputView; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A {@link CoderTypeSerializer} that prefixes every record with its encoded length. + * + * <p>Beam coders do not delimit their own output, so {@link CoderTypeSerializer#copy(DataInputView, + * DataOutputView)} has to decode and re-encode a record to find where it ends. Flink calls that + * method for every record when its batch sorter spills to disk. With a length prefix, copying is a + * plain byte copy. + * + * <p>The format differs from {@link CoderTypeSerializer}, so this serializer must only be used for + * data that is never persisted across job versions, i.e. in batch execution. + * + * <p>With {@code lazyKvValues}, records are deserialized as {@link LazyKvWindowedValue}s. This + * requires a {@link WindowedValues.FullWindowedValueCoder} of a {@link KvCoder}. Records whose + * value was never accessed are serialized again by copying their original bytes. + */ +@SuppressWarnings({ + "rawtypes", // TODO(https://github.com/apache/beam/issues/20447) + "nullness", // TODO(https://github.com/apache/beam/issues/20497) + "unchecked" +}) +public class LengthPrefixedCoderTypeSerializer<T> extends CoderTypeSerializer<T> { + + private static final long serialVersionUID = 1L; + + /** Scratch buffers above this size are released after use instead of being kept around. */ + private static final int MAX_RETAINED_BUFFER_SIZE = 1 << 20; + + private final Coder<T> coder; + + private final boolean fasterCopy; + + private final boolean lazyKvValues; + + /** Decodes the windowing information of a record, see {@link LazyKvWindowedValue}. */ + private final WindowedValues.@Nullable WindowedValueCoder<Void> headerCoder; + + private final @Nullable KvCoder<?, ?> kvCoder; + + private transient @Nullable ExposedByteArrayOutputStream encodeBuffer; + + private transient byte @Nullable [] decodeBuffer; + + public LengthPrefixedCoderTypeSerializer(Coder<T> coder, boolean fasterCopy) { + this(coder, fasterCopy, false); + } + + public LengthPrefixedCoderTypeSerializer( + Coder<T> coder, boolean fasterCopy, boolean lazyKvValues) { + super(coder, fasterCopy); + this.coder = coder; + this.fasterCopy = fasterCopy; + this.lazyKvValues = lazyKvValues; + if (lazyKvValues) { + Preconditions.checkArgument( + supportsLazyKvValues(coder), + "Lazy values require a FullWindowedValueCoder of a KvCoder, got %s", + coder); + WindowedValues.FullWindowedValueCoder<?> windowedCoder = + (WindowedValues.FullWindowedValueCoder<?>) coder; + this.headerCoder = windowedCoder.withValueCoder(VoidCoder.of()); + this.kvCoder = (KvCoder<?, ?>) windowedCoder.getValueCoder(); + } else { + this.headerCoder = null; + this.kvCoder = null; + } + } + + /** Whether {@code coder} can be used with lazily decoded values. */ + public static boolean supportsLazyKvValues(Coder<?> coder) { + return coder instanceof WindowedValues.FullWindowedValueCoder + && ((WindowedValues.FullWindowedValueCoder<?>) coder).getValueCoder() instanceof KvCoder; + } + + @Override + public LengthPrefixedCoderTypeSerializer<T> duplicate() { + return new LengthPrefixedCoderTypeSerializer<>(coder, fasterCopy, lazyKvValues); + } + + @Override + public void serialize(T t, DataOutputView dataOutputView) throws IOException { + byte[] encodedRecord = + t instanceof LazyKvWindowedValue ? ((LazyKvWindowedValue<?, ?>) t).encodedRecord() : null; + if (encodedRecord != null) { + // The value was never decoded, so the original bytes are still valid. + dataOutputView.writeInt(encodedRecord.length); + dataOutputView.write(encodedRecord); + return; + } + ExposedByteArrayOutputStream buffer = encodeBuffer; + if (buffer == null) { + buffer = new ExposedByteArrayOutputStream(); + encodeBuffer = buffer; + } + buffer.reset(); + coder.encode(t, buffer); + int length = buffer.size(); + dataOutputView.writeInt(length); + dataOutputView.write(buffer.buffer(), 0, length); + if (buffer.capacity() > MAX_RETAINED_BUFFER_SIZE) { + encodeBuffer = null; + } + } + + @Override + public T deserialize(DataInputView dataInputView) throws IOException { + int length = dataInputView.readInt(); + if (headerCoder != null) { + // The lazy value keeps the record, so it needs its own array. + byte[] encodedRecord = new byte[length]; + dataInputView.readFully(encodedRecord); + return (T) new LazyKvWindowedValue<>(encodedRecord, headerCoder, kvCoder); + } + byte[] buffer = decodeBuffer; + if (buffer == null || buffer.length < length) { + buffer = new byte[Math.max(length, 64)]; + } + dataInputView.readFully(buffer, 0, length); + decodeBuffer = buffer.length > MAX_RETAINED_BUFFER_SIZE ? null : buffer; + try { + return coder.decode(new ByteArrayInputStream(buffer, 0, length)); + } catch (CoderException e) { + Throwable cause = e.getCause(); + if (cause instanceof EOFException) { + throw (EOFException) cause; + } else { + throw e; + } + } + } + + @Override + public void copy(DataInputView dataInputView, DataOutputView dataOutputView) throws IOException { + int length = dataInputView.readInt(); + dataOutputView.writeInt(length); + dataOutputView.write(dataInputView, length); + } Review Comment: `LengthPrefixedCoderTypeSerializer` overrides `copy(DataInputView, DataOutputView)`, but inherits the in-memory `copy(T t)` / `copy(T t, T reuse)` methods from `CoderTypeSerializer`. When `fasterCopy` is `false` (the default), `CoderTypeSerializer#copy(T t)` runs `CoderUtils.clone(coder, t)`. Calling `CoderUtils.clone(coder, t)` on an untouched `LazyKvWindowedValue` invokes `t.getValue()` during encoding—which eagerly decodes the `KV`, **nulls out `t.encodedRecord` on the input instance `t`**, re-encodes the full record into a byte array, and decodes it back into an eager `WindowedValue`. Since `encodedRecord` is never mutated in place, should we override `copy(T t)` here to preserve laziness without decoding or re-encoding? ```java @Override public T copy(T t) { if (fasterCopy) { return t; } byte[] encodedRecord = t instanceof LazyKvWindowedValue ? ((LazyKvWindowedValue<?, ?>) t).encodedRecord() : null; if (encodedRecord != null && headerCoder != null && kvCoder != null) { return (T) new LazyKvWindowedValue<>(encodedRecord, headerCoder, kvCoder); } return super.copy(t); } ``` ########## runners/flink/src/main/java/org/apache/beam/runners/flink/translation/types/LazyKvWindowedValue.java: ########## @@ -0,0 +1,213 @@ +/* + * 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.runners.flink.translation.types; + +import io.opentelemetry.context.Context; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; +import org.apache.beam.runners.flink.translation.wrappers.streaming.FlinkKeyUtils; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.CausedByDrain; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link WindowedValue} of a {@link KV} that is kept in its encoded form until it is accessed. + * + * <p>Created by {@link LengthPrefixedCoderTypeSerializer} for records that Flink only moves around, + * e.g. between the network, its sorter and spill files. While the value is untouched, the record is + * serialized again by writing back its original bytes. The windowing information is only decoded + * when it is accessed, and the grouping key is read from the encoded record without decoding the + * value. + * + * <p>The first call to {@link #getValue()} decodes the value and drops the encoded record, as the + * caller may mutate the returned value. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public final class LazyKvWindowedValue<K, V> implements WindowedValue<KV<K, V>> { + + /** Decodes the windowing information of a record, and none of its value. */ + private final WindowedValues.WindowedValueCoder<Void> headerCoder; + + private final KvCoder<K, V> kvCoder; + + /** The encoded record, or null once the value has been decoded. */ + private byte @Nullable [] encodedRecord; + + /** The windowing information and metadata, with a null value. Decoded on first access. */ + private @Nullable WindowedValue<Void> header; + + /** + * Offset of the encoded {@link KV} in {@link #encodedRecord}, valid once the header is decoded. + */ + private int valueOffset; + + private @Nullable KV<K, V> value; + + LazyKvWindowedValue( + byte[] encodedRecord, + WindowedValues.WindowedValueCoder<Void> headerCoder, + KvCoder<K, V> kvCoder) { + this.encodedRecord = encodedRecord; + this.headerCoder = headerCoder; + this.kvCoder = kvCoder; + } + + /** The encoded record if the value was never decoded, null otherwise. */ + byte @Nullable [] encodedRecord() { + return encodedRecord; + } + + private WindowedValue<Void> header() { + WindowedValue<Void> h = header; + if (h == null) { + byte[] bytes = encodedRecord; + ByteArrayInputStream in = new ByteArrayInputStream(bytes); + try { + // The value coder is a VoidCoder, so this stops right before the encoded KV. + h = headerCoder.decode(in); + } catch (IOException e) { + throw new RuntimeException("Failed to decode windowing information", e); + } + valueOffset = bytes.length - in.available(); + header = h; + } + return h; + } + + /** The key, encoded as by {@link FlinkKeyUtils#encodeKey}. */ + public byte[] getEncodedKey() { + if (encodedRecord == null) { + return FlinkKeyUtils.encodeKey(value.getKey(), kvCoder.getKeyCoder()).array(); + } + header(); + byte[] bytes = encodedRecord; + // KvCoder encodes the key first, the same way as FlinkKeyUtils.encodeKey. + ByteArrayInputStream in = new ByteArrayInputStream(bytes, valueOffset, bytes.length); Review Comment: The 3rd parameter to `ByteArrayInputStream(byte[] buf, int offset, int length)` is `length`, not `endIndex` (it only works because `ByteArrayInputStream` internally clamps `Math.min(offset + length, buf.length)`). For consistency with `getValue()` on line 129: ```suggestion ByteArrayInputStream in = new ByteArrayInputStream(bytes, valueOffset, bytes.length - valueOffset); ``` ########## runners/flink/src/main/java/org/apache/beam/runners/flink/FlinkPipelineOptions.java: ########## @@ -367,6 +367,26 @@ public Long create(PipelineOptions options) { void setFasterCopy(Boolean fasterCopy); + @Description( + "Batch only. Prefix every record exchanged or sorted by Flink with its encoded length, so " + + "that Flink can copy records (e.g. when spilling sorted data to disk) without " + + "decoding and re-encoding them. Has no effect in streaming mode. " + + "Enabled by default; set to false to disable.") + @Default.Boolean(true) + Boolean getLengthPrefixBatchRecords(); + + void setLengthPrefixBatchRecords(Boolean lengthPrefixBatchRecords); + + @Description( + "Batch only. Decode the values of records sent to GroupByKey and Combine operators lazily, " + + "so that records are passed from the network to Flink's sorter without decoding and " + + "re-encoding their values. Has no effect in streaming mode. " + + "Enabled by default; set to false to disable.") + @Default.Boolean(true) + Boolean getLazyBatchGroupingInputs(); Review Comment: Nit: In Flink 1.x (`FlinkPipelineExecutionEnvironment`), non-portable batch pipelines only use `FlinkStreamingPipelineTranslator` (and therefore `FlinkStreamingAggregationsTranslators` where `.withLazyKvValues()` is wired) when `--useDataStreamForBatch=true` (otherwise they use `FlinkBatchPipelineTranslator` / `DataSet`). If intentional, it might be worth noting in the Flink 1.x `@Description` here that this applies when using `DataStream` batch mode (`--useDataStreamForBatch`). ########## runners/flink/src/main/java/org/apache/beam/runners/flink/translation/types/LazyKvWindowedValue.java: ########## @@ -0,0 +1,213 @@ +/* + * 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.runners.flink.translation.types; + +import io.opentelemetry.context.Context; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.util.Arrays; +import java.util.Collection; +import org.apache.beam.runners.flink.translation.wrappers.streaming.FlinkKeyUtils; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.transforms.windowing.PaneInfo; +import org.apache.beam.sdk.values.CausedByDrain; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.ValueKind; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowedValues; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.joda.time.Instant; + +/** + * A {@link WindowedValue} of a {@link KV} that is kept in its encoded form until it is accessed. + * + * <p>Created by {@link LengthPrefixedCoderTypeSerializer} for records that Flink only moves around, + * e.g. between the network, its sorter and spill files. While the value is untouched, the record is + * serialized again by writing back its original bytes. The windowing information is only decoded + * when it is accessed, and the grouping key is read from the encoded record without decoding the + * value. + * + * <p>The first call to {@link #getValue()} decodes the value and drops the encoded record, as the + * caller may mutate the returned value. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public final class LazyKvWindowedValue<K, V> implements WindowedValue<KV<K, V>> { + + /** Decodes the windowing information of a record, and none of its value. */ + private final WindowedValues.WindowedValueCoder<Void> headerCoder; + + private final KvCoder<K, V> kvCoder; + + /** The encoded record, or null once the value has been decoded. */ + private byte @Nullable [] encodedRecord; + + /** The windowing information and metadata, with a null value. Decoded on first access. */ + private @Nullable WindowedValue<Void> header; + + /** + * Offset of the encoded {@link KV} in {@link #encodedRecord}, valid once the header is decoded. + */ + private int valueOffset; + + private @Nullable KV<K, V> value; + + LazyKvWindowedValue( + byte[] encodedRecord, + WindowedValues.WindowedValueCoder<Void> headerCoder, + KvCoder<K, V> kvCoder) { + this.encodedRecord = encodedRecord; + this.headerCoder = headerCoder; + this.kvCoder = kvCoder; + } + + /** The encoded record if the value was never decoded, null otherwise. */ + byte @Nullable [] encodedRecord() { + return encodedRecord; + } + + private WindowedValue<Void> header() { + WindowedValue<Void> h = header; + if (h == null) { + byte[] bytes = encodedRecord; + ByteArrayInputStream in = new ByteArrayInputStream(bytes); + try { + // The value coder is a VoidCoder, so this stops right before the encoded KV. + h = headerCoder.decode(in); + } catch (IOException e) { + throw new RuntimeException("Failed to decode windowing information", e); + } + valueOffset = bytes.length - in.available(); + header = h; + } + return h; + } + + /** The key, encoded as by {@link FlinkKeyUtils#encodeKey}. */ + public byte[] getEncodedKey() { + if (encodedRecord == null) { + return FlinkKeyUtils.encodeKey(value.getKey(), kvCoder.getKeyCoder()).array(); + } + header(); + byte[] bytes = encodedRecord; + // KvCoder encodes the key first, the same way as FlinkKeyUtils.encodeKey. + ByteArrayInputStream in = new ByteArrayInputStream(bytes, valueOffset, bytes.length); + try { + kvCoder.getKeyCoder().decode(in); + } catch (IOException e) { + throw new RuntimeException("Failed to decode key", e); + } + return Arrays.copyOfRange(bytes, valueOffset, bytes.length - in.available()); Review Comment: Each call to `getEncodedKey()` runs `kvCoder.getKeyCoder().decode(in)` and slices `Arrays.copyOfRange(...)`. If `KvToFlinkKeyKeySelector#getKey` is invoked more than once on the same `LazyKvWindowedValue` (e.g., once for key-group partitioning and once when inserting into Flink's `SortingDataInput`), it will re-decode the key via `keyCoder.decode(in)` each time. Would it make sense to cache the key's end offset (or the encoded key slice) after the first `getEncodedKey()` call? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
