Repository: apache/flink
File: flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/operators/aggregate/window/buffers/RecordsWindowBuffer.java (Java)

Code before:
        this.recordsBuffer =
                new WindowBytesMultiMap(
                        operatorOwner, memoryManager, memorySize, keySer, inputSer.getArity());
        this.recordSerializer = inputSer;
        this.reuseWindowKey = new WindowKeySerializer(keySer).createInstance();
        this.requiresCopy = requiresCopy;
        this.shiftTimeZone = shiftTimeZone;
    }

    @Override
    public void addElement(RowData key, long sliceEnd, RowData element) throws Exception {
        // track the lowest trigger time, if watermark exceeds the trigger time,
        // it means there are some elements in the buffer belong to a window going to be fired,
        // and we need to flush the buffer into state for firing.
        minSliceEnd = Math.min(sliceEnd, minSliceEnd);

        reuseWindowKey.replace(sliceEnd, key);
        while (true) {
            LookupInfo<WindowKey, Iterator<RowData>> lookup = recordsBuffer.lookup(reuseWindowKey);
            try {
                recordsBuffer.append(lookup, recordSerializer.toBinaryRow(element));
                break;
            } catch (EOFException e) {
                if (recordsBuffer.getNumKeys() == 0) {
                    // Buffer is empty, retry won't help (record is too large for the buffer)
                    throw e;
                }
                // Buffer has data, flush and retry
                flush();
                checkState(
                        recordsBuffer.getNumKeys() == 0,
                        "The recordsBuffer should be empty after flushing.");
            }
        }
    }

    @Override
    public void advanceProgress(long progress) throws Exception {
        if (isWindowFired(minSliceEnd, progress, shiftTimeZone)) {
            // there should be some window to be fired, flush buffer to state first
            flush();
        }
    }

    @Override
    public void flush() throws Exception {
        if (recordsBuffer.getNumKeys() > 0) {
            KeyValueIterator<WindowKey, Iterator<RowData>> entryIterator =
                    recordsBuffer.getEntryIterator(requiresCopy);
            while (entryIterator.advanceNext()) {
                combineFunction.combine(entryIterator.getKey(), entryIterator.getValue());

Review comment (suggestion):
Is `getNumKeys` correct and sufficient?
Or should `getNumElements` be used instead or in addition?

In `append()`, I see that `numKeys` is not incremented if `lookupInfo.found`.
I guess there's an implicit invariant that numKeys > 0 in this case, but I'm not sure.


Code after:
        this.recordsBuffer =
                new WindowBytesMultiMap(
                        operatorOwner, memoryManager, memorySize, keySer, inputSer.getArity());
        this.recordSerializer = inputSer;
        this.reuseWindowKey = new WindowKeySerializer(keySer).createInstance();
        this.requiresCopy = requiresCopy;
        this.shiftTimeZone = shiftTimeZone;
    }

    @Override
    public void addElement(RowData key, long sliceEnd, RowData element) throws Exception {
        reuseWindowKey.replace(sliceEnd, key);
        while (true) {
            LookupInfo<WindowKey, Iterator<RowData>> lookup = recordsBuffer.lookup(reuseWindowKey);
            try {
                recordsBuffer.append(lookup, recordSerializer.toBinaryRow(element));
                // Track the lowest trigger time. If watermark exceeds the trigger time,
                // it means there are some elements in the buffer belong to a window going
                // to be fired, and we need to flush the buffer into state for firing.
                minSliceEnd = Math.min(sliceEnd, minSliceEnd);
                break;
            } catch (EOFException e) {
                if (recordsBuffer.getNumKeys() == 0) {
                    // Buffer is empty, retry won't help (record is too large for the buffer)
                    throw e;
                }
                // Buffer has data, flush and retry
                flush();
                checkState(
                        recordsBuffer.getNumKeys() == 0,
                        "The recordsBuffer should be empty after flushing.");
            }
        }
    }

    @Override
    public void advanceProgress(long progress) throws Exception {
        if (isWindowFired(minSliceEnd, progress, shiftTimeZone)) {
            // there should be some window to be fired, flush buffer to state first
            flush();
        }
    }

    @Override
    public void flush() throws Exception {
        if (recordsBuffer.getNumKeys() > 0) {
            KeyValueIterator<WindowKey, Iterator<RowData>> entryIterator =
                    recordsBuffer.getEntryIterator(requiresCopy);
            while (entryIterator.advanceNext()) {
                combineFunction.combine(entryIterator.getKey(), entryIterator.getValue());
            }