net.snowflake.client.jdbc.JsonResultChunk Maven / Gradle / Ivy
/*
* Copyright (c) 2012-2019 Snowflake Computing Inc. All rights reserved.
*/
package net.snowflake.client.jdbc;
import com.fasterxml.jackson.databind.JsonNode;
import java.lang.ref.SoftReference;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.BitSet;
import java.util.LinkedList;
import java.util.List;
import net.snowflake.client.core.SFBaseSession;
import net.snowflake.client.log.SFLogger;
import net.snowflake.client.log.SFLoggerFactory;
import net.snowflake.common.core.SqlState;
public class JsonResultChunk extends SnowflakeResultChunk {
private static final int NULL_VALUE = Integer.MIN_VALUE;
private static final SFLogger logger = SFLoggerFactory.getLogger(JsonResultChunk.class);
private ResultChunkData data;
private int currentRow;
private SFBaseSession session;
public JsonResultChunk(
String url, int rowCount, int colCount, int uncompressedSize, SFBaseSession session) {
super(url, rowCount, colCount, uncompressedSize);
data = new BlockResultChunkDataV2(computeCharactersNeeded(), rowCount, colCount, session);
this.session = session;
}
public static Object extractCell(JsonNode resultData, int rowIdx, int colIdx) {
JsonNode currentRow = resultData.get(rowIdx);
JsonNode colNode = currentRow.get(colIdx);
if (colNode.isTextual()) {
return colNode.asText();
} else if (colNode.isNumber()) {
return colNode.numberValue();
} else if (colNode.isNull()) {
return null;
}
throw new RuntimeException("Unknow json type");
}
public void tryReuse(ResultChunkDataCache cache) {
// Allocate chunk data, double necessary amount for later reuse
cache.reuseOrCreateResultData(data);
}
/**
* Creates a String object for the given cell
*
* @param rowIdx zero based row
* @param colIdx zero based column
* @return String
*/
public final Object getCell(int rowIdx, int colIdx) {
return data.get(colCount * rowIdx + colIdx);
}
public final void addRow(Object[] row) throws SnowflakeSQLException {
if (row.length != colCount) {
throw new SnowflakeSQLLoggedException(
this.session,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
SqlState.INTERNAL_ERROR,
"Exception: expected " + colCount + " columns and received " + row.length);
}
for (Object cell : row) {
if (cell == null) {
data.add(null);
} else {
if (cell instanceof String) {
data.add((String) cell);
} else if (cell instanceof Boolean) {
data.add((boolean) cell ? "1" : "0");
} else {
throw new SnowflakeSQLLoggedException(
this.session,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
SqlState.INTERNAL_ERROR,
"unknown data type in JSON row " + cell.getClass().toString());
}
}
}
currentRow++;
}
/**
* Checks that all data has been added after parsing.
*
* @throws SnowflakeSQLException when rows are not all downloaded
*/
public final void ensureRowsComplete() throws SnowflakeSQLException {
// Check that all the rows have been decoded, raise an error if not
if (rowCount != currentRow) {
throw new SnowflakeSQLException(
SqlState.INTERNAL_ERROR,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
"Exception: expected " + rowCount + " rows and received " + currentRow);
}
}
@Override
public void reset() {
this.currentRow = 0;
this.data.reset();
}
/**
* Compute the memory necessary to store the data of this chunk
*
* @return necessary memory in bytes
*/
@Override
public final long computeNeededChunkMemory() {
if (data != null) {
return data.computeNeededChunkMemory();
}
return 0;
}
@Override
public final void freeData() {
if (data != null) {
data.freeData();
}
}
public int computeCharactersNeeded() {
// remove [ , ] characters, they won't be stored
return uncompressedSize
- (rowCount * 2) // opening [ and komma separating rows
- (rowCount * colCount); // komma separating cells and closing ]
}
public void addOffset(int offset) throws SnowflakeSQLException {
data.addOffset(offset);
}
public void setIsNull() throws SnowflakeSQLException {
data.setIsNull();
}
public void setLastLength(int len) throws SnowflakeSQLException {
data.setLastLength(len);
}
public void nextIndex() throws SnowflakeSQLException {
data.nextIndex();
}
public byte get(int offset) throws SnowflakeSQLException {
return data.getByte(offset);
}
public void addByte(byte b, int pos) throws SnowflakeSQLException {
data.addByte(b, pos);
}
public void addBytes(byte[] src, int offset, int pos, int length) throws SnowflakeSQLException {
data.addBytes(src, offset, pos, length);
}
/**
* This class abstracts the storage of the strings in one chunk. To the user the class behaves
* similar to an ArrayList.
*/
private interface ResultChunkData {
/**
* Add the string to the data list
*
* @param string value to add
*/
void add(String string) throws SnowflakeSQLException;
/**
* Access an element by an index
*
* @param index determines the element
* @return String containing the same data as the one passed to add()
*/
String get(int index);
/**
* Compute the necessary memory to store this chunk
*
* @return memory in bytes
*/
long computeNeededChunkMemory();
/** Let GC collect the memory */
void freeData();
/**
* add offset into offset buffer
*
* @param offset
* @throws SnowflakeSQLException
*/
void addOffset(int offset) throws SnowflakeSQLException;
/**
* label current value as null
*
* @throws SnowflakeSQLException
*/
void setIsNull() throws SnowflakeSQLException;
/**
* increase index
*
* @throws SnowflakeSQLException
*/
void nextIndex() throws SnowflakeSQLException;
/**
* set the last length
*
* @param len
* @throws SnowflakeSQLException
*/
void setLastLength(int len) throws SnowflakeSQLException;
/**
* get one byte from the byte array
*
* @param offset
* @return
* @throws SnowflakeSQLException
*/
byte getByte(int offset) throws SnowflakeSQLException;
/**
* add one byte to the byte array at nextIndex
*
* @param b
* @param pos
* @throws SnowflakeSQLException
*/
void addByte(byte b, int pos) throws SnowflakeSQLException;
/**
* add bytes to the byte array
*
* @param src
* @param src_offset
* @param pos
* @param length
* @throws SnowflakeSQLException
*/
void addBytes(byte[] src, int src_offset, int pos, int length) throws SnowflakeSQLException;
void reset();
}
/**
* BlockResultChunkDataV2: This implementation copies the strings to byte arrays and stores the
* offsets and bitmaps. This design can save half of the memory usage compared to the original one
*/
private static class BlockResultChunkDataV2 implements ResultChunkData {
BlockResultChunkDataV2(int totalLength, int rowCount, int colCount, SFBaseSession session) {
this.blockCount = getBlock(totalLength - 1) + 1;
this.rowCount = rowCount;
this.colCount = colCount;
this.metaBlockCount = getMetaBlock(this.rowCount * this.colCount - 1) + 1;
this.session = session;
}
@Override
public void reset() {
freeData();
this.lastLength = 0;
this.nextIndex = 0;
}
@Override
public void addOffset(int offset) {
if (data.size() < blockCount || offsets.size() < metaBlockCount) {
allocateArrays();
}
offsets.get(getMetaBlock(nextIndex))[getMetaBlockIndex(nextIndex)] = offset;
}
@Override
public void setIsNull() {
isNulls.get(getMetaBlock(nextIndex)).set(getMetaBlockIndex(nextIndex));
}
@Override
public void setLastLength(int len) {
lastLength = len;
}
@Override
public byte getByte(int offset) {
return data.get(getBlock(offset))[getBlockOffset(offset)];
}
@Override
public void addByte(byte b, int pos) {
if (data.size() < blockCount || offsets.size() < metaBlockCount) {
allocateArrays();
}
data.get(getBlock(pos))[getBlockOffset(pos)] = b;
}
@Override
public void addBytes(byte[] src, int src_offset, int pos, int length) {
if (data.size() < blockCount || offsets.size() < metaBlockCount) {
allocateArrays();
}
final int offset = pos;
// copy string to the char array
int copied = 0;
if (spaceLeftOnBlock(offset) < length) {
while (copied < length) {
final int copySize = Math.min(length - copied, spaceLeftOnBlock(offset + copied));
System.arraycopy(
src,
src_offset + copied,
data.get(getBlock(offset + copied)),
getBlockOffset(offset + copied),
copySize);
copied += copySize;
}
} else {
System.arraycopy(
src, src_offset, data.get(getBlock(offset)), getBlockOffset(offset), length);
}
}
@Override
public void nextIndex() {
nextIndex++;
}
@Override
public void add(String string) throws SnowflakeSQLException {
throw new SnowflakeSQLLoggedException(
this.session,
ErrorCode.INTERNAL_ERROR.getMessageCode(),
SqlState.INTERNAL_ERROR,
"Unimplemented");
}
private int getLength(int index, int offset) {
if (index == rowCount * colCount - 1) {
// last one
return lastLength;
} else {
int nextOffset = offsets.get(getMetaBlock(index + 1))[getMetaBlockIndex(index + 1)];
return nextOffset - offset;
}
}
@Override
public String get(int index) {
final boolean isNull = isNulls.get(getMetaBlock(index)).get(getMetaBlockIndex(index));
if (isNull) {
return null;
} else {
final int offset = offsets.get(getMetaBlock(index))[getMetaBlockIndex(index)];
final int length = getLength(index, offset);
// Create string from the char arrays
if (spaceLeftOnBlock(offset) < length) {
int copied = 0;
byte[] cell = new byte[length];
while (copied < length) {
final int copySize = Math.min(length - copied, spaceLeftOnBlock(offset + copied));
System.arraycopy(
data.get(getBlock(offset + copied)),
getBlockOffset(offset + copied),
cell,
copied,
copySize);
copied += copySize;
}
return new String(cell, StandardCharsets.UTF_8);
} else {
return new String(
data.get(getBlock(offset)), getBlockOffset(offset), length, StandardCharsets.UTF_8);
}
}
}
@Override
public long computeNeededChunkMemory() {
long dataRequirement = blockCount * blockLength * 1L;
long metadataRequirement =
metaBlockCount * metaBlockLength * 4L // offsets
+ metaBlockCount * metaBlockLength / 8L // isNulls
+ 1L; // lastLength
return dataRequirement + metadataRequirement;
}
@Override
public void freeData() {
data.clear();
offsets.clear();
isNulls.clear();
}
private static int getBlock(int offset) {
return offset >> blockLengthBits;
}
private static int getBlockOffset(int offset) {
return offset & (blockLength - 1);
}
private static int spaceLeftOnBlock(int offset) {
return blockLength - getBlockOffset(offset);
}
private static int getMetaBlock(int index) {
return index >> metaBlockLengthBits;
}
private static int getMetaBlockIndex(int index) {
return index & (metaBlockLength - 1);
}
private void allocateArrays() {
logger.debug("allocating {} B for ResultChunk", computeNeededChunkMemory());
while (data.size() < blockCount) {
data.add(new byte[1 << blockLengthBits]);
}
while (offsets.size() < metaBlockCount) {
offsets.add(new int[1 << metaBlockLengthBits]);
isNulls.add(new BitSet(1 << metaBlockLengthBits));
}
logger.debug("allocated {} B for ResultChunk", computeNeededChunkMemory());
}
// blocks for storing the string data
int blockCount;
private static final int blockLengthBits = 23;
private static int blockLength = 1 << blockLengthBits;
private final ArrayList data = new ArrayList<>();
SFBaseSession session;
// blocks for storing offsets and lengths
int metaBlockCount;
private static int metaBlockLengthBits = 15;
private static int metaBlockLength = 1 << metaBlockLengthBits;
private final ArrayList offsets = new ArrayList<>();
private final ArrayList isNulls = new ArrayList<>();
private int lastLength;
private int rowCount, colCount;
private int nextIndex = 0;
}
/** Cache the data, offset and length blocks */
static class ResultChunkDataCache {
/**
* Add the data to the cache. CAUTION: The result chunk is not usable afterward
*
* @param chunk add this to the cache
*/
void add(JsonResultChunk chunk) {
cache.add(new SoftReference<>(chunk.data));
chunk.data = null;
}
/**
* Creates a new ResultChunkData which reuses as much blocks as possible
*
* @param data fill this with reused blocks
*/
void reuseOrCreateResultData(ResultChunkData data) {
List> remove = new ArrayList<>();
try {
for (SoftReference ref : cache) {
ResultChunkData dat = ref.get();
if (dat == null) {
remove.add(ref);
continue;
}
if (dat instanceof BlockResultChunkDataV2) {
BlockResultChunkDataV2 bTargetData = (BlockResultChunkDataV2) data;
BlockResultChunkDataV2 bCachedDat = (BlockResultChunkDataV2) dat;
if (bCachedDat.data.size() == 0 && bCachedDat.offsets.size() == 0) {
remove.add(ref);
continue;
}
while (bTargetData.data.size() < bTargetData.blockCount && bCachedDat.data.size() > 0) {
bTargetData.data.add(bCachedDat.data.remove(bCachedDat.data.size() - 1));
}
while (bTargetData.offsets.size() < bTargetData.metaBlockCount
&& bCachedDat.offsets.size() > 0) {
bTargetData.offsets.add(bCachedDat.offsets.remove(bCachedDat.offsets.size() - 1));
BitSet isNulls = bCachedDat.isNulls.remove(bCachedDat.isNulls.size() - 1);
isNulls.clear(); // SNOW-80208 have to clear isNulls explicitly
bTargetData.isNulls.add(isNulls);
}
if (bTargetData.data.size() == bTargetData.blockCount
&& bTargetData.offsets.size() == bTargetData.metaBlockCount) {
return;
}
} else {
remove.add(ref);
}
}
} finally {
cache.removeAll(remove);
}
}
/** Let GC collect all data hold by the cache */
void clear() {
for (SoftReference ref : cache) {
ResultChunkData dat = ref.get();
if (dat != null) {
dat.freeData();
}
}
cache.clear();
}
List> cache = new LinkedList<>();
}
}
© 2015 - 2024 Weber Informatics LLC | Privacy Policy