Parquet input: add a boolean type support

This commit is contained in:
Ivan Zhukov 2018-05-17 01:39:17 +03:00
parent f8dae2bc76
commit 4076ae77b4
2 changed files with 23 additions and 2 deletions

View File

@ -95,6 +95,22 @@ void ParquetBlockInputStream::fillColumnWithStringData(std::shared_ptr<arrow::Co
}
}
void ParquetBlockInputStream::fillColumnWithBooleanData(std::shared_ptr<arrow::Column> & arrow_column, MutableColumnPtr & internal_column)
{
PaddedPODArray<UInt8> & column_data = static_cast<ColumnVector<UInt8> &>(*internal_column).getData();
column_data.resize(arrow_column->length());
for (size_t chunk_i = 0; chunk_i != static_cast<size_t>(arrow_column->data()->num_chunks()); ++chunk_i)
{
arrow::BooleanArray & chunk = static_cast<arrow::BooleanArray &>(*(arrow_column->data()->chunk(chunk_i)));
/// buffers[0] is a null bitmap and buffers[1] are actual values
std::shared_ptr<arrow::Buffer> buffer = chunk.data()->buffers[1];
for (size_t bool_i = 0; bool_i != static_cast<size_t>(chunk.length()); ++bool_i)
column_data[bool_i] = chunk.Value(bool_i);
}
}
#define FOR_ARROW_NUMERIC_TYPES(M) \
M(arrow::Type::UINT8, UInt8) \
M(arrow::Type::INT8, Int8) \
@ -187,6 +203,10 @@ Block ParquetBlockInputStream::readImpl()
{
fillColumnWithStringData(arrow_column, read_column);
}
else if (arrow::Type::BOOL == arrow_type)
{
fillColumnWithBooleanData(arrow_column, read_column);
}
// TODO: check that values smaller than INT32 are being read correctly
#define DISPATCH(ARROW_NUMERIC_TYPE, CPP_NUMERIC_TYPE) \
else if (ARROW_NUMERIC_TYPE == arrow_type) \

View File

@ -33,6 +33,7 @@ private:
void fillColumnWithNumericData(std::shared_ptr<arrow::Column> & arrow_column, MutableColumnPtr & internal_column);
void fillColumnWithStringData(std::shared_ptr<arrow::Column> & arrow_column, MutableColumnPtr & internal_column);
void fillColumnWithBooleanData(std::shared_ptr<arrow::Column> & arrow_column, MutableColumnPtr & internal_column);
std::unordered_map<arrow::Type::type, std::shared_ptr<IDataType>> arrow_type_to_internal_type = {
{arrow::Type::UINT8, std::make_shared<DataTypeUInt8>()},
@ -46,9 +47,9 @@ private:
{arrow::Type::FLOAT, std::make_shared<DataTypeFloat32>()},
{arrow::Type::DOUBLE, std::make_shared<DataTypeFloat64>()},
{arrow::Type::STRING, std::make_shared<DataTypeString>()}//,
{arrow::Type::STRING, std::make_shared<DataTypeString>()},
{arrow::Type::BOOL, std::make_shared<DataTypeUInt8>()}//,
// TODO:
/* {arrow::Type::BOOL, std::make_shared<DataTypeUInt8>()}, */
/* {arrow::Type::DATE32, Date32, Int32Type}, */
/* {arrow::Type::DATE64, Date64, Int32Type}//, */
// TODO: add other types