DuckDBDataChunkReader.java

package org.duckdb;

import static org.duckdb.DuckDBBindings.*;

import java.nio.ByteBuffer;
import java.util.stream.LongStream;
import org.duckdb.DuckDBFunctions.FunctionException;

/**
 * Reader over callback input data chunks.
 *
 * <p>Column index violations throw {@link IndexOutOfBoundsException}.
 */
public final class DuckDBDataChunkReader {
    private final ByteBuffer chunkRef;
    private final long rowCount;
    private final long columnCount;
    private final DuckDBReadableVector[] vectors;

    DuckDBDataChunkReader(ByteBuffer chunkRef) {
        if (chunkRef == null) {
            throw new FunctionException("Invalid data chunk reference");
        }
        this.chunkRef = chunkRef;
        this.rowCount = duckdb_data_chunk_get_size(chunkRef);
        this.columnCount = duckdb_data_chunk_get_column_count(chunkRef);
        this.vectors = new DuckDBReadableVector[Math.toIntExact(columnCount)];

        for (long columnIndex = 0; columnIndex < columnCount; columnIndex++) {
            ByteBuffer vectorRef = duckdb_data_chunk_get_vector(chunkRef, columnIndex);
            int arrayIndex = Math.toIntExact(columnIndex);
            vectors[arrayIndex] = new DuckDBReadableVector(vectorRef, rowCount);
        }
    }

    public long rowCount() {
        return rowCount;
    }

    public long columnCount() {
        return columnCount;
    }

    public LongStream stream() {
        return LongStream.range(0, rowCount);
    }

    public DuckDBReadableVector vector(long columnIndex) {
        if (columnIndex < 0 || columnIndex >= columnCount) {
            throw new IndexOutOfBoundsException("Column index out of bounds: " + columnIndex);
        }
        int arrayIndex = Math.toIntExact(columnIndex);
        return vectors[arrayIndex];
    }
}