BufferEntry.java

/*
 * Hivemall: Hive scalable Machine Learning Library
 *
 * Copyright (C) 2015 Makoto YUI
 *
 * Licensed 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 hivemall.utils.buffer;

import javax.annotation.Nonnull;

public class BufferEntry {

    @Nonnull
    private final Buffers buffers;
    private final int size;

    private int pointer = -1;

    BufferEntry(@Nonnull Buffers buffers) {
        this.buffers = buffers;
        this.size = buffers.getEntrySize();
    }

    public int getSize() {
        return size;
    }
    
    /**
     * @see Buffers#fillEntry(int, BufferEntry)
     */
    void locate(int pointer) {
        this.pointer = pointer;
    }
    
    private void ensureLocated() {
        if (pointer == -1) {
            throw new IllegalStateException("locate() is not called");
        }
    }

}

Buffers.java

/*
 * Hivemall: Hive scalable Machine Learning Library
 *
 * Copyright (C) 2015 Makoto YUI
 *
 * Licensed 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 hivemall.utils.buffer;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;

import javax.annotation.Nonnull;

public final class Buffers {
    /* 4 MiB */
    public static final int DEFAULT_CHUNK_SIZE = 4 * 1024 * 1024;

    private final int chunkSize;
    private final int entrySize;
    private final boolean allocateDirect;

    private final List<ByteBuffer> chunks;
    private int allocatedBytes;

    public Buffers(int chunkSize, int entrySize, final boolean direct) {
        this.chunkSize = chunkSize;
        this.entrySize = entrySize;
        this.allocateDirect = direct;
        this.chunks = new ArrayList<ByteBuffer>(32);
        this.allocatedBytes = 0;
    }

    int getEntrySize() {
        return entrySize;
    }

    public int allocatedBytes() {
        return allocatedBytes;
    }

    @Nonnull
    public BufferEntry newEntry() {
        return new BufferEntry(this);
    }

    @Nonnull
    public BufferEntry newEntry(final int entryIndex) {
        BufferEntry entry = new BufferEntry(this);
        fillEntry(entryIndex, entry);
        return entry;
    }

    public void fillEntry(final int entryIndex, @Nonnull final BufferEntry entry) {
        int entryPtr = entryPointer(entryIndex);
        int chunkIdx = chunkIndex(entryPtr);
        ensureCapacity(chunkIdx + 1);

        entry.locate(entryPtr);
    }

    private void ensureCapacity(final int requiredChunks) {
        int currentNumChunks = chunks.size();
        if (currentNumChunks >= requiredChunks) {
            return;
        }

        final int numAllocate = requiredChunks - currentNumChunks;
        for (int i = 0; i < numAllocate; i++) {
            chunks.add(allocateBuffer());
        }
    }

    @Nonnull
    private ByteBuffer allocateBuffer() {
        ByteBuffer buf = allocateDirect ? ByteBuffer.allocateDirect(chunkSize)
                : ByteBuffer.allocate(chunkSize);
        this.allocatedBytes += chunkSize;
        return buf;
    }

    /**
     * @param index logical entry index
     * @return physical entry index
     */
    private int entryPointer(final int entryIndex) {
        return entryIndex * entrySize;
    }

    /**
     * @param pointer physical pointer
     * @return chunk index
     */
    private int chunkIndex(final int pointer) {
        return pointer / chunkSize;
    }

    private int entryIndex(final int pointer) {
        return pointer % chunkSize;
    }

    public byte get(final int pointer) {
        int chunkIdx = chunkIndex(pointer);
        ByteBuffer buf = getChunk(chunkIdx);

        int entryIndex = entryIndex(pointer);
        return buf.get(entryIndex);
    }

    public byte get(final int pointer, @Nonnull final ByteBuffer chunk) {
        int entryIndex = entryIndex(pointer);
        return chunk.get(entryIndex);
    }

    @Nonnull
    public ByteBuffer getChunk(final int chunkIdx) {
        int currentNumChunks = chunks.size();
        if (chunkIdx >= currentNumChunks) {
            throw new IllegalStateException("Chunk #" + chunkIdx + " is not allocated yet");
        }

        ByteBuffer buf = chunks.get(chunkIdx);
        if (buf == null) {
            throw new IllegalStateException("Chunk #" + chunkIdx + " is null");
        }
        return buf;
    }

}