@dobesv/parquets
Version:
TypeScript implementation of the Parquet file format, based on parquet.js
512 lines • 19.2 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.ParquetTransformer = exports.ParquetEnvelopeWriter = exports.ParquetWriter = void 0;
const stream_1 = require("stream");
const codec_1 = require("./codec");
const Compression = require("./compression");
const thrift_1 = require("./thrift");
const Util = require("./util");
const Int64 = require("node-int64");
const shred_1 = require("./shred");
/**
* Parquet File Magic String
*/
const PARQUET_MAGIC = 'PAR1';
/**
* Parquet File Format Version
*/
const PARQUET_VERSION = 1;
/**
* Default Page and Row Group sizes
*/
const PARQUET_DEFAULT_PAGE_SIZE = 8192;
const PARQUET_DEFAULT_ROW_GROUP_SIZE = 4096;
/**
* Repetition and Definition Level Encoding
*/
const PARQUET_RDLVL_TYPE = 'INT32';
const PARQUET_RDLVL_ENCODING = 'RLE';
/**
* Write a parquet file to an output stream. The ParquetWriter will perform
* buffering/batching for performance, so close() must be called after all rows
* are written.
*/
class ParquetWriter {
/**
* Create a new buffered parquet writer for a given envelope writer
*/
constructor(schema, envelopeWriter, opts) {
this.schema = schema;
this.envelopeWriter = envelopeWriter;
this.rowBuffer = new shred_1.ParquetWriteBuffer(schema);
this.rowGroupSize = opts.rowGroupSize || PARQUET_DEFAULT_ROW_GROUP_SIZE;
this.closed = false;
this.headerWritten = false;
this.userMetadata = {};
}
/**
* Convenience method to create a new buffered parquet writer that writes to
* the specified file
*/
static async openFile(schema, path, opts) {
const outputStream = await Util.osopen(path, opts);
return ParquetWriter.openStream(schema, outputStream, opts);
}
/**
* Convenience method to create a new buffered parquet writer that writes to
* the specified stream
*/
static async openStream(schema, outputStream, opts) {
if (!opts) {
// tslint:disable-next-line:no-parameter-reassignment
opts = {};
}
const envelopeWriter = await ParquetEnvelopeWriter.openStream(schema, outputStream, opts);
return new ParquetWriter(schema, envelopeWriter, opts);
}
/**
* Write the header if it was not already written
*/
async ensureHeaderWritten() {
if (!this.headerWritten) {
try {
// Set the flag before making the call so that a concurrent call while the header
// is being written will not write the header a second time
this.headerWritten = true;
// Go ahead and write the header
await this.envelopeWriter.writeHeader();
}
catch (err) {
this.envelopeWriter.close();
throw err;
}
}
}
/**
* Append a single row to the parquet file. Rows are buffered in memory until
* rowGroupSize rows are in the buffer or close() is called
*/
async appendRow(row) {
if (this.closed) {
throw new Error('writer was closed');
}
(0, shred_1.shredRecord)(this.schema, row, this.rowBuffer);
if (this.rowBuffer.rowCount >= this.rowGroupSize) {
await this.ensureHeaderWritten();
await this.envelopeWriter.writeRowGroup(this.rowBuffer);
this.rowBuffer = new shred_1.ParquetWriteBuffer(this.schema);
}
}
/**
* Finish writing the parquet file and commit the footer to disk. This method
* MUST be called after you are finished adding rows. You must not call this
* method twice on the same object or add any rows after the close() method has
* been called
*/
async close(callback) {
if (this.closed) {
throw new Error('writer was closed');
}
this.closed = true;
// Make sure we have written the header even if the file is empty
await this.ensureHeaderWritten();
if (this.rowBuffer.rowCount > 0 ||
this.rowBuffer.rowCount >= this.rowGroupSize) {
await this.envelopeWriter.writeRowGroup(this.rowBuffer);
this.rowBuffer = new shred_1.ParquetWriteBuffer(this.schema);
}
await this.envelopeWriter.writeFooter(this.userMetadata);
await this.envelopeWriter.close();
this.envelopeWriter = null;
if (callback) {
callback();
}
}
/**
* Add key<>value metadata to the file
*/
setMetadata(key, value) {
// TODO: value to be any, obj -> JSON
this.userMetadata[String(key)] = String(value);
}
/**
* Set the parquet row group size. This values controls the maximum number
* of rows that are buffered in memory at any given time as well as the number
* of rows that are co-located on disk. A higher value is generally better for
* read-time I/O performance at the tradeoff of write-time memory usage.
*/
setRowGroupSize(cnt) {
this.rowGroupSize = cnt;
}
/**
* Set the parquet data page size. The data page size controls the maximum
* number of column values that are written to disk as a consecutive array
*/
setPageSize(cnt) {
this.envelopeWriter.setPageSize(cnt);
}
}
exports.ParquetWriter = ParquetWriter;
/**
* Create a parquet file from a schema and a number of row groups. This class
* performs direct, unbuffered writes to the underlying output stream and is
* intendend for advanced and internal users; the writeXXX methods must be
* called in the correct order to produce a valid file.
*/
class ParquetEnvelopeWriter {
constructor(schema, writeFn, closeFn, fileOffset, opts) {
this.schema = schema;
this.write = writeFn;
this.close = closeFn;
this.offset = fileOffset;
this.rowCount = 0;
this.rowGroups = [];
this.pageSize = opts.pageSize || PARQUET_DEFAULT_PAGE_SIZE;
this.useDataPageV2 = 'useDataPageV2' in opts ? opts.useDataPageV2 : false;
}
/**
* Create a new parquet envelope writer that writes to the specified stream
*/
static async openStream(schema, outputStream, opts) {
const writeFn = Util.oswrite.bind(undefined, outputStream);
const closeFn = Util.osclose.bind(undefined, outputStream);
return new ParquetEnvelopeWriter(schema, writeFn, closeFn, 0, opts);
}
writeSection(buf) {
this.offset += buf.length;
return this.write(buf);
}
/**
* Encode the parquet file header
*/
writeHeader() {
return this.writeSection(Buffer.from(PARQUET_MAGIC));
}
/**
* Encode a parquet row group. The records object should be created using the
* shredRecord method
*/
writeRowGroup(records) {
const rowGroup = encodeRowGroup(this.schema, records, {
baseOffset: this.offset,
pageSize: this.pageSize,
useDataPageV2: this.useDataPageV2,
});
this.rowCount += records.rowCount;
this.rowGroups.push(rowGroup.metadata);
return this.writeSection(rowGroup.body);
}
/**
* Write the parquet file footer
*/
writeFooter(userMetadata) {
if (!userMetadata) {
// tslint:disable-next-line:no-parameter-reassignment
userMetadata = {};
}
return this.writeSection(encodeFooter(this.schema, this.rowCount, this.rowGroups, userMetadata));
}
/**
* Set the parquet data page size. The data page size controls the maximum
* number of column values that are written to disk as a consecutive array
*/
setPageSize(cnt) {
this.pageSize = cnt;
}
}
exports.ParquetEnvelopeWriter = ParquetEnvelopeWriter;
/**
* Create a parquet transform stream
*/
class ParquetTransformer extends stream_1.Transform {
constructor(schema, opts = {}) {
super({ objectMode: true });
this.waiting = [];
const writeFn = (function (t) {
return function (b) {
if (!t.push(b)) {
// stop writing until the readable is ready again
return new Promise((resolve, reject) => {
t.waiting.push([resolve, reject]);
});
}
return Promise.resolve();
};
})(this);
const closeFn = (function (t) {
return function () {
t.push(null);
return Promise.resolve();
};
})(this);
this.writer = new ParquetWriter(schema, new ParquetEnvelopeWriter(schema, writeFn, closeFn, 0, opts), opts);
}
// If I/O was delayed due to backpressure and then the stream is destroyed,
// propagate an error back to the callee of the I/O operation(s)
// tslint:disable-next-line:function-name
_destroy(error, callback) {
try {
if (this.waiting.length) {
const waiting = this.waiting;
this.waiting = [];
waiting.forEach(([resolve, reject]) => error ? reject(error) : resolve());
}
callback(null);
}
catch (err) {
callback(err);
}
}
// If we get backpressure we will delay returning from a call to write until
// the next call to _read
// tslint:disable-next-line:function-name
_read(arg) {
if (this.waiting.length) {
const waiting = this.waiting;
this.waiting = [];
waiting.forEach(([resolve]) => resolve());
}
return super._read(arg);
}
// tslint:disable-next-line:function-name
_transform(row, encoding, callback) {
if (row) {
this.writer.appendRow(row).then(() => callback(), err => callback(err));
}
else {
callback();
}
}
// tslint:disable-next-line:function-name
_flush(callback) {
this.writer.close(callback);
}
}
exports.ParquetTransformer = ParquetTransformer;
/**
* Encode a consecutive array of data using one of the parquet encodings
*/
function encodeValues(type, encoding, values, opts) {
if (!(encoding in codec_1.PARQUET_CODEC)) {
throw new Error(`invalid encoding: ${encoding}`);
}
return codec_1.PARQUET_CODEC[encoding].encodeValues(type, values, opts);
}
/**
* Encode a parquet data page
*/
function encodeDataPage(column, data) {
/* encode repetition and definition levels */
let rLevelsBuf = Buffer.alloc(0);
if (column.rLevelMax > 0) {
rLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.rLevels, {
bitWidth: Util.getBitWidth(column.rLevelMax),
// disableEnvelope: false
});
}
let dLevelsBuf = Buffer.alloc(0);
if (column.dLevelMax > 0) {
dLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.dLevels, {
bitWidth: Util.getBitWidth(column.dLevelMax),
// disableEnvelope: false
});
}
/* encode values */
const valuesBuf = encodeValues(column.primitiveType, column.encoding, data.values, { typeLength: column.typeLength, bitWidth: column.typeLength });
const dataBuf = Buffer.concat([rLevelsBuf, dLevelsBuf, valuesBuf]);
// compression = column.compression === 'UNCOMPRESSED' ? (compression || 'UNCOMPRESSED') : column.compression;
const compressedBuf = Compression.deflate(column.compression, dataBuf);
/* build page header */
const header = new thrift_1.PageHeader({
type: thrift_1.PageType.DATA_PAGE,
data_page_header: new thrift_1.DataPageHeader({
num_values: data.count,
encoding: thrift_1.Encoding[column.encoding],
definition_level_encoding: thrift_1.Encoding[PARQUET_RDLVL_ENCODING],
repetition_level_encoding: thrift_1.Encoding[PARQUET_RDLVL_ENCODING], // [PARQUET_RDLVL_ENCODING]
}),
uncompressed_page_size: dataBuf.length,
compressed_page_size: compressedBuf.length,
});
/* concat page header, repetition and definition levels and values */
const headerBuf = Util.serializeThrift(header);
const page = Buffer.concat([headerBuf, compressedBuf]);
return { header, headerSize: headerBuf.length, page };
}
/**
* Encode a parquet data page (v2)
*/
function encodeDataPageV2(column, data, rowCount) {
/* encode values */
const valuesBuf = encodeValues(column.primitiveType, column.encoding, data.values, {
typeLength: column.typeLength,
bitWidth: column.typeLength,
});
// compression = column.compression === 'UNCOMPRESSED' ? (compression || 'UNCOMPRESSED') : column.compression;
const compressedBuf = Compression.deflate(column.compression, valuesBuf);
/* encode repetition and definition levels */
let rLevelsBuf = Buffer.alloc(0);
if (column.rLevelMax > 0) {
rLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.rLevels, {
bitWidth: Util.getBitWidth(column.rLevelMax),
disableEnvelope: true,
});
}
let dLevelsBuf = Buffer.alloc(0);
if (column.dLevelMax > 0) {
dLevelsBuf = encodeValues(PARQUET_RDLVL_TYPE, PARQUET_RDLVL_ENCODING, data.dLevels, {
bitWidth: Util.getBitWidth(column.dLevelMax),
disableEnvelope: true,
});
}
/* build page header */
const header = new thrift_1.PageHeader({
type: thrift_1.PageType.DATA_PAGE_V2,
data_page_header_v2: new thrift_1.DataPageHeaderV2({
num_values: data.count,
num_nulls: data.count - data.values.length,
num_rows: rowCount,
encoding: thrift_1.Encoding[column.encoding],
definition_levels_byte_length: dLevelsBuf.length,
repetition_levels_byte_length: rLevelsBuf.length,
is_compressed: column.compression !== 'UNCOMPRESSED',
}),
uncompressed_page_size: rLevelsBuf.length + dLevelsBuf.length + valuesBuf.length,
compressed_page_size: rLevelsBuf.length + dLevelsBuf.length + compressedBuf.length,
});
/* concat page header, repetition and definition levels and values */
const headerBuf = Util.serializeThrift(header);
const page = Buffer.concat([
headerBuf,
rLevelsBuf,
dLevelsBuf,
compressedBuf,
]);
return { header, headerSize: headerBuf.length, page };
}
/**
* Encode an array of values into a parquet column chunk
*/
function encodeColumnChunk(column, buffer, offset, opts) {
const data = buffer.columnData[column.path.join()];
const baseOffset = (opts.baseOffset || 0) + offset;
/* encode data page(s) */
// const pages: Buffer[] = [];
let pageBuf;
// tslint:disable-next-line:variable-name
let total_uncompressed_size = 0;
// tslint:disable-next-line:variable-name
let total_compressed_size = 0;
{
let result;
if (opts.useDataPageV2) {
result = encodeDataPageV2(column, data, buffer.rowCount);
}
else {
result = encodeDataPage(column, data);
}
// pages.push(result.page);
pageBuf = result.page;
total_uncompressed_size +=
result.header.uncompressed_page_size + result.headerSize;
total_compressed_size +=
result.header.compressed_page_size + result.headerSize;
}
// const pagesBuf = Buffer.concat(pages);
// const compression = column.compression === 'UNCOMPRESSED' ? (opts.compression || 'UNCOMPRESSED') : column.compression;
/* prepare metadata header */
const metadata = new thrift_1.ColumnMetaData({
path_in_schema: column.path,
num_values: data.count,
data_page_offset: baseOffset,
encodings: [],
total_uncompressed_size,
total_compressed_size,
type: thrift_1.Type[column.primitiveType],
codec: thrift_1.CompressionCodec[column.compression],
});
/* list encodings */
metadata.encodings.push(thrift_1.Encoding[PARQUET_RDLVL_ENCODING]);
metadata.encodings.push(thrift_1.Encoding[column.encoding]);
/* concat metadata header and data pages */
const metadataOffset = baseOffset + pageBuf.length;
const body = Buffer.concat([pageBuf, Util.serializeThrift(metadata)]);
return { body, metadata, metadataOffset };
}
/**
* Encode a list of column values into a parquet row group
*/
function encodeRowGroup(schema, data, opts) {
const metadata = new thrift_1.RowGroup({
num_rows: data.rowCount,
columns: [],
total_byte_size: 0,
});
let body = Buffer.alloc(0);
for (const field of schema.fieldList) {
if (field.isNested) {
continue;
}
const cchunkData = encodeColumnChunk(field, data, body.length, opts);
const cchunk = new thrift_1.ColumnChunk({
file_offset: cchunkData.metadataOffset,
meta_data: cchunkData.metadata,
});
metadata.columns.push(cchunk);
metadata.total_byte_size = new Int64(+metadata.total_byte_size + cchunkData.body.length);
body = Buffer.concat([body, cchunkData.body]);
}
return { body, metadata };
}
/**
* Encode a parquet file metadata footer
*/
function encodeFooter(schema, rowCount, rowGroups, userMetadata) {
const metadata = new thrift_1.FileMetaData({
version: PARQUET_VERSION,
created_by: 'parquets',
num_rows: rowCount,
row_groups: rowGroups,
schema: [],
key_value_metadata: [],
});
for (const key in userMetadata) {
const kv = new thrift_1.KeyValue({
key,
value: userMetadata[key],
});
metadata.key_value_metadata.push(kv);
}
{
const schemaRoot = new thrift_1.SchemaElement({
name: 'root',
num_children: Object.keys(schema.fields).length,
});
metadata.schema.push(schemaRoot);
}
for (const field of schema.fieldList) {
const relt = thrift_1.FieldRepetitionType[field.repetitionType];
const schemaElem = new thrift_1.SchemaElement({
name: field.name,
repetition_type: relt,
});
if (field.isNested) {
schemaElem.num_children = field.fieldCount;
}
else {
schemaElem.type = thrift_1.Type[field.primitiveType];
}
if (field.originalType) {
schemaElem.converted_type = thrift_1.ConvertedType[field.originalType];
}
schemaElem.type_length = field.typeLength;
metadata.schema.push(schemaElem);
}
const metadataEncoded = Util.serializeThrift(metadata);
const footerEncoded = Buffer.alloc(metadataEncoded.length + 8);
metadataEncoded.copy(footerEncoded);
footerEncoded.writeUInt32LE(metadataEncoded.length, metadataEncoded.length);
footerEncoded.write(PARQUET_MAGIC, metadataEncoded.length + 4);
return footerEncoded;
}
//# sourceMappingURL=writer.js.map