File size: 4,487 Bytes
9ae1216 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | /**
* @import {DataReader, DecodedArray, ParquetType} from '../src/types.js'
*/
import { readVarInt } from './thrift.js'
/**
* Read values from a run-length encoded/bit-packed hybrid encoding.
*
* If length is zero, then read int32 length at the start.
*
* @param {DataReader} reader
* @param {number} width - bitwidth
* @param {DecodedArray} output
* @param {number} [length] - length of the encoded data
*/
export function readRleBitPackedHybrid(reader, width, output, length) {
if (length === undefined) {
length = reader.view.getUint32(reader.offset, true)
reader.offset += 4
}
const startOffset = reader.offset
let seen = 0
while (seen < output.length) {
const header = readVarInt(reader)
if (header & 1) {
// bit-packed
seen = readBitPacked(reader, header, width, output, seen)
} else {
// rle
const count = header >>> 1
readRle(reader, count, width, output, seen)
seen += count
}
}
reader.offset = startOffset + length // duckdb writes an empty block
}
/**
* Run-length encoding: read value with bitWidth and repeat it count times.
*
* @param {DataReader} reader
* @param {number} count
* @param {number} bitWidth
* @param {DecodedArray} output
* @param {number} seen
*/
function readRle(reader, count, bitWidth, output, seen) {
const width = bitWidth + 7 >> 3
let value = 0
for (let i = 0; i < width; i++) {
value |= reader.view.getUint8(reader.offset++) << (i << 3)
}
// assert(value < 1 << bitWidth)
// repeat value count times
for (let i = 0; i < count; i++) {
output[seen + i] = value
}
}
/**
* Read a bit-packed run of the rle/bitpack hybrid.
* Supports width > 8 (crossing bytes).
*
* @param {DataReader} reader
* @param {number} header - bit-pack header
* @param {number} bitWidth
* @param {DecodedArray} output
* @param {number} seen
* @returns {number} total output values so far
*/
function readBitPacked(reader, header, bitWidth, output, seen) {
let count = header >> 1 << 3 // values to read
const mask = (1 << bitWidth) - 1
let data = 0
if (reader.offset < reader.view.byteLength) {
data = reader.view.getUint8(reader.offset++)
} else if (mask) {
// sometimes out-of-bounds reads are masked out
throw new Error(`parquet bitpack offset ${reader.offset} out of range`)
}
let left = 8
let right = 0
// read values
while (count) {
// if we have crossed a byte boundary, shift the data
if (right > 8) {
right -= 8
left -= 8
data >>>= 8
} else if (left - right < bitWidth) {
// if we don't have bitWidth number of bits to read, read next byte
data |= reader.view.getUint8(reader.offset) << left
reader.offset++
left += 8
} else {
if (seen < output.length) {
// emit value
output[seen++] = data >> right & mask
}
count--
right += bitWidth
}
}
return seen
}
/**
* @param {DataReader} reader
* @param {number} count
* @param {ParquetType} type
* @param {number | undefined} typeLength
* @returns {DecodedArray}
*/
export function byteStreamSplit(reader, count, type, typeLength) {
const width = byteWidth(type, typeLength)
const bytes = new Uint8Array(count * width)
for (let b = 0; b < width; b++) {
for (let i = 0; i < count; i++) {
bytes[i * width + b] = reader.view.getUint8(reader.offset++)
}
}
// interpret bytes as typed array
if (type === 'FLOAT') return new Float32Array(bytes.buffer)
else if (type === 'DOUBLE') return new Float64Array(bytes.buffer)
else if (type === 'INT32') return new Int32Array(bytes.buffer)
else if (type === 'INT64') return new BigInt64Array(bytes.buffer)
else if (type === 'FIXED_LEN_BYTE_ARRAY') {
// split into arrays of typeLength
const split = new Array(count)
for (let i = 0; i < count; i++) {
split[i] = bytes.subarray(i * width, (i + 1) * width)
}
return split
}
throw new Error(`parquet byte_stream_split unsupported type: ${type}`)
}
/**
* @param {ParquetType} type
* @param {number | undefined} typeLength
* @returns {number}
*/
function byteWidth(type, typeLength) {
switch (type) {
case 'INT32':
case 'FLOAT':
return 4
case 'INT64':
case 'DOUBLE':
return 8
case 'FIXED_LEN_BYTE_ARRAY':
if (!typeLength) throw new Error('parquet byteWidth missing type_length')
return typeLength
default:
throw new Error(`parquet unsupported type: ${type}`)
}
}
|