Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Include row group level stats when writing ORC files #10041

Merged
merged 19 commits into from
Jan 19, 2022
Merged

Include row group level stats when writing ORC files #10041

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
19 commits
Select commit Hold shift + click to select a range
b8ae756
separate stats by level
vuule Jan 11, 2022
d6c18c8
encode rg stats
vuule Jan 13, 2022
cb2b972
rename putb
vuule Jan 13, 2022
95f018c
add put_bytes
vuule Jan 13, 2022
c5b62b9
include rg stats in rg index entries
vuule Jan 13, 2022
f1f6958
docs
vuule Jan 13, 2022
c078879
fix; don't use optional
vuule Jan 14, 2022
61733f9
Merge branch 'branch-22.02' of https://github.com/rapidsai/cudf into …
vuule Jan 14, 2022
4bfe814
return written size from put_byte and put_bytes
vuule Jan 14, 2022
2535754
reuse encode_field_number in protobuf writer
vuule Jan 15, 2022
683a016
copyright year
vuule Jan 15, 2022
61a8bec
comment
vuule Jan 15, 2022
fecf4d5
Merge branch 'branch-22.02' of https://github.com/rapidsai/cudf into …
vuule Jan 18, 2022
509138b
Merge branch 'branch-22.02' of https://github.com/rapidsai/cudf into …
vuule Jan 18, 2022
febcd04
Merge branch 'branch-22.02' of https://github.com/rapidsai/cudf into …
vuule Jan 19, 2022
4c82d89
host_span; static_assert
vuule Jan 19, 2022
2afedce
ProtobufType enum
vuule Jan 19, 2022
46a4a3c
style
vuule Jan 19, 2022
6ea0a50
copyright year
vuule Jan 19, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 34 additions & 23 deletions cpp/src/io/orc/orc.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2021, NVIDIA CORPORATION.
* Copyright (c) 2019-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -209,43 +209,54 @@ void ProtobufWriter::put_row_index_entry(int32_t present_blk,
int32_t data_ofs,
int32_t data2_blk,
int32_t data2_ofs,
TypeKind kind)
TypeKind kind,
ColStatsBlob const* stats)
{
size_t sz = 0, lpos;
putb(1 * 8 + PB_TYPE_FIXEDLEN); // 1:RowIndex.entry
put_uint(encode_field_number(1, PB_TYPE_FIXEDLEN)); // 1:RowIndex.entry
lpos = m_buf->size();
putb(0xcd); // sz+2
putb(1 * 8 + PB_TYPE_FIXEDLEN); // 1:positions[packed=true]
putb(0xcd); // sz
put_byte(0xcd); // sz+2
put_uint(encode_field_number(1, PB_TYPE_FIXEDLEN)); // 1:positions[packed=true]
put_byte(0xcd); // sz
if (present_blk >= 0) sz += put_uint(present_blk);
if (present_ofs >= 0) {
sz += put_uint(present_ofs) + 2;
putb(0); // run pos = 0
putb(0); // bit pos = 0
sz += put_uint(present_ofs);
sz += put_byte(0); // run pos = 0
sz += put_byte(0); // bit pos = 0
}
if (data_blk >= 0) { sz += put_uint(data_blk); }
if (data_ofs >= 0) {
sz += put_uint(data_ofs);
if (kind != STRING && kind != FLOAT && kind != DOUBLE && kind != DECIMAL) {
putb(0); // RLE run pos always zero (assumes RLE aligned with row index boundaries)
sz++;
// RLE run pos always zero (assumes RLE aligned with row index boundaries)
sz += put_byte(0);
if (kind == BOOLEAN) {
putb(0); // bit position in byte, always zero
sz++;
// bit position in byte, always zero
sz += put_byte(0);
}
}
}
if (kind !=
INT) // INT kind can be passed in to bypass 2nd stream index (dictionary length streams)
{
// INT kind can be passed in to bypass 2nd stream index (dictionary length streams)
if (kind != INT) {
if (data2_blk >= 0) { sz += put_uint(data2_blk); }
if (data2_ofs >= 0) {
sz += put_uint(data2_ofs) + 1;
putb(0); // RLE run pos always zero (assumes RLE aligned with row index boundaries)
sz += put_uint(data2_ofs);
// RLE run pos always zero (assumes RLE aligned with row index boundaries)
sz += put_byte(0);
}
}
m_buf->data()[lpos] = (uint8_t)(sz + 2);
// size of the field 1
m_buf->data()[lpos + 2] = (uint8_t)(sz);

if (stats != nullptr) {
sz += put_uint(encode_field_number<decltype(*stats)>(2)); // 2: statistics
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: maybe field number (2 in this case) should be an enum. I see that it's used in a lot of places though, so maybe a followup.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's doable, but there would need to be different enums for each ORC message type, since the numbers are not unique between messages (see https://orc.apache.org/specification/ORCv1/). We can have the set of enums (non-class) and still pass them as int. I would really need to do this in a follow up for this one to make it into 22.02.

// Statistics field contains its length as varint and dtype specific data (encoded on the GPU)
sz += put_uint(stats->size());
sz += put_bytes(*stats);
}

// size of the whole row index entry
m_buf->data()[lpos] = (uint8_t)(sz + 2);
}

size_t ProtobufWriter::write(const PostScript& s)
Expand All @@ -256,7 +267,7 @@ size_t ProtobufWriter::write(const PostScript& s)
if (s.compression != NONE) { w.field_uint(3, s.compressionBlockSize); }
w.field_packed_uint(4, s.version);
w.field_uint(5, s.metadataLength);
w.field_string(8000, s.magic);
w.field_blob(8000, s.magic);
return w.value();
}

Expand Down Expand Up @@ -300,8 +311,8 @@ size_t ProtobufWriter::write(const SchemaType& s)
size_t ProtobufWriter::write(const UserMetadataItem& s)
{
ProtobufFieldWriter w(this);
w.field_string(1, s.name);
w.field_string(2, s.value);
w.field_blob(1, s.name);
w.field_blob(2, s.value);
return w.value();
}

Expand All @@ -310,7 +321,7 @@ size_t ProtobufWriter::write(const StripeFooter& s)
ProtobufFieldWriter w(this);
w.field_repeated_struct(1, s.streams);
w.field_repeated_struct(2, s.columns);
if (s.writerTimezone != "") { w.field_string(3, s.writerTimezone); }
if (s.writerTimezone != "") { w.field_blob(3, s.writerTimezone); }
return w.value();
}

Expand Down
137 changes: 78 additions & 59 deletions cpp/src/io/orc/orc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2021, NVIDIA CORPORATION.
* Copyright (c) 2019-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -131,6 +131,67 @@ struct Metadata {
std::vector<StripeStatistics> stripeStats;
};

int static constexpr encode_field_number(int field_number, uint8_t field_type) noexcept
vuule marked this conversation as resolved.
Show resolved Hide resolved
{
return (field_number * 8) + field_type;
}

namespace {
template <typename base_t,
typename std::enable_if_t<!std::is_arithmetic<base_t>::value and
!std::is_enum<base_t>::value>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return encode_field_number(field_number, PB_TYPE_FIXEDLEN);
}

template <typename base_t,
typename std::enable_if_t<std::is_integral<base_t>::value or
std::is_enum<base_t>::value>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return encode_field_number(field_number, PB_TYPE_VARINT);
}

template <typename base_t, typename std::enable_if_t<std::is_same_v<base_t, float>>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return encode_field_number(field_number, PB_TYPE_FIXED32);
}

template <typename base_t, typename std::enable_if_t<std::is_same_v<base_t, double>>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return encode_field_number(field_number, PB_TYPE_FIXED64);
}
}; // namespace

template <
typename T,
typename std::enable_if_t<!std::is_class<T>::value or std::is_same_v<T, std::string>>* = nullptr>
int static constexpr encode_field_number(int field_number) noexcept
{
return encode_field_number_base<T>(field_number);
}

// containters change the field number encoding
template <
typename T,
typename std::enable_if_t<std::is_same<T, std::vector<typename T::value_type>>::value>* = nullptr>
int static constexpr encode_field_number(int field_number) noexcept
{
return encode_field_number_base<T>(field_number);
}

// optional fields don't change the field number encoding
template <typename T,
typename std::enable_if_t<
std::is_same<T, std::optional<typename T::value_type>>::value>* = nullptr>
int static constexpr encode_field_number(int field_number) noexcept
{
return encode_field_number_base<typename T::value_type>(field_number);
}

/**
* @brief Class for parsing Orc's Protocol Buffers encoded metadata
*/
Expand Down Expand Up @@ -181,60 +242,6 @@ class ProtobufReader {
template <typename T, typename... Operator>
void function_builder(T& s, size_t maxlen, std::tuple<Operator...>& op);

template <typename base_t,
typename std::enable_if_t<!std::is_arithmetic<base_t>::value and
!std::is_enum<base_t>::value>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return (field_number * 8) + PB_TYPE_FIXEDLEN;
}

template <typename base_t,
typename std::enable_if_t<std::is_integral<base_t>::value or
std::is_enum<base_t>::value>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return (field_number * 8) + PB_TYPE_VARINT;
}

template <typename base_t, typename std::enable_if_t<std::is_same_v<base_t, float>>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return (field_number * 8) + PB_TYPE_FIXED32;
}

template <typename base_t, typename std::enable_if_t<std::is_same_v<base_t, double>>* = nullptr>
int static constexpr encode_field_number_base(int field_number) noexcept
{
return (field_number * 8) + PB_TYPE_FIXED64;
}

template <typename T,
typename std::enable_if_t<!std::is_class<T>::value or std::is_same_v<T, std::string>>* =
nullptr>
int static constexpr encode_field_number(int field_number) noexcept
{
return encode_field_number_base<T>(field_number);
}

// containters change the field number encoding
template <typename T,
typename std::enable_if_t<
std::is_same<T, std::vector<typename T::value_type>>::value>* = nullptr>
int static constexpr encode_field_number(int field_number) noexcept
{
return encode_field_number_base<T>(field_number);
}

// optional fields don't change the field number encoding
template <typename T,
typename std::enable_if_t<
std::is_same<T, std::optional<typename T::value_type>>::value>* = nullptr>
int static constexpr encode_field_number(int field_number) noexcept
{
return encode_field_number_base<typename T::value_type>(field_number);
}

uint32_t read_field_size(const uint8_t* end);

template <typename T, typename std::enable_if_t<std::is_integral<T>::value>* = nullptr>
Expand Down Expand Up @@ -470,16 +477,27 @@ class ProtobufWriter {
public:
ProtobufWriter() { m_buf = nullptr; }
ProtobufWriter(std::vector<uint8_t>* output) { m_buf = output; }
void putb(uint8_t v) { m_buf->push_back(v); }
uint32_t put_byte(uint8_t v)
{
m_buf->push_back(v);
return 1;
}
template <typename T>
uint32_t put_bytes(T const& values)
vuule marked this conversation as resolved.
Show resolved Hide resolved
{
m_buf->reserve(m_buf->size() + values.size());
m_buf->insert(m_buf->end(), values.cbegin(), values.cend());
return values.size();
}
uint32_t put_uint(uint64_t v)
{
int l = 1;
while (v > 0x7f) {
putb(static_cast<uint8_t>(v | 0x80));
put_byte(static_cast<uint8_t>(v | 0x80));
v >>= 7;
l++;
}
putb(static_cast<uint8_t>(v));
put_byte(static_cast<uint8_t>(v));
return l;
}
uint32_t put_int(int64_t v)
Expand All @@ -493,7 +511,8 @@ class ProtobufWriter {
int32_t data_ofs,
int32_t data2_blk,
int32_t data2_ofs,
TypeKind kind);
TypeKind kind,
ColStatsBlob const* stats);

public:
size_t write(const PostScript&);
Expand Down
36 changes: 11 additions & 25 deletions cpp/src/io/orc/orc_field_writer.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2020, NVIDIA CORPORATION.
* Copyright (c) 2020-2022, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -41,7 +41,7 @@ struct ProtobufWriter::ProtobufFieldWriter {
template <typename T>
void field_uint(int field, const T& value)
{
struct_size += p->put_uint(field * 8 + PB_TYPE_VARINT);
struct_size += p->put_uint(encode_field_number<T>(field));
struct_size += p->put_uint(static_cast<uint64_t>(value));
}

Expand All @@ -52,9 +52,9 @@ struct ProtobufWriter::ProtobufFieldWriter {
template <typename T>
void field_packed_uint(int field, const std::vector<T>& value)
{
struct_size += p->put_uint(field * 8 + PB_TYPE_FIXEDLEN);
struct_size += p->put_uint(encode_field_number<std::vector<T>>(field));
auto lpos = p->m_buf->size();
p->putb(0);
p->put_byte(0);
auto sz = std::accumulate(value.begin(), value.end(), 0, [p = this->p](size_t sum, auto val) {
return sum + p->put_uint(val);
});
Expand All @@ -65,29 +65,15 @@ struct ProtobufWriter::ProtobufFieldWriter {
(*(p->m_buf))[lpos] = static_cast<uint8_t>(sz);
}

/**
* @brief Function to write a string to the internal buffer
*/
void field_string(int field, const std::string& value)
{
size_t len = value.length();
struct_size += p->put_uint(field * 8 + PB_TYPE_FIXEDLEN);
struct_size += p->put_uint(len) + len;
for (size_t i = 0; i < len; i++)
p->putb(value[i]);
}

/**
* @brief Function to write a blob to the internal buffer
*/
template <typename T>
void field_blob(int field, const std::vector<T>& value)
void field_blob(int field, T const& values)
{
size_t len = value.size();
struct_size += p->put_uint(field * 8 + PB_TYPE_FIXEDLEN);
struct_size += p->put_uint(len) + len;
for (size_t i = 0; i < len; i++)
p->putb(value[i]);
struct_size += p->put_uint(encode_field_number<T>(field));
struct_size += p->put_uint(values.size());
struct_size += p->put_bytes(values);
}

/**
Expand All @@ -96,9 +82,9 @@ struct ProtobufWriter::ProtobufFieldWriter {
template <typename T>
void field_struct(int field, const T& value)
{
struct_size += p->put_uint((field)*8 + PB_TYPE_FIXEDLEN);
struct_size += p->put_uint(encode_field_number(field, PB_TYPE_FIXEDLEN));
auto lpos = p->m_buf->size();
p->putb(0);
p->put_byte(0);
auto sz = p->write(value);
struct_size += sz + 1;
for (; sz > 0x7f; sz >>= 7, struct_size++)
Expand All @@ -112,7 +98,7 @@ struct ProtobufWriter::ProtobufFieldWriter {
void field_repeated_string(int field, const std::vector<std::string>& value)
{
for (const auto& elem : value)
field_string(field, elem);
field_blob(field, elem);
}

/**
Expand Down
Loading