From 29616cca0926100dfc9d0616fb85dc75d66348a8 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Fri, 4 Sep 2020 23:34:37 +0000 Subject: [PATCH 01/23] ARROW-8494: [C++] Full support for mixed lista and structs Also: ARROW-9810 (generalize rep/def level conversion to list lengths/bitmaps) This adds helper methods for reconstructing all necessary metadata for arrow types. For now this doesn't handle null_slot_usage (i.e. children of FixedSizeList), it throws exceptions when nulls are encountered in this case. It uses there for generic reconstruction. The unit tests demonstrate how to use the helper methods in combination with LevelInfo (generated from parquet/arrow/schema.h) to reconstruct the metadata. The arrow reader.cc is now rewritten to use these method. - Refactors necessary APIs to use LevelInfo and makes use of them in column_reader - Adds implementations for reconstructing list validity bitmaps (uses rep/def levels) - Adds implementations for reconstruction list lengths (uses rep/def levels.). - Adds dynamic dispatch for level comparison algorithms for AVX2 and BMI2. - Adds a pextract alternative that uses BitRunReader that can be used as a fallback. - Fixes some bugs in detailed reconstruction to array tests. --- cpp/cmake_modules/SetupCxxFlags.cmake | 6 +- cpp/src/arrow/util/cpu_info.h | 5 + cpp/src/parquet/CMakeLists.txt | 14 + .../parquet/arrow/arrow_reader_writer_test.cc | 12 +- cpp/src/parquet/arrow/reader.cc | 433 +++++++++--------- cpp/src/parquet/arrow/reader_internal.cc | 113 ----- cpp/src/parquet/arrow/reader_internal.h | 7 - .../arrow/reconstruct_internal_test.cc | 151 +++--- cpp/src/parquet/arrow/schema.cc | 1 - cpp/src/parquet/column_reader.cc | 114 +++-- cpp/src/parquet/column_reader.h | 23 +- cpp/src/parquet/level_comparison.cc | 71 +++ cpp/src/parquet/level_comparison.h | 91 ++++ cpp/src/parquet/level_comparison_avx2.cc | 33 ++ cpp/src/parquet/level_conversion.cc | 267 ++++++----- cpp/src/parquet/level_conversion.h | 82 ++-- cpp/src/parquet/level_conversion_benchmark.cc | 6 +- cpp/src/parquet/level_conversion_bmi2.cc | 35 ++ cpp/src/parquet/level_conversion_inc.h | 130 ++++++ cpp/src/parquet/level_conversion_test.cc | 238 +++++++++- 20 files changed, 1153 insertions(+), 679 deletions(-) create mode 100644 cpp/src/parquet/level_comparison.cc create mode 100644 cpp/src/parquet/level_comparison.h create mode 100644 cpp/src/parquet/level_comparison_avx2.cc create mode 100644 cpp/src/parquet/level_conversion_bmi2.cc create mode 100644 cpp/src/parquet/level_conversion_inc.h diff --git a/cpp/cmake_modules/SetupCxxFlags.cmake b/cpp/cmake_modules/SetupCxxFlags.cmake index a3d4062f0eb..96e79c69fc4 100644 --- a/cpp/cmake_modules/SetupCxxFlags.cmake +++ b/cpp/cmake_modules/SetupCxxFlags.cmake @@ -50,7 +50,7 @@ if(ARROW_CPU_FLAG STREQUAL "x86") # skylake-avx512 consists of AVX512F,AVX512BW,AVX512VL,AVX512CD,AVX512DQ set(ARROW_AVX512_FLAG "-march=skylake-avx512 -mbmi2") # Append the avx2/avx512 subset option also, fix issue ARROW-9877 for homebrew-cpp - set(ARROW_AVX2_FLAG "${ARROW_AVX2_FLAG} -mavx2") + set(ARROW_AVX2_FLAG "${ARROW_AVX2_FLAG} -mavx2 -mbmi2") set(ARROW_AVX512_FLAG "${ARROW_AVX512_FLAG} -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw") check_cxx_compiler_flag(${ARROW_SSE4_2_FLAG} CXX_SUPPORTS_SSE4_2) @@ -68,10 +68,10 @@ if(ARROW_CPU_FLAG STREQUAL "x86") add_definitions(-DARROW_HAVE_RUNTIME_SSE4_2) endif() if(CXX_SUPPORTS_AVX2 AND ARROW_RUNTIME_SIMD_LEVEL MATCHES "^(AVX2|AVX512|MAX)$") - add_definitions(-DARROW_HAVE_RUNTIME_AVX2) + add_definitions(-DARROW_HAVE_RUNTIME_AVX2 -DARROW_HAVE_RUNTIME_BMI2) endif() if(CXX_SUPPORTS_AVX512 AND ARROW_RUNTIME_SIMD_LEVEL MATCHES "^(AVX512|MAX)$") - add_definitions(-DARROW_HAVE_RUNTIME_AVX512) + add_definitions(-DARROW_HAVE_RUNTIME_AVX512 -DARROW_HAVE_RUNTIME_BMI2) endif() elseif(ARROW_CPU_FLAG STREQUAL "ppc") # power compiler flags, gcc/clang only diff --git a/cpp/src/arrow/util/cpu_info.h b/cpp/src/arrow/util/cpu_info.h index 73695b7742e..a57ffd29467 100644 --- a/cpp/src/arrow/util/cpu_info.h +++ b/cpp/src/arrow/util/cpu_info.h @@ -106,6 +106,11 @@ class ARROW_EXPORT CpuInfo { /// Returns the vendor of the cpu. Vendor vendor() const { return vendor_; } + bool HasEfficientBmi2() const { + // BMI2 (pext, pdep) is only efficient on Intel X86 processors. + return vendor() == Vendor::Intel && IsSupported(BMI2); + } + private: CpuInfo(); diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index b70fe6b168d..3722a229338 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -188,6 +188,7 @@ set(PARQUET_SRCS file_writer.cc internal_file_decryptor.cc internal_file_encryptor.cc + level_comparison.cc level_conversion.cc metadata.cc murmur3.cc @@ -202,6 +203,19 @@ set(PARQUET_SRCS stream_writer.cc types.cc) +if(CXX_SUPPORTS_AVX2) + # AVX2 is used as a proxy for BMI2. + list(APPEND PARQUET_SRCS level_comparison_avx2.cc level_conversion_bmi2.cc) + set_source_files_properties(level_comparison_avx2.cc + level_conversion_bmi2.cc + PROPERTIES + SKIP_PRECOMPILE_HEADERS + ON + COMPILE_FLAGS + "${ARROW_AVX2_FLAG} -DARROW_HAVE_BMI2") + +endif() + if(PARQUET_REQUIRE_ENCRYPTION) set(PARQUET_SRCS ${PARQUET_SRCS} encryption_internal.cc) else() diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index 476d82f7fac..a23f6e559a2 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -416,7 +416,6 @@ void DoSimpleRoundtrip(const std::shared_ptr& table, bool use_threads, ::arrow::default_memory_pool(), &reader)); reader->set_use_threads(use_threads); - if (column_subset.size() > 0) { ASSERT_OK_NO_THROW(reader->ReadTable(column_subset, out)); } else { @@ -2362,7 +2361,7 @@ TEST(ArrowReadWrite, SingleColumnNullableStruct) { } // Disabled until implementation can be finished. -TEST(TestArrowReadWrite, DISABLED_CanonicalNestedRoundTrip) { +TEST(TestArrowReadWrite, CanonicalNestedRoundTrip) { auto doc_id = field("DocId", ::arrow::int64(), /*nullable=*/false); auto links = field( "Links", @@ -2391,7 +2390,7 @@ TEST(TestArrowReadWrite, DISABLED_CanonicalNestedRoundTrip) { // string literals implemented properly auto name_array = ::arrow::ArrayFromJSON( name->type(), - "([[{\"Language\": [{\"Code\": \"en_us\", \"Country\":\"us\"}," + "[[{\"Language\": [{\"Code\": \"en_us\", \"Country\":\"us\"}," "{\"Code\": \"en_us\", \"Country\": null}]," "\"Url\": \"http://A\"}," "{\"Url\": \"http://B\"}," @@ -2810,12 +2809,6 @@ TEST_F(TestNestedSchemaRead, ReadTablePartial) { ASSERT_NO_FATAL_FAILURE(ValidateTableArrayTypes(*table)); } -TEST_F(TestNestedSchemaRead, StructAndListTogetherUnsupported) { - ASSERT_NO_FATAL_FAILURE(CreateSimpleNestedParquet(Repetition::REPEATED)); - std::shared_ptr
table; - ASSERT_RAISES(NotImplemented, reader_->ReadTable(&table)); -} - TEST_P(TestNestedSchemaRead, DeepNestedSchemaRead) { #ifdef PARQUET_VALGRIND const int num_trees = 3; @@ -2994,7 +2987,6 @@ TEST_P(TestArrowReaderAdHocSparkAndHvr, ReadDecimals) { ASSERT_OK(builder.Append(value)); } ASSERT_OK(builder.Finish(&expected_array)); - AssertArraysEqual(*expected_array, *chunk); } diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 5f13259058d..2da90ab1e18 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -43,6 +43,7 @@ #include "parquet/schema.h" using arrow::Array; +using arrow::ArrayData; using arrow::BooleanArray; using arrow::ChunkedArray; using arrow::DataType; @@ -59,10 +60,6 @@ using arrow::Table; using arrow::TimestampArray; using arrow::internal::Iota; -using parquet::schema::GroupNode; -using parquet::schema::Node; -using parquet::schema::PrimitiveNode; - // Help reduce verbosity using ParquetReader = parquet::ParquetFileReader; @@ -77,18 +74,43 @@ using parquet::internal::RecordReader; namespace parquet { namespace arrow { +namespace { + +::arrow::Result> ChunksToSingle(const ChunkedArray& chunked) { + switch (chunked.num_chunks()) { + case 0: { + ARROW_ASSIGN_OR_RAISE(std::shared_ptr array, + ::arrow::MakeArrayOfNull(chunked.type(), 0)); + return array->data(); + } + case 1: + return chunked.chunk(0)->data(); + default: + // ARROW-3762(wesm): If item reader yields a chunked array, we reject as + // this is not yet implemented + return Status::NotImplemented( + "Nested data conversions not implemented for chunked array outputs"); + } +} + +} // namespace class ColumnReaderImpl : public ColumnReader { public: - enum ReaderType { PRIMITIVE, LIST, STRUCT }; - virtual Status GetDefLevels(const int16_t** data, int64_t* length) = 0; virtual Status GetRepLevels(const int16_t** data, int64_t* length) = 0; virtual const std::shared_ptr field() = 0; + ::arrow::Status NextBatch(int64_t batch_size, + std::shared_ptr<::arrow::ChunkedArray>* out) final { + RETURN_NOT_OK(LoadBatch(batch_size)); + return BuildArray(batch_size, out); + } - virtual const ColumnDescriptor* descr() const = 0; + virtual ::arrow::Status LoadBatch(int64_t num_records) = 0; - virtual ReaderType type() const = 0; + virtual ::arrow::Status BuildArray(int64_t number_of_slots, + std::shared_ptr<::arrow::ChunkedArray>* out) = 0; + virtual bool IsOrHasRepeatedChild() const = 0; }; std::shared_ptr> VectorToSharedSet( @@ -373,30 +395,34 @@ class RowGroupReaderImpl : public RowGroupReader { class LeafReader : public ColumnReaderImpl { public: LeafReader(std::shared_ptr ctx, std::shared_ptr field, - std::unique_ptr input) + std::unique_ptr input, + ::parquet::internal::LevelInfo leaf_info) : ctx_(std::move(ctx)), field_(std::move(field)), input_(std::move(input)), descr_(input_->descr()) { record_reader_ = RecordReader::Make( - descr_, ctx_->pool, field_->type()->id() == ::arrow::Type::DICTIONARY); + descr_, leaf_info, ctx_->pool, field_->type()->id() == ::arrow::Type::DICTIONARY); NextRowGroup(); } - Status GetDefLevels(const int16_t** data, int64_t* length) override { + Status GetDefLevels(const int16_t** data, int64_t* length) final { *data = record_reader_->def_levels(); *length = record_reader_->levels_position(); return Status::OK(); } - Status GetRepLevels(const int16_t** data, int64_t* length) override { + Status GetRepLevels(const int16_t** data, int64_t* length) final { *data = record_reader_->rep_levels(); *length = record_reader_->levels_position(); return Status::OK(); } - Status NextBatch(int64_t records_to_read, std::shared_ptr* out) override { + bool IsOrHasRepeatedChild() const final { return false; } + + Status LoadBatch(int64_t records_to_read) final { BEGIN_PARQUET_CATCH_EXCEPTIONS + out_ = nullptr; record_reader_->Reset(); // Pre-allocation gives much better performance for flat columns record_reader_->Reserve(records_to_read); @@ -411,17 +437,22 @@ class LeafReader : public ColumnReaderImpl { } } RETURN_NOT_OK(TransferColumnData(record_reader_.get(), field_->type(), descr_, - ctx_->pool, out)); + ctx_->pool, &out_)); return Status::OK(); END_PARQUET_CATCH_EXCEPTIONS } - const std::shared_ptr field() override { return field_; } - const ColumnDescriptor* descr() const override { return descr_; } + ::arrow::Status BuildArray(int64_t expected_array_entries, + std::shared_ptr<::arrow::ChunkedArray>* out) final { + *out = out_; + return Status::OK(); + ; + } - ReaderType type() const override { return PRIMITIVE; } + const std::shared_ptr field() override { return field_; } private: + std::shared_ptr out_; void NextRowGroup() { std::unique_ptr page_reader = input_->NextChunk(); record_reader_->SetPageReader(std::move(page_reader)); @@ -434,16 +465,16 @@ class LeafReader : public ColumnReaderImpl { std::shared_ptr record_reader_; }; -class NestedListReader : public ColumnReaderImpl { +template +class ListReader : public ColumnReaderImpl { public: - NestedListReader(std::shared_ptr ctx, std::shared_ptr field, - int16_t max_definition_level, int16_t max_repetition_level, - std::unique_ptr item_reader) + ListReader(std::shared_ptr ctx, std::shared_ptr field, + ::parquet::internal::LevelInfo level_info, + std::unique_ptr child_reader) : ctx_(std::move(ctx)), field_(std::move(field)), - max_definition_level_(max_definition_level), - max_repetition_level_(max_repetition_level), - item_reader_(std::move(item_reader)) {} + level_info_(level_info), + item_reader_(std::move(child_reader)) {} Status GetDefLevels(const int16_t** data, int64_t* length) override { return item_reader_->GetDefLevels(data, length); @@ -453,229 +484,222 @@ class NestedListReader : public ColumnReaderImpl { return item_reader_->GetRepLevels(data, length); } - Status NextBatch(int64_t records_to_read, std::shared_ptr* out) override { - if (item_reader_->type() == ColumnReaderImpl::STRUCT) { - return Status::Invalid("Mix of struct and list types not yet supported"); - } - - RETURN_NOT_OK(item_reader_->NextBatch(records_to_read, out)); + bool IsOrHasRepeatedChild() const final { return true; } - std::shared_ptr item_chunk; - switch ((*out)->num_chunks()) { - case 0: { - ARROW_ASSIGN_OR_RAISE(item_chunk, ::arrow::MakeArrayOfNull((*out)->type(), 0)); - break; - } - case 1: - item_chunk = (*out)->chunk(0); - break; - default: - // ARROW-3762(wesm): If item reader yields a chunked array, we reject as - // this is not yet implemented - return Status::NotImplemented( - "Nested data conversions not implemented for chunked array outputs"); - } + Status LoadBatch(int64_t number_of_records) final { + return item_reader_->LoadBatch(number_of_records); + } + Status BuildArray(int64_t number_of_slots, + std::shared_ptr* out) override { const int16_t* def_levels; const int16_t* rep_levels; int64_t num_levels; RETURN_NOT_OK(item_reader_->GetDefLevels(&def_levels, &num_levels)); RETURN_NOT_OK(item_reader_->GetRepLevels(&rep_levels, &num_levels)); - std::shared_ptr result; - RETURN_NOT_OK(ReconstructNestedList(item_chunk, field_, max_definition_level_, - max_repetition_level_, def_levels, rep_levels, - num_levels, ctx_->pool, &result)); + std::shared_ptr validity_buffer; + ::parquet::internal::ValidityBitmapInputOutput validity_io; + if (field_->nullable()) { + ARROW_ASSIGN_OR_RAISE(validity_buffer, AllocateBitmap(number_of_slots, ctx_->pool)); + + validity_io.valid_bits = validity_buffer->mutable_data(); + } + ARROW_ASSIGN_OR_RAISE( + std::shared_ptr lengths_buffer, + AllocateBuffer(sizeof(IndexType) * std::max(int64_t{2}, number_of_slots + 1), + ctx_->pool)); + // ensure zero initialization in case we have reached a zero length list (and + // because first entry is always zero). + IndexType* length_data = reinterpret_cast(lengths_buffer->mutable_data()); + std::fill(length_data, length_data + 2, 0); + BEGIN_PARQUET_CATCH_EXCEPTIONS + ::parquet::internal::ConvertDefRepLevelsToList( + def_levels, rep_levels, num_levels, level_info_, &validity_io, length_data); + END_PARQUET_CATCH_EXCEPTIONS + for (int x = 1; x <= validity_io.values_read; x++) { + length_data[x] += length_data[x - 1]; + } + // We might return less then the requested slot (i.e. reaching an end of a file) + // ensure we've set all the bits here. + if (validity_io.values_read < number_of_slots) { + // + 1 because arrays lengths are values + 1 + std::fill(length_data + validity_io.values_read + 1, + length_data + number_of_slots + 1, 0); + if (validity_io.valid_bits != nullptr) { + std::fill(validity_io.valid_bits + BitUtil::BytesForBits(validity_io.values_read), + validity_io.valid_bits + BitUtil::BytesForBits(number_of_slots), 0); + } + } + + RETURN_NOT_OK(item_reader_->BuildArray(length_data[validity_io.values_read], out)); + ARROW_ASSIGN_OR_RAISE(std::shared_ptr item_chunk, ChunksToSingle(**out)); + + std::vector> buffers{ + validity_io.null_count > 0 ? validity_buffer : std::shared_ptr(), + lengths_buffer}; + auto data = std::make_shared( + field_->type(), + /*length=*/validity_io.values_read, std::move(buffers), + std::vector>{item_chunk}, validity_io.null_count); + + std::shared_ptr result = ::arrow::MakeArray(data); + RETURN_NOT_OK(result->Validate()); *out = std::make_shared(result); return Status::OK(); } const std::shared_ptr field() override { return field_; } - const ColumnDescriptor* descr() const override { return nullptr; } - - ReaderType type() const override { return LIST; } - private: std::shared_ptr ctx_; std::shared_ptr field_; - int16_t max_definition_level_; - int16_t max_repetition_level_; + ::parquet::internal::LevelInfo level_info_; std::unique_ptr item_reader_; }; class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { public: explicit StructReader(std::shared_ptr ctx, - const SchemaField& schema_field, std::shared_ptr filtered_field, + ::parquet::internal::LevelInfo level_info, std::vector>&& children) : ctx_(std::move(ctx)), - schema_field_(schema_field), filtered_field_(std::move(filtered_field)), - struct_def_level_(schema_field.level_info.def_level), - children_(std::move(children)) {} + level_info_(level_info), + children_(std::move(children)) { + has_repeated_child_ = IsOrHasRepeatedChild(); + } - Status NextBatch(int64_t records_to_read, std::shared_ptr* out) override; + bool IsOrHasRepeatedChild() const final { + if (children_.empty()) { + return false; + } + // Needs to be kept consistent with how rep/def levels are chosen. + return children_.front()->IsOrHasRepeatedChild(); + } + + Status LoadBatch(int64_t records_to_read) override { + for (const std::unique_ptr& reader : children_) { + RETURN_NOT_OK(reader->LoadBatch(records_to_read)); + } + return Status::OK(); + } + Status BuildArray(int64_t records_to_read, std::shared_ptr* out) override; Status GetDefLevels(const int16_t** data, int64_t* length) override; Status GetRepLevels(const int16_t** data, int64_t* length) override; const std::shared_ptr field() override { return filtered_field_; } - const ColumnDescriptor* descr() const override { return nullptr; } - ReaderType type() const override { return STRUCT; } private: - std::shared_ptr ctx_; - SchemaField schema_field_; - std::shared_ptr filtered_field_; - int16_t struct_def_level_; - std::vector> children_; - std::shared_ptr def_levels_buffer_; - Status DefLevelsToNullArray(std::shared_ptr* null_bitmap, int64_t* null_count); + const std::shared_ptr ctx_; + const std::shared_ptr filtered_field_; + const ::parquet::internal::LevelInfo level_info_; + const std::vector> children_; + bool has_repeated_child_; }; -Status StructReader::DefLevelsToNullArray(std::shared_ptr* null_bitmap_out, - int64_t* null_count_out) { - auto null_count = 0; - const int16_t* def_levels_data; - int64_t def_levels_length; - RETURN_NOT_OK(GetDefLevels(&def_levels_data, &def_levels_length)); - ARROW_ASSIGN_OR_RAISE(auto null_bitmap, - AllocateEmptyBitmap(def_levels_length, ctx_->pool)); - uint8_t* null_bitmap_ptr = null_bitmap->mutable_data(); - for (int64_t i = 0; i < def_levels_length; i++) { - if (def_levels_data[i] < struct_def_level_) { - // Mark null - null_count += 1; - } else { - DCHECK_EQ(def_levels_data[i], struct_def_level_); - ::arrow::BitUtil::SetBit(null_bitmap_ptr, i); - } - } - - *null_count_out = null_count; - *null_bitmap_out = (null_count == 0) ? nullptr : null_bitmap; - return Status::OK(); -} - -// TODO(itaiin): Consider caching the results of this calculation - -// note that this is only used once for each read for now Status StructReader::GetDefLevels(const int16_t** data, int64_t* length) { *data = nullptr; if (children_.size() == 0) { - // Empty struct *length = 0; - return Status::OK(); + return Status::Invalid("StructReader had no childre"); } - // We have at least one child - const int16_t* child_def_levels; - int64_t child_length = 0; - bool found_nullable_child = false; - int16_t* result_levels = nullptr; - - int child_index = 0; - while (child_index < static_cast(children_.size())) { - if (!children_[child_index]->field()->nullable()) { - ++child_index; - continue; - } - RETURN_NOT_OK(children_[child_index]->GetDefLevels(&child_def_levels, &child_length)); - auto size = child_length * sizeof(int16_t); - ARROW_ASSIGN_OR_RAISE(def_levels_buffer_, AllocateResizableBuffer(size, ctx_->pool)); - // Initialize with the minimal def level - std::memset(def_levels_buffer_->mutable_data(), -1, size); - result_levels = reinterpret_cast(def_levels_buffer_->mutable_data()); - found_nullable_child = true; - break; - } + // This method should only be called when this struct or one of its parents + // are optional/repeated. Meaning all children must have rep/def levels associated + // with them. Furthermore, it shouldn't matter which definition levels we + // choose at this point, because callers sitting above this on the call graph + // should all have identical levels for the purposes of decoding (technically + // we could optimize by choosing a child with the least number of rep/def levels + // but we can leave that for future work). + RETURN_NOT_OK(children_[0]->GetDefLevels(data, length)); - if (!found_nullable_child) { + if (*data == nullptr) { + // Only happens if there are actually 0 rows available. *data = nullptr; *length = 0; - return Status::OK(); } + return Status::OK(); +} - // Look at the rest of the children - - // When a struct is defined, all of its children def levels are at least at - // nesting level, and def level equals nesting level. - // When a struct is not defined, all of its children def levels are less than - // the nesting level, and the def level equals max(children def levels) - // All other possibilities are malformed definition data. - for (; child_index < static_cast(children_.size()); ++child_index) { - // Child is non-nullable, and therefore has no definition levels - if (!children_[child_index]->field()->nullable()) { - continue; - } +Status StructReader::GetRepLevels(const int16_t** data, int64_t* length) { + *data = nullptr; + if (children_.size() == 0) { + *length = 0; + return Status::Invalid("StructReader had no childre"); + } - auto& child = children_[child_index]; - int64_t current_child_length; - RETURN_NOT_OK(child->GetDefLevels(&child_def_levels, ¤t_child_length)); - - if (child_length != current_child_length) { - std::stringstream ss; - ss << "Parquet struct decoding error. Expected to decode " << child_length - << " definition levels" - << " from child field \"" << child->field()->ToString() << "\" in parent \"" - << this->field()->ToString() << "\" but was only able to decode " - << current_child_length; - return Status::IOError(ss.str()); - } + // This method should only be called when this struct or one of its parents + // are optional/repeated. Meaning all children must have rep/def levels associated + // with them. Furthermore, it shouldn't matter which repetition levels we + // choose at this point, because callers sitting above this on the call graph + // should all have identical levels for the purposes of decoding them. + RETURN_NOT_OK(children_[0]->GetRepLevels(data, length)); - DCHECK_EQ(child_length, current_child_length); - for (int64_t i = 0; i < child_length; i++) { - // Check that value is either uninitialized, or current - // and previous children def levels agree on the struct level - DCHECK((result_levels[i] == -1) || ((result_levels[i] >= struct_def_level_) == - (child_def_levels[i] >= struct_def_level_))); - result_levels[i] = - std::max(result_levels[i], std::min(child_def_levels[i], struct_def_level_)); - } + if (data == nullptr) { + // Only happens if there are actually 0 rows available. + *data = nullptr; + *length = 0; } - *data = reinterpret_cast(def_levels_buffer_->data()); - *length = static_cast(child_length); return Status::OK(); } -Status StructReader::GetRepLevels(const int16_t** data, int64_t* length) { - return Status::NotImplemented("GetRepLevels is not implemented for struct"); -} - -Status StructReader::NextBatch(int64_t records_to_read, - std::shared_ptr* out) { - std::vector> children_arrays; +Status StructReader::BuildArray(int64_t records_to_read, + std::shared_ptr* out) { + std::vector> children_array_data; std::shared_ptr null_bitmap; - int64_t null_count; + int64_t null_count = 0; + int64_t values_read = records_to_read; + + BEGIN_PARQUET_CATCH_EXCEPTIONS + const int16_t* def_levels; + const int16_t* rep_levels; + int64_t num_levels; + + if (has_repeated_child_) { + ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); + RETURN_NOT_OK(children_.front()->GetDefLevels(&def_levels, &num_levels)); + RETURN_NOT_OK(children_.front()->GetRepLevels(&rep_levels, &num_levels)); + + ::parquet::internal::ValidityBitmapInputOutput validity_io; + validity_io.valid_bits = null_bitmap->mutable_data(); + ConvertDefRepLevelsToBitmap(def_levels, rep_levels, num_levels, level_info_, + &validity_io); + null_count = validity_io.null_count; + values_read = validity_io.values_read; + } else if (filtered_field_->nullable()) { + ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); + RETURN_NOT_OK(children_.front()->GetDefLevels(&def_levels, &num_levels)); + DefinitionLevelsToBitmap(def_levels, num_levels, level_info_, &values_read, + &null_count, null_bitmap->mutable_data(), + /*valid_bits_offset=*/0); + // Ensure all values are initialized + } + if (BitUtil::BytesForBits(values_read) < BitUtil::BytesForBits(records_to_read)) { + std::fill(null_bitmap->mutable_data() + BitUtil::BytesForBits(values_read), + null_bitmap->mutable_data() + BitUtil::BytesForBits(records_to_read), 0); + } + END_PARQUET_CATCH_EXCEPTIONS // Gather children arrays and def levels for (auto& child : children_) { - if (child->type() == ColumnReaderImpl::LIST) { - return Status::NotImplemented( - "Reading structs of lists from Parquet files not yet supported: ", - field()->ToString()); - } - std::shared_ptr field; - RETURN_NOT_OK(child->NextBatch(records_to_read, &field)); - - if (field->num_chunks() > 1) { - return Status::Invalid("Chunked field reads not yet supported with StructArray"); - } - children_arrays.push_back(field->chunk(0)); + RETURN_NOT_OK(child->BuildArray(values_read, &field)); + ARROW_ASSIGN_OR_RAISE(std::shared_ptr array_data, ChunksToSingle(*field)); + children_array_data.push_back(std::move(array_data)); } - RETURN_NOT_OK(DefLevelsToNullArray(&null_bitmap, &null_count)); - - int64_t struct_length = children_arrays[0]->length(); - for (size_t i = 1; i < children_arrays.size(); ++i) { - if (children_arrays[i]->length() != struct_length) { - // TODO(wesm): This should really only occur if the Parquet file is - // malformed. Should this be a DCHECK? - return Status::Invalid("Struct children had different lengths"); - } + if (!filtered_field_->nullable() && !has_repeated_child_) { + values_read = children_array_data.front()->length; } - auto result = std::make_shared(field()->type(), struct_length, - children_arrays, null_bitmap, null_count); + std::vector> buffers{ + null_count > 0 ? null_bitmap : std::shared_ptr()}; + auto data = std::make_shared(filtered_field_->type(), + /*length=*/values_read, std::move(buffers), + std::move(children_array_data)); + std::shared_ptr result = ::arrow::MakeArray(data); + RETURN_NOT_OK(result->Validate()); + *out = std::make_shared(result); return Status::OK(); } @@ -692,40 +716,28 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& if (!field.is_leaf()) { return Status::Invalid("Parquet non-leaf node has no children"); } + if (!ctx->IncludesLeaf(field.column_index)) { + *out = nullptr; + return Status::OK(); + } std::unique_ptr input( ctx->iterator_factory(field.column_index, ctx->reader)); - out->reset(new LeafReader(ctx, field.field, std::move(input))); + out->reset(new LeafReader(ctx, field.field, std::move(input), field.level_info)); } else if (type_id == ::arrow::Type::LIST) { - // We can only read lists-of-lists or structs at the moment auto list_field = field.field; auto child = &field.children[0]; - while (child->field->type()->id() == ::arrow::Type::LIST) { - child = &child->children[0]; - } - if (child->field->type()->id() == ::arrow::Type::STRUCT) { - return Status::NotImplemented( - "Reading lists of structs from Parquet files " - "not yet supported: ", - field.field->ToString()); - } - if (!ctx->IncludesLeaf(child->column_index)) { + std::unique_ptr child_reader; + RETURN_NOT_OK(GetReader(*child, ctx, &child_reader)); + if (child_reader == nullptr) { *out = nullptr; return Status::OK(); } - std::unique_ptr child_reader; - RETURN_NOT_OK(GetReader(*child, ctx, &child_reader)); - // Use the max definition/repetition level of the leaf here - out->reset(new NestedListReader(ctx, list_field, child->level_info.def_level, - child->level_info.rep_level, - std::move(child_reader))); + out->reset(new ListReader(ctx, list_field, field.level_info, + std::move(child_reader))); } else if (type_id == ::arrow::Type::STRUCT) { std::vector> child_fields; std::vector> child_readers; for (const auto& child : field.children) { - if (child.is_leaf() && !ctx->IncludesLeaf(child.column_index)) { - // Excluded leaf - continue; - } std::unique_ptr child_reader; RETURN_NOT_OK(GetReader(child, ctx, &child_reader)); if (!child_reader) { @@ -742,7 +754,8 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& auto filtered_field = ::arrow::field(field.field->name(), ::arrow::struct_(child_fields), field.field->nullable(), field.field->metadata()); - out->reset(new StructReader(ctx, field, filtered_field, std::move(child_readers))); + out->reset(new StructReader(ctx, filtered_field, field.level_info, + std::move(child_readers))); } else { return Status::Invalid("Unsupported nested type: ", field.field->ToString()); } diff --git a/cpp/src/parquet/arrow/reader_internal.cc b/cpp/src/parquet/arrow/reader_internal.cc index 903cbabaae2..7ab401bd9cf 100644 --- a/cpp/src/parquet/arrow/reader_internal.cc +++ b/cpp/src/parquet/arrow/reader_internal.cc @@ -795,118 +795,5 @@ Status TransferColumnData(RecordReader* reader, std::shared_ptr value_ return Status::OK(); } -Status ReconstructNestedList(const std::shared_ptr& arr, - std::shared_ptr field, int16_t max_def_level, - int16_t max_rep_level, const int16_t* def_levels, - const int16_t* rep_levels, int64_t total_levels, - ::arrow::MemoryPool* pool, std::shared_ptr* out) { - // Walk downwards to extract nullability - std::vector item_names; - std::vector nullable; - std::vector> field_metadata; - std::vector> offset_builders; - std::vector> valid_bits_builders; - nullable.push_back(field->nullable()); - while (field->type()->num_fields() > 0) { - if (field->type()->num_fields() > 1) { - return Status::NotImplemented("Fields with more than one child are not supported."); - } else { - if (field->type()->id() != ::arrow::Type::LIST) { - return Status::NotImplemented("Currently only nesting with Lists is supported."); - } - field = field->type()->field(0); - } - item_names.push_back(field->name()); - offset_builders.emplace_back( - std::make_shared<::arrow::Int32Builder>(::arrow::int32(), pool)); - valid_bits_builders.emplace_back( - std::make_shared<::arrow::BooleanBuilder>(::arrow::boolean(), pool)); - nullable.push_back(field->nullable()); - field_metadata.push_back(field->metadata()); - } - - int64_t list_depth = offset_builders.size(); - // This describes the minimal definition that describes a level that - // reflects a value in the primitive values array. - int16_t values_def_level = max_def_level; - if (nullable[nullable.size() - 1]) { - values_def_level--; - } - - // The definition levels that are needed so that a list is declared - // as empty and not null. - std::vector empty_def_level(list_depth); - int def_level = 0; - for (int i = 0; i < list_depth; i++) { - if (nullable[i]) { - def_level++; - } - empty_def_level[i] = static_cast(def_level); - def_level++; - } - - int32_t values_offset = 0; - std::vector null_counts(list_depth, 0); - for (int64_t i = 0; i < total_levels; i++) { - int16_t rep_level = rep_levels[i]; - if (rep_level < max_rep_level) { - for (int64_t j = rep_level; j < list_depth; j++) { - if (j == (list_depth - 1)) { - RETURN_NOT_OK(offset_builders[j]->Append(values_offset)); - } else { - RETURN_NOT_OK(offset_builders[j]->Append( - static_cast(offset_builders[j + 1]->length()))); - } - - if (((empty_def_level[j] - 1) == def_levels[i]) && (nullable[j])) { - RETURN_NOT_OK(valid_bits_builders[j]->Append(false)); - null_counts[j]++; - break; - } else { - RETURN_NOT_OK(valid_bits_builders[j]->Append(true)); - if (empty_def_level[j] == def_levels[i]) { - break; - } - } - } - } - if (def_levels[i] >= values_def_level) { - values_offset++; - } - } - // Add the final offset to all lists - for (int64_t j = 0; j < list_depth; j++) { - if (j == (list_depth - 1)) { - RETURN_NOT_OK(offset_builders[j]->Append(values_offset)); - } else { - RETURN_NOT_OK(offset_builders[j]->Append( - static_cast(offset_builders[j + 1]->length()))); - } - } - - std::vector> offsets; - std::vector> valid_bits; - std::vector list_lengths; - for (int64_t j = 0; j < list_depth; j++) { - list_lengths.push_back(offset_builders[j]->length() - 1); - std::shared_ptr array; - RETURN_NOT_OK(offset_builders[j]->Finish(&array)); - offsets.emplace_back(std::static_pointer_cast(array)->values()); - RETURN_NOT_OK(valid_bits_builders[j]->Finish(&array)); - valid_bits.emplace_back(std::static_pointer_cast(array)->values()); - } - - *out = arr; - - // TODO(wesm): Use passed-in field - for (int64_t j = list_depth - 1; j >= 0; j--) { - auto list_type = ::arrow::list(::arrow::field(item_names[j], (*out)->type(), - nullable[j + 1], field_metadata[j])); - *out = std::make_shared<::arrow::ListArray>(list_type, list_lengths[j], offsets[j], - *out, valid_bits[j], null_counts[j]); - } - return Status::OK(); -} - } // namespace arrow } // namespace parquet diff --git a/cpp/src/parquet/arrow/reader_internal.h b/cpp/src/parquet/arrow/reader_internal.h index 62eb166da01..8dd439f4fbb 100644 --- a/cpp/src/parquet/arrow/reader_internal.h +++ b/cpp/src/parquet/arrow/reader_internal.h @@ -104,13 +104,6 @@ Status TransferColumnData(::parquet::internal::RecordReader* reader, const ColumnDescriptor* descr, ::arrow::MemoryPool* pool, std::shared_ptr<::arrow::ChunkedArray>* out); -Status ReconstructNestedList(const std::shared_ptr<::arrow::Array>& arr, - std::shared_ptr<::arrow::Field> field, int16_t max_def_level, - int16_t max_rep_level, const int16_t* def_levels, - const int16_t* rep_levels, int64_t total_levels, - ::arrow::MemoryPool* pool, - std::shared_ptr<::arrow::Array>* out); - struct ReaderContext { ParquetFileReader* reader; ::arrow::MemoryPool* pool; diff --git a/cpp/src/parquet/arrow/reconstruct_internal_test.cc b/cpp/src/parquet/arrow/reconstruct_internal_test.cc index 6faa5a5c70a..f3850244ceb 100644 --- a/cpp/src/parquet/arrow/reconstruct_internal_test.cc +++ b/cpp/src/parquet/arrow/reconstruct_internal_test.cc @@ -40,15 +40,6 @@ #include "parquet/file_writer.h" #include "parquet/properties.h" -// Set to 1 to see failures in failing tests -#define RUN_FAILING_TESTS 0 - -#if RUN_FAILING_TESTS -#define FAILING(test_name) test_name -#else -#define FAILING(test_name) DISABLED_##test_name -#endif - using arrow::Array; using arrow::ArrayFromJSON; using arrow::AssertArraysEqual; @@ -316,7 +307,7 @@ TEST_F(TestReconstructColumn, PrimitiveRequired) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(PrimitiveRepeated)) { +TEST_F(TestReconstructColumn, PrimitiveRepeated) { // Arrow schema: list(int32 not null) not null this->SetParquetSchema( PrimitiveNode::Make("node_name", Repetition::REPEATED, ParquetType::INT32)); @@ -349,7 +340,7 @@ TEST_F(TestReconstructColumn, NestedRequiredRequired) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedOptionalRequired)) { +TEST_F(TestReconstructColumn, NestedOptionalRequired) { // Arrow schema: struct(a: int32 not null) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -419,7 +410,7 @@ TEST_F(TestReconstructColumn, NestedRequiredRequiredRequired) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedRequiredOptionalRequired)) { +TEST_F(TestReconstructColumn, NestedRequiredOptionalRequired) { // Arrow schema: struct(a: struct(b: int32 not null)) not null SetParquetSchema(GroupNode::Make( "parent", Repetition::REQUIRED, @@ -440,7 +431,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedRequiredOptionalRequired)) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedOptionalRequiredOptional)) { +TEST_F(TestReconstructColumn, NestedOptionalRequiredOptional) { // Arrow schema: struct(a: struct(b: int32) not null) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -526,7 +517,7 @@ TEST_F(TestReconstructColumn, NestedTwoFields2) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(NestedTwoFields3)) { +TEST_F(TestReconstructColumn, NestedTwoFields3) { // Arrow schema: struct(a: int32 not null, b: int64 not null) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -847,7 +838,7 @@ TEST_F(TestReconstructColumn, ThreeLevelListOptionalOptional) { // Legacy list encodings // -TEST_F(TestReconstructColumn, FAILING(TwoLevelListRequired)) { +TEST_F(TestReconstructColumn, TwoLevelListRequired) { // Arrow schema: list(int32 not null) not null SetParquetSchema(GroupNode::Make( "parent", Repetition::REQUIRED, @@ -860,11 +851,11 @@ TEST_F(TestReconstructColumn, FAILING(TwoLevelListRequired)) { // TODO should field name "element" (Parquet convention for List nodes) // be changed to "item" (Arrow convention for List types)? - auto expected = ArrayFromJSON(List(int32()), "[[], [4, 5], [6]]"); + auto expected = ArrayFromJSON(List(int32(), /*nullable=*/false), "[[], [4, 5], [6]]"); AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(TwoLevelListOptional)) { +TEST_F(TestReconstructColumn, TwoLevelListOptional) { // Arrow schema: list(int32 not null) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -884,12 +875,12 @@ TEST_F(TestReconstructColumn, FAILING(TwoLevelListOptional)) { // List-in-struct // -TEST_F(TestReconstructColumn, FAILING(NestedList1)) { +TEST_F(TestReconstructColumn, NestedList1) { // Arrow schema: struct(a: list(int32 not null) not null) not null SetParquetSchema(GroupNode::Make( - "a", Repetition::REQUIRED, + "a", Repetition::REQUIRED, // this {GroupNode::Make( - "parent", Repetition::REQUIRED, + "p", Repetition::REQUIRED, {GroupNode::Make("list", Repetition::REPEATED, {PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32)})}, @@ -899,20 +890,20 @@ TEST_F(TestReconstructColumn, FAILING(NestedList1)) { LevelVector rep_levels = {0, 0, 1, 0}; std::vector values = {4, 5, 6}; - auto type = OneFieldStruct("a", List(int32(), /*nullable=*/false), + auto type = OneFieldStruct("p", List(int32(), /*nullable=*/false), /*nullable=*/false); - auto expected = ArrayFromJSON(type, R"([{"a": []}, - {"a": [4, 5]}, - {"a": [6]}])"); + auto expected = ArrayFromJSON(type, R"([{"p": []}, + {"p": [4, 5]}, + {"p": [6]}])"); AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedList2)) { +TEST_F(TestReconstructColumn, NestedList2) { // Arrow schema: struct(a: list(int32 not null) not null) SetParquetSchema(GroupNode::Make( "a", Repetition::OPTIONAL, {GroupNode::Make( - "parent", Repetition::REQUIRED, + "p", Repetition::REQUIRED, {GroupNode::Make("list", Repetition::REPEATED, {PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32)})}, @@ -922,21 +913,21 @@ TEST_F(TestReconstructColumn, FAILING(NestedList2)) { LevelVector rep_levels = {0, 0, 0, 1, 0}; std::vector values = {4, 5, 6}; - auto type = OneFieldStruct("a", List(int32(), /*nullable=*/false), + auto type = OneFieldStruct("p", List(int32(), /*nullable=*/false), /*nullable=*/false); auto expected = ArrayFromJSON(type, R"([null, - {"a": []}, - {"a": [4, 5]}, - {"a": [6]}])"); + {"p": []}, + {"p": [4, 5]}, + {"p": [6]}])"); AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedList3)) { +TEST_F(TestReconstructColumn, NestedList3) { // Arrow schema: struct(a: list(int32 not null)) not null SetParquetSchema(GroupNode::Make( - "a", Repetition::REQUIRED, + "a", Repetition::REQUIRED, // column name (column a is a struct of) {GroupNode::Make( - "parent", Repetition::OPTIONAL, + "p", Repetition::OPTIONAL, // name in struct {GroupNode::Make("list", Repetition::REPEATED, {PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32)})}, @@ -946,20 +937,20 @@ TEST_F(TestReconstructColumn, FAILING(NestedList3)) { LevelVector rep_levels = {0, 0, 0, 1, 0}; std::vector values = {4, 5, 6}; - auto type = OneFieldStruct("a", List(int32())); - auto expected = ArrayFromJSON(type, R"([{"a": null}, - {"a": []}, - {"a": [4, 5]}, - {"a": [6]}])"); + auto type = OneFieldStruct("p", List(int32(), /*nullable=*/false)); + auto expected = ArrayFromJSON(type, R"([{"p": null}, + {"p": []}, + {"p": [4, 5]}, + {"p": [6]}])"); AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedList4)) { +TEST_F(TestReconstructColumn, NestedList4) { // Arrow schema: struct(a: list(int32 not null)) SetParquetSchema(GroupNode::Make( "a", Repetition::OPTIONAL, {GroupNode::Make( - "parent", Repetition::OPTIONAL, + "p", Repetition::OPTIONAL, {GroupNode::Make("list", Repetition::REPEATED, {PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32)})}, @@ -969,21 +960,21 @@ TEST_F(TestReconstructColumn, FAILING(NestedList4)) { LevelVector rep_levels = {0, 0, 0, 0, 1, 0}; std::vector values = {4, 5, 6}; - auto type = OneFieldStruct("a", List(int32())); + auto type = OneFieldStruct("p", List(int32(), /*nullable=*/false)); auto expected = ArrayFromJSON(type, R"([null, - {"a": null}, - {"a": []}, - {"a": [4, 5]}, - {"a": [6]}])"); + {"p": null}, + {"p": []}, + {"p": [4, 5]}, + {"p": [6]}])"); AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedList5)) { +TEST_F(TestReconstructColumn, NestedList5) { // Arrow schema: struct(a: list(int32) not null) SetParquetSchema(GroupNode::Make( "a", Repetition::OPTIONAL, {GroupNode::Make( - "parent", Repetition::REQUIRED, + "p", Repetition::REQUIRED, {GroupNode::Make("list", Repetition::REPEATED, {PrimitiveNode::Make("element", Repetition::OPTIONAL, ParquetType::INT32)})}, @@ -993,20 +984,20 @@ TEST_F(TestReconstructColumn, FAILING(NestedList5)) { LevelVector rep_levels = {0, 0, 0, 1, 0, 1}; std::vector values = {4, 5, 6}; - auto type = OneFieldStruct("a", List(int32()), /*nullable=*/false); + auto type = OneFieldStruct("p", List(int32()), /*nullable=*/false); auto expected = ArrayFromJSON(type, R"([null, - {"a": []}, - {"a": [4, null]}, - {"a": [5, 6]}])"); + {"p": []}, + {"p": [4, null]}, + {"p": [5, 6]}])"); AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(NestedList6)) { +TEST_F(TestReconstructColumn, NestedList6) { // Arrow schema: struct(a: list(int32)) SetParquetSchema(GroupNode::Make( "a", Repetition::OPTIONAL, {GroupNode::Make( - "parent", Repetition::OPTIONAL, + "p", Repetition::OPTIONAL, {GroupNode::Make("list", Repetition::REPEATED, {PrimitiveNode::Make("element", Repetition::OPTIONAL, ParquetType::INT32)})}, @@ -1016,12 +1007,12 @@ TEST_F(TestReconstructColumn, FAILING(NestedList6)) { LevelVector rep_levels = {0, 0, 0, 0, 1, 0, 1}; std::vector values = {4, 5, 6}; - auto type = OneFieldStruct("a", List(int32())); + auto type = OneFieldStruct("p", List(int32())); auto expected = ArrayFromJSON(type, R"([null, - {"a": null}, - {"a": []}, - {"a": [4, null]}, - {"a": [5, 6]}])"); + {"p": null}, + {"p": []}, + {"p": [4, null]}, + {"p": [5, 6]}])"); AssertReconstruct(*expected, def_levels, rep_levels, values); } @@ -1029,7 +1020,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedList6)) { // Struct-in-list // -TEST_F(TestReconstructColumn, FAILING(ListNested1)) { +TEST_F(TestReconstructColumn, ListNested1) { // Arrow schema: list(struct(a: int32 not null) not null) not null SetParquetSchema(GroupNode::Make( "parent", Repetition::REQUIRED, @@ -1052,7 +1043,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNested1)) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(ListNested2)) { +TEST_F(TestReconstructColumn, ListNested2) { // Arrow schema: list(struct(a: int32 not null) not null) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -1076,7 +1067,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNested2)) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(ListNested3)) { +TEST_F(TestReconstructColumn, ListNested3) { // Arrow schema: list(struct(a: int32 not null)) not null SetParquetSchema(GroupNode::Make( "parent", Repetition::REQUIRED, @@ -1098,7 +1089,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNested3)) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(ListNested4)) { +TEST_F(TestReconstructColumn, ListNested4) { // Arrow schema: list(struct(a: int32 not null)) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -1121,7 +1112,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNested4)) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(ListNested5)) { +TEST_F(TestReconstructColumn, ListNested5) { // Arrow schema: list(struct(a: int32) not null) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -1145,7 +1136,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNested5)) { AssertReconstruct(*expected, def_levels, rep_levels, values); } -TEST_F(TestReconstructColumn, FAILING(ListNested6)) { +TEST_F(TestReconstructColumn, ListNested6) { // Arrow schema: list(struct(a: int32)) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -1172,7 +1163,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNested6)) { // Struct (two fields)-in-list // -TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields1)) { +TEST_F(TestReconstructColumn, ListNestedTwoFields1) { // Arrow schema: list(struct(a: int32 not null, // b: int64 not null) not null) not null SetParquetSchema(GroupNode::Make( @@ -1202,7 +1193,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields1)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields2)) { +TEST_F(TestReconstructColumn, ListNestedTwoFields2) { // Arrow schema: list(struct(a: int32, // b: int64 not null) not null) not null SetParquetSchema(GroupNode::Make( @@ -1232,7 +1223,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields2)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields3)) { +TEST_F(TestReconstructColumn, ListNestedTwoFields3) { // Arrow schema: list(struct(a: int32 not null, // b: int64 not null)) not null SetParquetSchema(GroupNode::Make( @@ -1261,7 +1252,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields3)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields4)) { +TEST_F(TestReconstructColumn, ListNestedTwoFields4) { // Arrow schema: list(struct(a: int32, // b: int64 not null) not null) SetParquetSchema(GroupNode::Make( @@ -1292,7 +1283,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields4)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields5)) { +TEST_F(TestReconstructColumn, ListNestedTwoFields5) { // Arrow schema: list(struct(a: int32, // b: int64 not null)) SetParquetSchema(GroupNode::Make( @@ -1313,8 +1304,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields5)) { Int64Vector{7, 8})); auto type = - List(struct_({field("a", int32()), field("b", int64(), /*nullable=*/false)}), - /*nullable=*/false); + List(struct_({field("a", int32()), field("b", int64(), /*nullable=*/false)})); auto expected = ArrayFromJSON(type, R"([null, [], @@ -1323,7 +1313,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields5)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields6)) { +TEST_F(TestReconstructColumn, ListNestedTwoFields6) { // Arrow schema: list(struct(a: int32, // b: int64)) SetParquetSchema(GroupNode::Make( @@ -1343,9 +1333,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields6)) { ASSERT_OK(WriteInt64Column(DefLevels{0, 1, 3, 2, 4}, RepLevels{0, 0, 0, 1, 0}, Int64Vector{7})); - auto type = - List(struct_({field("a", int32()), field("b", int64(), /*nullable=*/false)}), - /*nullable=*/false); + auto type = List(struct_({field("a", int32()), field("b", int64())})); auto expected = ArrayFromJSON(type, R"([null, [], @@ -1358,7 +1346,7 @@ TEST_F(TestReconstructColumn, FAILING(ListNestedTwoFields6)) { // List-in-struct (two fields) // -TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList1)) { +TEST_F(TestReconstructColumn, NestedTwoFieldsList1) { // Arrow schema: struct(a: int64 not null, // b: list(int32 not null) not null // ) not null @@ -1388,7 +1376,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList1)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList2)) { +TEST_F(TestReconstructColumn, NestedTwoFieldsList2) { // Arrow schema: struct(a: int64 not null, // b: list(int32 not null) // ) not null @@ -1418,7 +1406,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList2)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList3)) { +TEST_F(TestReconstructColumn, NestedTwoFieldsList3) { // Arrow schema: struct(a: int64, // b: list(int32 not null) // ) not null @@ -1448,7 +1436,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList3)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList4)) { +TEST_F(TestReconstructColumn, NestedTwoFieldsList4) { // Arrow schema: struct(a: int64, // b: list(int32 not null) // ) @@ -1480,7 +1468,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList4)) { CheckColumn(/*column_index=*/0, *expected); } -TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList5)) { +TEST_F(TestReconstructColumn, NestedTwoFieldsList5) { // Arrow schema: struct(a: int64, b: list(int32)) SetParquetSchema(GroupNode::Make( "parent", Repetition::OPTIONAL, @@ -1499,8 +1487,7 @@ TEST_F(TestReconstructColumn, FAILING(NestedTwoFieldsList5)) { ASSERT_OK(WriteInt32Column(DefLevels{0, 1, 2, 4, 3, 4}, RepLevels{0, 0, 0, 0, 1, 0}, Int32Vector{7, 8})); - auto type = - struct_({field("a", int64()), field("b", List(int32(), /*nullable=*/false))}); + auto type = struct_({field("a", int64()), field("b", List(int32()))}); auto expected = ArrayFromJSON(type, R"([null, {"a": 4, "b": null}, diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index a25704a7115..8b2faddf2b4 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -447,7 +447,6 @@ Status GroupToStruct(const GroupNode& node, LevelInfo current_levels, SchemaField* out) { std::vector> arrow_fields; out->children.resize(node.field_count()); - // All level increments for the node are expected to happen by callers. // This is required because repeated elements need to have there own // SchemaField. diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 672b6e3708c..6b5450312b3 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -42,6 +42,7 @@ #include "parquet/encoding.h" #include "parquet/encryption_internal.h" #include "parquet/internal_file_decryptor.h" +#include "parquet/level_comparison.h" #include "parquet/level_conversion.h" #include "parquet/properties.h" #include "parquet/statistics.h" @@ -55,6 +56,25 @@ using arrow::internal::checked_cast; using arrow::internal::MultiplyWithOverflow; namespace parquet { +namespace { +inline bool HasSpacedValues(const ColumnDescriptor* descr) { + if (descr->max_repetition_level() > 0) { + // repeated+flat case + return !descr->schema_node()->is_required(); + } else { + // non-repeated+nested case + // Find if a node forces nulls in the lowest level along the hierarchy + const schema::Node* node = descr->schema_node().get(); + while (node) { + if (node->is_optional()) { + return true; + } + node = node->parent(); + } + return false; + } +} +} // namespace LevelDecoder::LevelDecoder() : num_values_remaining_(0) {} @@ -63,6 +83,7 @@ LevelDecoder::~LevelDecoder() {} int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, int num_buffered_values, const uint8_t* data, int32_t data_size) { + max_level_ = max_level; int32_t num_bytes = 0; encoding_ = encoding; num_values_remaining_ = num_buffered_values; @@ -110,6 +131,7 @@ int LevelDecoder::SetData(Encoding::type encoding, int16_t max_level, void LevelDecoder::SetDataV2(int32_t num_bytes, int16_t max_level, int num_buffered_values, const uint8_t* data) { + max_level_ = max_level; // Repetition and definition levels always uses RLE encoding // in the DataPageV2 format. if (num_bytes < 0) { @@ -135,6 +157,15 @@ int LevelDecoder::Decode(int batch_size, int16_t* levels) { } else { num_decoded = bit_packed_decoder_->GetBatch(bit_width_, levels, num_values); } + if (num_decoded > 0) { + internal::MinMax min_max = internal::FindMinMax(levels, num_decoded); + if (ARROW_PREDICT_FALSE(min_max.min < 0 || min_max.max > max_level_)) { + std::stringstream ss; + ss << "Malformed levels. min: " << min_max.min << " max: " << min_max.max + << " out of range. Max Level: " << max_level_; + throw ParquetException(ss.str()); + } + } num_values_remaining_ -= num_decoded; return num_decoded; } @@ -880,8 +911,7 @@ int64_t TypedColumnReaderImpl::ReadBatchSpaced( } } - const bool has_spaced_values = internal::HasSpacedValues(this->descr_); - + const bool has_spaced_values = HasSpacedValues(this->descr_); int64_t null_count = 0; if (!has_spaced_values) { int values_to_read = 0; @@ -896,9 +926,12 @@ int64_t TypedColumnReaderImpl::ReadBatchSpaced( /*bits_are_set=*/true); *values_read = total_values; } else { - internal::DefinitionLevelsToBitmap(def_levels, num_def_levels, this->max_def_level_, - this->max_rep_level_, values_read, &null_count, - valid_bits, valid_bits_offset); + internal::LevelInfo info; + info.repeated_ancestor_def_level = this->max_def_level_ - 1; + info.def_level = this->max_def_level_; + info.rep_level = this->max_rep_level_; + internal::DefinitionLevelsToBitmap(def_levels, num_def_levels, info, values_read, + &null_count, valid_bits, valid_bits_offset); total_values = this->ReadValuesSpaced(*values_read, values, static_cast(null_count), valid_bits, valid_bits_offset); @@ -1008,8 +1041,10 @@ class TypedRecordReader : public ColumnReaderImplBase, public: using T = typename DType::c_type; using BASE = ColumnReaderImplBase; - TypedRecordReader(const ColumnDescriptor* descr, MemoryPool* pool) : BASE(descr, pool) { - nullable_values_ = internal::HasSpacedValues(descr); + TypedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, MemoryPool* pool) + : BASE(descr, pool) { + leaf_info_ = leaf_info; + nullable_values_ = leaf_info.HasNullableValues(); at_record_start_ = true; records_read_ = 0; values_written_ = 0; @@ -1128,7 +1163,7 @@ class TypedRecordReader : public ColumnReaderImplBase, } std::shared_ptr ReleaseIsValid() override { - if (nullable_values_) { + if (leaf_info_.HasNullableValues()) { auto result = valid_bits_; PARQUET_THROW_NOT_OK(result->Resize(BitUtil::BytesForBits(values_written_), true)); valid_bits_ = AllocateBuffer(this->pool_); @@ -1170,13 +1205,7 @@ class TypedRecordReader : public ColumnReaderImplBase, break; } } - } else if (ARROW_PREDICT_FALSE(rep_level > this->max_rep_level_)) { - std::stringstream ss; - ss << "Malformed repetition levels, " << rep_level << " exceeded maximum " - << this->max_rep_level_ << " indicated by schema"; - throw ParquetException(ss.str()); } - // We have decided to consume the level at this position; therefore we // must advance until we find another record boundary at_record_start_ = false; @@ -1184,11 +1213,6 @@ class TypedRecordReader : public ColumnReaderImplBase, const int16_t def_level = *def_levels++; if (def_level == this->max_def_level_) { ++values_to_read; - } else if (ARROW_PREDICT_FALSE(def_level > this->max_def_level_)) { - std::stringstream ss; - ss << "Malformed definition levels, " << def_level << " exceeded maximum " - << this->max_def_level_ << " indicated by schema"; - throw ParquetException(ss.str()); } ++levels_position_; } @@ -1249,7 +1273,7 @@ class TypedRecordReader : public ColumnReaderImplBase, } values_capacity_ = new_values_capacity; } - if (nullable_values_) { + if (leaf_info_.HasNullableValues()) { int64_t valid_bytes_new = BitUtil::BytesForBits(values_capacity_); if (valid_bits_->size() < valid_bytes_new) { int64_t valid_bytes_old = BitUtil::BytesForBits(values_written_); @@ -1344,12 +1368,12 @@ class TypedRecordReader : public ColumnReaderImplBase, } int64_t null_count = 0; - if (nullable_values_) { + if (leaf_info_.HasNullableValues()) { int64_t values_with_nulls = 0; - DefinitionLevelsToBitmap( - def_levels() + start_levels_position, levels_position_ - start_levels_position, - this->max_def_level_, this->max_rep_level_, &values_with_nulls, &null_count, - valid_bits_->mutable_data(), values_written_); + DefinitionLevelsToBitmap(def_levels() + start_levels_position, + levels_position_ - start_levels_position, leaf_info_, + &values_with_nulls, &null_count, + valid_bits_->mutable_data(), values_written_); values_to_read = values_with_nulls - null_count; DCHECK_GE(values_to_read, 0); ReadValuesSpaced(values_with_nulls, null_count); @@ -1357,7 +1381,7 @@ class TypedRecordReader : public ColumnReaderImplBase, DCHECK_GE(values_to_read, 0); ReadValuesDense(values_to_read); } - if (this->max_def_level_ > 0) { + if (this->leaf_info_.def_level > 0) { // Optional, repeated, or some mix thereof this->ConsumeBufferedValues(levels_position_ - start_levels_position); } else { @@ -1415,13 +1439,15 @@ class TypedRecordReader : public ColumnReaderImplBase, T* ValuesHead() { return reinterpret_cast(values_->mutable_data()) + values_written_; } + LevelInfo leaf_info_; }; class FLBARecordReader : public TypedRecordReader, virtual public BinaryRecordReader { public: - FLBARecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : TypedRecordReader(descr, pool), builder_(nullptr) { + FLBARecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, + ::arrow::MemoryPool* pool) + : TypedRecordReader(descr, leaf_info, pool), builder_(nullptr) { DCHECK_EQ(descr_->physical_type(), Type::FIXED_LEN_BYTE_ARRAY); int byte_width = descr_->type_length(); std::shared_ptr<::arrow::DataType> type = ::arrow::fixed_size_binary(byte_width); @@ -1473,8 +1499,9 @@ class FLBARecordReader : public TypedRecordReader, class ByteArrayChunkedRecordReader : public TypedRecordReader, virtual public BinaryRecordReader { public: - ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, ::arrow::MemoryPool* pool) - : TypedRecordReader(descr, pool) { + ByteArrayChunkedRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, + ::arrow::MemoryPool* pool) + : TypedRecordReader(descr, leaf_info, pool) { DCHECK_EQ(descr_->physical_type(), Type::BYTE_ARRAY); accumulator_.builder.reset(new ::arrow::BinaryBuilder(pool)); } @@ -1513,9 +1540,9 @@ class ByteArrayChunkedRecordReader : public TypedRecordReader, class ByteArrayDictionaryRecordReader : public TypedRecordReader, virtual public DictionaryRecordReader { public: - ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, + ByteArrayDictionaryRecordReader(const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool) - : TypedRecordReader(descr, pool), builder_(pool) { + : TypedRecordReader(descr, leaf_info, pool), builder_(pool) { this->read_dictionary_ = true; } @@ -1602,35 +1629,36 @@ template <> void TypedRecordReader::DebugPrintState() {} std::shared_ptr MakeByteArrayRecordReader(const ColumnDescriptor* descr, + LevelInfo leaf_info, ::arrow::MemoryPool* pool, bool read_dictionary) { if (read_dictionary) { - return std::make_shared(descr, pool); + return std::make_shared(descr, leaf_info, pool); } else { - return std::make_shared(descr, pool); + return std::make_shared(descr, leaf_info, pool); } } std::shared_ptr RecordReader::Make(const ColumnDescriptor* descr, - MemoryPool* pool, + LevelInfo leaf_info, MemoryPool* pool, const bool read_dictionary) { switch (descr->physical_type()) { case Type::BOOLEAN: - return std::make_shared>(descr, pool); + return std::make_shared>(descr, leaf_info, pool); case Type::INT32: - return std::make_shared>(descr, pool); + return std::make_shared>(descr, leaf_info, pool); case Type::INT64: - return std::make_shared>(descr, pool); + return std::make_shared>(descr, leaf_info, pool); case Type::INT96: - return std::make_shared>(descr, pool); + return std::make_shared>(descr, leaf_info, pool); case Type::FLOAT: - return std::make_shared>(descr, pool); + return std::make_shared>(descr, leaf_info, pool); case Type::DOUBLE: - return std::make_shared>(descr, pool); + return std::make_shared>(descr, leaf_info, pool); case Type::BYTE_ARRAY: - return MakeByteArrayRecordReader(descr, pool, read_dictionary); + return MakeByteArrayRecordReader(descr, leaf_info, pool, read_dictionary); case Type::FIXED_LEN_BYTE_ARRAY: - return std::make_shared(descr, pool); + return std::make_shared(descr, leaf_info, pool); default: { // PARQUET-1481: This can occur if the file is corrupt std::stringstream ss; diff --git a/cpp/src/parquet/column_reader.h b/cpp/src/parquet/column_reader.h index 7b5ee1b722a..60c44ffa6d2 100644 --- a/cpp/src/parquet/column_reader.h +++ b/cpp/src/parquet/column_reader.h @@ -23,6 +23,7 @@ #include #include "parquet/exception.h" +#include "parquet/level_conversion.h" #include "parquet/platform.h" #include "parquet/schema.h" #include "parquet/types.h" @@ -75,6 +76,7 @@ class PARQUET_EXPORT LevelDecoder { Encoding::type encoding_; std::unique_ptr<::arrow::util::RleDecoder> rle_decoder_; std::unique_ptr<::arrow::BitUtil::BitReader> bit_packed_decoder_; + int16_t max_level_; }; struct CryptoContext { @@ -208,7 +210,7 @@ namespace internal { class RecordReader { public: static std::shared_ptr Make( - const ColumnDescriptor* descr, + const ColumnDescriptor* descr, LevelInfo leaf_info, ::arrow::MemoryPool* pool = ::arrow::default_memory_pool(), const bool read_dictionary = false); @@ -314,25 +316,6 @@ class DictionaryRecordReader : virtual public RecordReader { virtual std::shared_ptr<::arrow::ChunkedArray> GetResult() = 0; }; -// TODO(itaiin): another code path split to merge when the general case is done -static inline bool HasSpacedValues(const ColumnDescriptor* descr) { - if (descr->max_repetition_level() > 0) { - // repeated+flat case - return !descr->schema_node()->is_required(); - } else { - // non-repeated+nested case - // Find if a node forces nulls in the lowest level along the hierarchy - const schema::Node* node = descr->schema_node().get(); - while (node) { - if (node->is_optional()) { - return true; - } - node = node->parent(); - } - return false; - } -} - } // namespace internal using BoolReader = TypedColumnReader; diff --git a/cpp/src/parquet/level_comparison.cc b/cpp/src/parquet/level_comparison.cc new file mode 100644 index 00000000000..2b68fef42c1 --- /dev/null +++ b/cpp/src/parquet/level_comparison.cc @@ -0,0 +1,71 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#define IMPL_NAMESPACE standard +#include "parquet/level_comparison.h" + +#include + +#include "arrow/util/dispatch.h" + +namespace parquet { +namespace internal { +namespace { + +using ::arrow::internal::DispatchLevel; +using ::arrow::internal::DynamicDispatch; + +struct GreaterThanDynamicFunction { + using FunctionType = decltype(&GreaterThanBitmap); + + static std::vector> implementations() { + return { + { DispatchLevel::NONE, standard::GreaterThanBitmapImpl } +#if defined(ARROW_HAVE_RUNTIME_AVX2) + , { DispatchLevel::AVX2, GreaterThanBitmapAvx2 } +#endif + }; + } +}; + +struct MinMaxDynamicFunction { + using FunctionType = decltype(&FindMinMax); + + static std::vector> implementations() { + return { + { DispatchLevel::NONE, standard::FindMinMaxImpl } +#if defined(ARROW_HAVE_RUNTIME_AVX2) + , { DispatchLevel::AVX2, FindMinMaxAvx2 } +#endif + }; + } +}; + +} // namespace + +uint64_t GreaterThanBitmap(const int16_t* levels, int64_t num_levels, int16_t rhs) { + static DynamicDispatch dispatch; + return dispatch.func(levels, num_levels, rhs); +} + +MinMax FindMinMax(const int16_t* levels, int64_t num_levels) { + static DynamicDispatch dispatch; + return dispatch.func(levels, num_levels); +} + +} // namespace internal +} // namespace parquet diff --git a/cpp/src/parquet/level_comparison.h b/cpp/src/parquet/level_comparison.h new file mode 100644 index 00000000000..92e5f0d3cb3 --- /dev/null +++ b/cpp/src/parquet/level_comparison.h @@ -0,0 +1,91 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +#pragma once + +#include +#include + +#include "arrow/util/bit_util.h" + +namespace parquet { +namespace internal { + +// These APIs are likely to be revised as part of ARROW-8494 to reduce duplicate code. +// They currently represent minimal functionality for vectorized computation of definition +// levels. + +/// Builds a bitmap by applying predicate to the level vector provided. +/// +/// \param[in] levels Rep or def level array. +/// \param[in] num_levels The number of levels to process (must be [0, 64]) +/// \param[in] predicate The predicate to apply (must have the signature `bool +/// predicate(int16_t)`. +/// \returns The bitmap using least significant "bit" ordering. +/// +/// N.B. Correct byte ordering is dependent on little-endian architectures. +/// +template +inline uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, + Predicate predicate) { + // Both clang and GCC can vectorize this automatically with SSE4/AVX2. + uint64_t mask = 0; + for (int x = 0; x < num_levels; x++) { + mask |= static_cast(predicate(levels[x]) ? 1 : 0) << x; + } + return ::arrow::BitUtil::ToLittleEndian(mask); +} + +/// Builds a bitmap where each set bit indicates the corresponding level is greater +/// than rhs. +uint64_t GreaterThanBitmap(const int16_t* levels, int64_t num_levels, int16_t rhs); + +#if defined(ARROW_HAVE_RUNTIME_AVX2) +uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_t rhs); +#endif + +struct MinMax { + int16_t min; + int16_t max; +}; + +MinMax FindMinMax(const int16_t* levels, int64_t num_levels); + +#if defined(ARROW_HAVE_RUNTIME_AVX2) +MinMax FindMinMaxAvx2(const int16_t* levels, int64_t num_levels); +#endif + +// Used to make sure ODR rule isn't violated. +namespace IMPL_NAMESPACE { +inline MinMax FindMinMaxImpl(const int16_t* levels, int64_t num_levels) { + MinMax out{std::numeric_limits::max(), std::numeric_limits::min()}; + for (int x = 0; x < num_levels; x++) { + out.min = std::min(levels[x], out.min); + out.max = std::max(levels[x], out.max); + } + return out; +} + +inline uint64_t GreaterThanBitmapImpl(const int16_t* levels, int64_t num_levels, + int16_t rhs) { + return LevelsToBitmap(levels, num_levels, [rhs](int16_t value) { return value > rhs; }); +} + +} // namespace IMPL_NAMESPACE + +} // namespace internal + +} // namespace parquet diff --git a/cpp/src/parquet/level_comparison_avx2.cc b/cpp/src/parquet/level_comparison_avx2.cc new file mode 100644 index 00000000000..6cd6d0cc1f7 --- /dev/null +++ b/cpp/src/parquet/level_comparison_avx2.cc @@ -0,0 +1,33 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#define IMPL_NAMESPACE avx2 +#include "parquet/level_comparison.h" +#undef IMPL_NAMESPACE + +namespace parquet { +namespace internal { + +uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_t rhs) { + return avx2::GreaterThanBitmapImpl(levels, num_levels, rhs); +} + +MinMax FindMinMaxAvx2(const int16_t* levels, int64_t num_levels) { + return avx2::FindMinMaxImpl(levels, num_levels); +} +} // namespace internal +} // namespace parquet diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index cfa5df1a7e0..ecea37dd4aa 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -18,176 +18,169 @@ #include #include -#if defined(ARROW_HAVE_BMI2) -#include -#endif +#include "arrow/util/bit_run_reader.h" #include "arrow/util/bit_util.h" +#include "arrow/util/cpu_info.h" #include "arrow/util/logging.h" #include "parquet/exception.h" +#include "parquet/level_comparison.h" + +#define BMI_RUNTIME_VERSION standard +#include "parquet/level_conversion_inc.h" +#undef BMI_RUNTIME_VERSION namespace parquet { namespace internal { namespace { -inline void CheckLevelRange(const int16_t* levels, int64_t num_levels, - const int16_t max_expected_level) { - int16_t min_level = std::numeric_limits::max(); - int16_t max_level = std::numeric_limits::min(); - for (int x = 0; x < num_levels; x++) { - min_level = std::min(levels[x], min_level); - max_level = std::max(levels[x], max_level); - } - if (ARROW_PREDICT_FALSE(num_levels > 0 && - (min_level < 0 || max_level > max_expected_level))) { - throw ParquetException("definition level exceeds maximum"); - } -} - -#if !defined(ARROW_HAVE_AVX512) -inline void DefinitionLevelsToBitmapScalar( - const int16_t* def_levels, int64_t num_def_levels, const int16_t max_definition_level, - const int16_t max_repetition_level, int64_t* values_read, int64_t* null_count, - uint8_t* valid_bits, int64_t valid_bits_offset) { - // We assume here that valid_bits is large enough to accommodate the - // additional definition levels and the ones that have already been written - ::arrow::internal::BitmapWriter valid_bits_writer(valid_bits, valid_bits_offset, - num_def_levels); - - // TODO(itaiin): As an interim solution we are splitting the code path here - // between repeated+flat column reads, and non-repeated+nested reads. - // Those paths need to be merged in the future - for (int i = 0; i < num_def_levels; ++i) { - if (def_levels[i] == max_definition_level) { +using ::arrow::internal::CpuInfo; + +#if !defined(ARROW_HAVE_RUNTIME_BMI2) +void DefinitionLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, int64_t* values_read, + int64_t* null_count, uint8_t* valid_bits, + int64_t valid_bits_offset) { + ::arrow::internal::FirstTimeBitmapWriter valid_bits_writer( + valid_bits, + /*start_offset=*/valid_bits_offset, + /*length=*/num_def_levels); + for (int x = 0; x < num_def_levels; x++) { + if (def_levels[x] < level_info.repeated_ancestor_def_level) { + continue; + } + if (def_levels[x] >= level_info.def_level) { valid_bits_writer.Set(); - } else if (max_repetition_level > 0) { - // repetition+flat case - if (def_levels[i] == (max_definition_level - 1)) { - valid_bits_writer.Clear(); - *null_count += 1; - } else { - continue; - } } else { - // non-repeated+nested case - if (def_levels[i] < max_definition_level) { - valid_bits_writer.Clear(); - *null_count += 1; - } else { - throw ParquetException("definition level exceeds maximum"); - } + valid_bits_writer.Clear(); + *null_count += 1; } - valid_bits_writer.Next(); } valid_bits_writer.Finish(); *values_read = valid_bits_writer.position(); + if (*null_count > 0 && level_info.null_slot_usage > 1) { + throw ParquetException( + "Null values with null_slot_usage > 1 not supported." + "(i.e. FixedSizeLists with null values are not supported"); + } } #endif -template -int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_size, - const int16_t required_definition_level, - ::arrow::internal::FirstTimeBitmapWriter* writer) { - CheckLevelRange(def_levels, batch_size, required_definition_level); - uint64_t defined_bitmap = - internal::GreaterThanBitmap(def_levels, batch_size, required_definition_level - 1); - - DCHECK_LE(batch_size, 64); - if (has_repeated_parent) { -#if defined(ARROW_HAVE_BMI2) - // This is currently a specialized code path assuming only (nested) lists - // present through the leaf (i.e. no structs). Upper level code only calls - // this method when the leaf-values are nullable (otherwise no spacing is needed), - // Because only nested lists exists it is sufficient to know that the field - // was either null or included it (i.e. definition level > max_definitation_level - // -2) If there where structs mixed in, we need to know the def_level of the - // repeated parent so we can check for def_level > "def level of repeated parent". - uint64_t present_bitmap = internal::GreaterThanBitmap(def_levels, batch_size, - required_definition_level - 2); - uint64_t selected_bits = _pext_u64(defined_bitmap, present_bitmap); - writer->AppendWord(selected_bits, ::arrow::BitUtil::PopCount(present_bitmap)); - return ::arrow::BitUtil::PopCount(selected_bits); -#else - assert(false && "must not execute this without BMI2"); -#endif - } else { - writer->AppendWord(defined_bitmap, batch_size); - return ::arrow::BitUtil::PopCount(defined_bitmap); +template +void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output, LengthType* lengths) { + LengthType* orig_pos = lengths; + std::unique_ptr<::arrow::internal::FirstTimeBitmapWriter> valid_bits_writer; + if (output->valid_bits) { + valid_bits_writer.reset(new ::arrow::internal::FirstTimeBitmapWriter( + output->valid_bits, output->valid_bits_offset, num_def_levels)); } -} + for (int x = 0; x < num_def_levels; x++) { + // Skip items that belong to empty ancenstor lists and futher nested lists. + if (def_levels[x] < level_info.repeated_ancestor_def_level || + rep_levels[x] > level_info.rep_level) { + continue; + } + if (rep_levels[x] == level_info.rep_level) { + // A continuation of an existing list. + if (lengths != nullptr) { + *lengths += 1; + } + } else { + // current_rep < list rep_level i.e. start of a list (ancenstor empty lists are + // filtered out above). + if (lengths != nullptr) { + ++lengths; + *lengths = (def_levels[x] >= level_info.def_level) ? 1 : 0; + } -template -void DefinitionLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_levels, - const int16_t required_definition_level, - int64_t* values_read, int64_t* null_count, - uint8_t* valid_bits, int64_t valid_bits_offset) { - constexpr int64_t kBitMaskSize = 64; - ::arrow::internal::FirstTimeBitmapWriter writer(valid_bits, - /*start_offset=*/valid_bits_offset, - /*length=*/num_def_levels); - int64_t set_count = 0; - *values_read = 0; - while (num_def_levels > kBitMaskSize) { - set_count += DefinitionLevelsBatchToBitmap( - def_levels, kBitMaskSize, required_definition_level, &writer); - def_levels += kBitMaskSize; - num_def_levels -= kBitMaskSize; + if (valid_bits_writer != nullptr) { + // the level_info def level for lists reflects element present level. + // the prior level distinguishes between empty lists. + if (def_levels[x] >= level_info.def_level - 1) { + valid_bits_writer->Set(); + } else { + output->null_count++; + valid_bits_writer->Clear(); + } + valid_bits_writer->Next(); + } + } } - set_count += DefinitionLevelsBatchToBitmap( - def_levels, num_def_levels, required_definition_level, &writer); - - *values_read = writer.position(); - *null_count += *values_read - set_count; - writer.Finish(); -} - -void DefinitionLevelsToBitmapLittleEndian( - const int16_t* def_levels, int64_t num_def_levels, const int16_t max_definition_level, - const int16_t max_repetition_level, int64_t* values_read, int64_t* null_count, - uint8_t* valid_bits, int64_t valid_bits_offset) { - if (max_repetition_level > 0) { -// This is a short term hack to prevent using the pext BMI2 instructions -// on non-intel platforms where performance is subpar. -// In the medium term we will hopefully be able to runtime dispatch -// to use this on intel only platforms that support pext. -#if defined(ARROW_HAVE_AVX512) - // BMI2 is required for efficient bit extraction. - DefinitionLevelsToBitmapSimd( - def_levels, num_def_levels, max_definition_level, values_read, null_count, - valid_bits, valid_bits_offset); -#else - DefinitionLevelsToBitmapScalar(def_levels, num_def_levels, max_definition_level, - max_repetition_level, values_read, null_count, - valid_bits, valid_bits_offset); -#endif // ARROW_HAVE_BMI2 - - } else { - // No BMI2 intsturctions are used for non-repeated case. - DefinitionLevelsToBitmapSimd( - def_levels, num_def_levels, max_definition_level, values_read, null_count, - valid_bits, valid_bits_offset); + if (valid_bits_writer != nullptr) { + valid_bits_writer->Finish(); + } + if (lengths != nullptr) { + output->values_read = lengths - orig_pos; + } else if (valid_bits_writer != nullptr) { + output->values_read = valid_bits_writer->position(); + } + if (output->null_count > 0 && level_info.null_slot_usage > 1) { + throw ParquetException( + "Null values with null_slot_usage > 1 not supported." + "(i.e. FixedSizeLists with null values are not supported)"); } } } // namespace void DefinitionLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, - const int16_t max_definition_level, - const int16_t max_repetition_level, int64_t* values_read, + LevelInfo level_info, int64_t* values_read, int64_t* null_count, uint8_t* valid_bits, int64_t valid_bits_offset) { -#if ARROW_LITTLE_ENDIAN - DefinitionLevelsToBitmapLittleEndian(def_levels, num_def_levels, max_definition_level, - max_repetition_level, values_read, null_count, - valid_bits, valid_bits_offset); - + if (level_info.rep_level > 0) { +#if defined(ARROW_HAVE_RUNTIME_BMI2) + using FunctionType = decltype(&standard::DefinitionLevelsToBitmapSimd); + static FunctionType fn = + CpuInfo::GetInstance()->HasEfficientBmi2() + ? DefinitionLevelsToBitmapBmi2WithRepeatedParent + : standard::DefinitionLevelsToBitmapSimd; + fn(def_levels, num_def_levels, level_info, values_read, null_count, valid_bits, + valid_bits_offset); #else - DefinitionLevelsToBitmapScalar(def_levels, num_def_levels, max_definition_level, - max_repetition_level, values_read, null_count, - valid_bits, valid_bits_offset); + DefinitionLevelsToBitmapScalar(def_levels, num_def_levels, level_info, values_read, + null_count, valid_bits, valid_bits_offset); #endif + } else { + standard::DefinitionLevelsToBitmapSimd( + def_levels, num_def_levels, level_info, values_read, null_count, valid_bits, + valid_bits_offset); + } +} + +uint64_t RunBasedExtract(uint64_t bitmap, uint64_t select_bitmap) { + return standard::RunBasedExtractImpl(bitmap, select_bitmap); +} + +void ConvertDefRepLevelsToList(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output, + ::arrow::util::variant lengths) { + if (arrow::util::holds_alternative(lengths)) { + auto int32_lengths = ::arrow::util::get(lengths); + DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, + output, int32_lengths); + } else if (arrow::util::holds_alternative(lengths)) { + auto int64_lengths = ::arrow::util::get(lengths); + DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, + output, int64_lengths); + } else { + throw ParquetException("Unrecognized variant"); + } +} + +void ConvertDefRepLevelsToBitmap(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output) { + // DefReplevelsToListInfo assumes it for the actual list method and this + // method is for parent structs, so we need to bump def and ref level. + level_info.rep_level += 1; + level_info.def_level += 1; + DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, + output, /*lengths=*/nullptr); } } // namespace internal diff --git a/cpp/src/parquet/level_conversion.h b/cpp/src/parquet/level_conversion.h index dbecb3171cf..c133f632c95 100644 --- a/cpp/src/parquet/level_conversion.h +++ b/cpp/src/parquet/level_conversion.h @@ -19,6 +19,9 @@ #include +#include "arrow/util/bitmap.h" +#include "arrow/util/optional.h" +#include "arrow/util/variant.h" #include "parquet/platform.h" #include "parquet/schema.h" @@ -41,6 +44,8 @@ struct PARQUET_EXPORT LevelInfo { repeated_ancestor_def_level == b.repeated_ancestor_def_level; } + bool HasNullableValues() const { return repeated_ancestor_def_level < def_level; } + // How many slots an undefined but present (i.e. null) element in // parquet consumes when decoding to Arrow. // "Slot" is used in the same context as the Arrow specification @@ -132,43 +137,48 @@ struct PARQUET_EXPORT LevelInfo { } }; -void PARQUET_EXPORT DefinitionLevelsToBitmap( - const int16_t* def_levels, int64_t num_def_levels, const int16_t max_definition_level, - const int16_t max_repetition_level, int64_t* values_read, int64_t* null_count, - uint8_t* valid_bits, int64_t valid_bits_offset); - -// These APIs are likely to be revised as part of ARROW-8494 to reduce duplicate code. -// They currently represent minimal functionality for vectorized computation of definition -// levels. - -#if defined(ARROW_LITTLE_ENDIAN) -/// Builds a bitmap by applying predicate to the level vector provided. -/// -/// \param[in] levels Rep or def level array. -/// \param[in] num_levels The number of levels to process (must be [0, 64]) -/// \param[in] predicate The predicate to apply (must have the signature `bool -/// predicate(int16_t)`. -/// \returns The bitmap using least significant "bit" ordering. -/// -/// N.B. Correct byte ordering is dependent on little-endian architectures. -/// -template -uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, Predicate predicate) { - // Both clang and GCC can vectorize this automatically with SSE4/AVX2. - uint64_t mask = 0; - for (int x = 0; x < num_levels; x++) { - mask |= static_cast(predicate(levels[x]) ? 1 : 0) << x; - } - return mask; -} - -/// Builds a bitmap where each set bit indicates the corresponding level is greater -/// than rhs. -static inline uint64_t GreaterThanBitmap(const int16_t* levels, int64_t num_levels, - int16_t rhs) { - return LevelsToBitmap(levels, num_levels, [rhs](int16_t value) { return value > rhs; }); -} +/// Converts def_levels to validity bitmaps for non-list arrays. +/// TODO: use input/output parameter below instead of individual return variables. +void PARQUET_EXPORT DefinitionLevelsToBitmap(const int16_t* def_levels, + int64_t num_def_levels, LevelInfo level_info, + int64_t* values_read, int64_t* null_count, + uint8_t* valid_bits, + int64_t valid_bits_offset); + +/// Input/Output structure for reconstructed validity bitmaps. +struct PARQUET_EXPORT ValidityBitmapInputOutput { + /// The number of values added to the bitmap. + int64_t values_read = 0; + /// The number of nulls encountered. + int64_t null_count = 0; + // The validity bitmp to populate. Can only be null + // for DefRepLevelsToListInfo (if all that is needed is list lengths). + uint8_t* valid_bits = NULLPTR; + /// Input only, offset into valid_bits to start at. + int64_t valid_bits_offset = 0; +}; +/// Reconstructs a validity bitmap and list lengths for a ListArray based on +/// def/rep levels. +void PARQUET_EXPORT ConvertDefRepLevelsToList( + const int16_t* def_levels, const int16_t* rep_levels, int64_t num_def_levels, + LevelInfo level_info, ValidityBitmapInputOutput* output, + ::arrow::util::variant lengths); + +/// Reconstructs a validity bitmap for a struct that has nested children. +void PARQUET_EXPORT ConvertDefRepLevelsToBitmap(const int16_t* def_levels, + const int16_t* rep_levels, + int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output); + +uint64_t RunBasedExtract(uint64_t bitmap, uint64_t selection); + +#if defined(ARROW_HAVE_RUNTIME_BMI2) +void PARQUET_EXPORT DefinitionLevelsToBitmapBmi2WithRepeatedParent( + const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, + int64_t* values_read, int64_t* null_count, uint8_t* valid_bits, + int64_t valid_bits_offset); #endif } // namespace internal diff --git a/cpp/src/parquet/level_conversion_benchmark.cc b/cpp/src/parquet/level_conversion_benchmark.cc index 4f15838d339..dccc922b53c 100644 --- a/cpp/src/parquet/level_conversion_benchmark.cc +++ b/cpp/src/parquet/level_conversion_benchmark.cc @@ -38,10 +38,12 @@ std::vector RunDefinitionLevelsToBitmap(const std::vector& def int64_t null_count = 0; std::vector bitmap(/*count=*/def_levels.size(), 0); int rep = 0; + parquet::internal::LevelInfo info; + info.def_level = kHasRepeatedElements; + info.repeated_ancestor_def_level = kPresentDefLevel; for (auto _ : *state) { parquet::internal::DefinitionLevelsToBitmap( - def_levels.data(), def_levels.size(), /*max_definition_level=*/kPresentDefLevel, - /*max_repetition_level=*/kHasRepeatedElements, &values_read, &null_count, + def_levels.data(), def_levels.size(), info, &values_read, &null_count, bitmap.data(), /*valid_bits_offset=*/(rep++ % 8) * def_levels.size()); } state->SetBytesProcessed(int64_t(state->iterations()) * def_levels.size()); diff --git a/cpp/src/parquet/level_conversion_bmi2.cc b/cpp/src/parquet/level_conversion_bmi2.cc new file mode 100644 index 00000000000..1c9a6d6c692 --- /dev/null +++ b/cpp/src/parquet/level_conversion_bmi2.cc @@ -0,0 +1,35 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +#include "parquet/level_conversion.h" + +#define BMI_RUNTIME_VERSION bmi2 +#include "parquet/level_conversion_inc.h" +#undef BMI_RUNTIME_VERSION + +namespace parquet { +namespace internal { +void DefinitionLevelsToBitmapBmi2WithRepeatedParent( + const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, + int64_t* values_read, int64_t* null_count, uint8_t* valid_bits, + int64_t valid_bits_offset) { + bmi2::DefinitionLevelsToBitmapSimd( + def_levels, num_def_levels, level_info, values_read, null_count, valid_bits, + valid_bits_offset); +} + +} // namespace internal +} // namespace parquet diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h new file mode 100644 index 00000000000..11c901b2c16 --- /dev/null +++ b/cpp/src/parquet/level_conversion_inc.h @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +#pragma once + +#include "parquet/level_conversion.h" + +#include +#include +#if defined(ARROW_HAVE_BMI2) +#include +#endif + +#include "arrow/util/bit_run_reader.h" +#include "arrow/util/bit_util.h" +#include "arrow/util/logging.h" +#include "parquet/exception.h" +#include "parquet/level_comparison.h" + +namespace parquet { +namespace internal { +namespace BMI_RUNTIME_VERSION { + +using ::arrow::internal::BitRun; +using ::arrow::internal::BitRunReader; + +/// Algorithm to simulate pext using BitRunReader for cases where all bits +/// not set or set. +uint64_t RunBasedExtractMixed(uint64_t bitmap, uint64_t select_bitmap) { + bitmap = arrow::BitUtil::FromLittleEndian(bitmap); + uint64_t new_bitmap = 0; + ::arrow::internal::BitRunReader selection(reinterpret_cast(&select_bitmap), + /*start_offset=*/0, /*length=*/64); + ::arrow::internal::BitRun run = selection.NextRun(); + int64_t selected_bits = 0; + while (run.length != 0) { + if (run.set) { + new_bitmap |= (bitmap & ::arrow::BitUtil::LeastSignficantBitMask(run.length)) + << selected_bits; + selected_bits += run.length; + } + bitmap = bitmap >> run.length; + run = selection.NextRun(); + } + return arrow::BitUtil::ToLittleEndian(new_bitmap); +} + +inline uint64_t RunBasedExtractImpl(uint64_t bitmap, uint64_t select_bitmap) { + /// These checks should be inline and are likely to be common cases. + if (select_bitmap == ~uint64_t{0}) { + return bitmap; + } else if (select_bitmap == 0) { + return 0; + } + /// Fallback to the slow method. + return RunBasedExtractMixed(bitmap, select_bitmap); +} + +inline uint64_t ExtractBits(uint64_t bitmap, uint64_t select_bitmap) { +#if defined(ARROW_HAVE_BMI2) + return _pext_u64(bitmap, select_bitmap); +#else + return RunBasedExtractImpl(bitmap, select_bitmap); +#endif +} + +template +int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_size, + LevelInfo level_info, + ::arrow::internal::FirstTimeBitmapWriter* writer) { + // Greater than level_info.def_level - 1 implies >= the def_level + uint64_t defined_bitmap = + internal::GreaterThanBitmap(def_levels, batch_size, level_info.def_level - 1); + + DCHECK_LE(batch_size, 64); + if (has_repeated_parent) { + // Greater than level_info.repeated_ancestor_def_level - 1 implies >= the + // repeated_ancenstor_def_level + uint64_t present_bitmap = internal::GreaterThanBitmap( + def_levels, batch_size, level_info.repeated_ancestor_def_level - 1); + uint64_t selected_bits = ExtractBits(defined_bitmap, present_bitmap); + writer->AppendWord(selected_bits, ::arrow::BitUtil::PopCount(present_bitmap)); + return ::arrow::BitUtil::PopCount(selected_bits); + } else { + writer->AppendWord(defined_bitmap, batch_size); + return ::arrow::BitUtil::PopCount(defined_bitmap); + } +} + +template +void DefinitionLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, int64_t* values_read, + int64_t* null_count, uint8_t* valid_bits, + int64_t valid_bits_offset) { + constexpr int64_t kBitMaskSize = 64; + ::arrow::internal::FirstTimeBitmapWriter writer(valid_bits, + /*start_offset=*/valid_bits_offset, + /*length=*/num_def_levels); + int64_t set_count = 0; + *values_read = 0; + while (num_def_levels > kBitMaskSize) { + set_count += DefinitionLevelsBatchToBitmap( + def_levels, kBitMaskSize, level_info, &writer); + def_levels += kBitMaskSize; + num_def_levels -= kBitMaskSize; + } + set_count += DefinitionLevelsBatchToBitmap( + def_levels, num_def_levels, level_info, &writer); + + *values_read = writer.position(); + *null_count += *values_read - set_count; + writer.Finish(); +} + +} // namespace BMI_RUNTIME_VERSION +} // namespace internal +} // namespace parquet diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index d4f3719289d..c967bcf26d7 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -16,6 +16,7 @@ // under the License. #include "parquet/level_conversion.h" +#include "parquet/level_comparison.h" #include #include @@ -25,10 +26,12 @@ #include "arrow/util/bit_util.h" #include "arrow/util/bitmap.h" +#include "arrow/util/ubsan.h" namespace parquet { namespace internal { +using ::arrow::internal::Bitmap; using ::testing::ElementsAreArray; std::string BitmapToString(const uint8_t* bitmap, int64_t bit_count) { @@ -45,13 +48,14 @@ TEST(TestColumnReader, DefinitionLevelsToBitmap) { std::vector valid_bits(2, 0); - const int max_def_level = 3; - const int max_rep_level = 1; + LevelInfo level_info; + level_info.def_level = 3; + level_info.rep_level = 1; int64_t values_read = -1; int64_t null_count = 0; - internal::DefinitionLevelsToBitmap(def_levels.data(), 9, max_def_level, max_rep_level, - &values_read, &null_count, valid_bits.data(), + internal::DefinitionLevelsToBitmap(def_levels.data(), 9, level_info, &values_read, + &null_count, valid_bits.data(), 0 /* valid_bits_offset */); ASSERT_EQ(9, values_read); ASSERT_EQ(1, null_count); @@ -59,8 +63,8 @@ TEST(TestColumnReader, DefinitionLevelsToBitmap) { // Call again with 0 definition levels, make sure that valid_bits is unmodified const uint8_t current_byte = valid_bits[1]; null_count = 0; - internal::DefinitionLevelsToBitmap(def_levels.data(), 0, max_def_level, max_rep_level, - &values_read, &null_count, valid_bits.data(), + internal::DefinitionLevelsToBitmap(def_levels.data(), 0, level_info, &values_read, + &null_count, valid_bits.data(), 9 /* valid_bits_offset */); ASSERT_EQ(0, values_read); ASSERT_EQ(0, null_count); @@ -74,16 +78,17 @@ TEST(TestColumnReader, DefinitionLevelsToBitmapPowerOfTwo) { std::vector def_levels = {3, 3, 3, 2, 3, 3, 3, 3}; std::vector valid_bits(1, 0); - const int max_def_level = 3; - const int max_rep_level = 1; + LevelInfo level_info; + level_info.rep_level = 1; + level_info.def_level = 3; int64_t values_read = -1; int64_t null_count = 0; // Read the latter half of the validity bitmap - internal::DefinitionLevelsToBitmap(def_levels.data() + 4, 4, max_def_level, - max_rep_level, &values_read, &null_count, - valid_bits.data(), 4 /* valid_bits_offset */); + internal::DefinitionLevelsToBitmap(def_levels.data() + 4, 4, level_info, &values_read, + &null_count, valid_bits.data(), + 4 /* valid_bits_offset */); ASSERT_EQ(4, values_read); ASSERT_EQ(0, null_count); } @@ -111,12 +116,15 @@ TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { int64_t null_count = 5; int64_t values_read = 1; + LevelInfo level_info; + level_info.repeated_ancestor_def_level = 1; + level_info.def_level = 2; + level_info.rep_level = 1; // All zeros should be ignored, ones should be unset in the bitmp and 2 should be set. std::vector def_levels = {0, 0, 0, 2, 2, 1, 0, 2}; - DefinitionLevelsToBitmap( - def_levels.data(), def_levels.size(), /*max_definition_level=*/2, - /*max_repetition_level=*/1, &values_read, &null_count, validity_bitmap.data(), - /*valid_bits_offset=*/1); + DefinitionLevelsToBitmap(def_levels.data(), def_levels.size(), level_info, &values_read, + &null_count, validity_bitmap.data(), + /*valid_bits_offset=*/1); EXPECT_EQ(BitmapToString(validity_bitmap, /*bit_count=*/8), "01101000"); for (size_t x = 1; x < validity_bitmap.size(); x++) { @@ -126,5 +134,205 @@ TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { EXPECT_EQ(values_read, 4); // value should get overwritten. } +template +void DefRepLevelsToListLengths(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + LengthType* lengths) { + for (int x = 0; x < num_def_levels; x++) { + if (rep_levels[x] < level_info.rep_level) { + // A value less than the current rep_level indicates either a start of a + // new list or an empty list. + if (def_levels[x] >= level_info.repeated_ancestor_def_level) { + ++lengths; + *lengths = def_levels[x] >= level_info.def_level ? 1 : 0; + } + } else { + // The current list length only increases when the rep level is equal + // to the current one (greater rep levels belong to the next length list. + lengths += rep_levels[x] == level_info.rep_level ? 1 : 0; + } + } +} + +class MultiLevelTestData { + public: + // Triply nested list values borrow from write_path + // [null, [[1 , null, 3], []], []], + // [[[]], [[], [1, 2]], null, [[3]]], + // null, + // [] + std::vector def_levels_{2, 7, 6, 7, 5, 3, // first row + 5, 5, 7, 7, 2, 7, // second row + 0, // third row + 1}; + std::vector rep_levels_{0, 1, 3, 3, 2, 1, // first row + 0, 1, 2, 3, 1, 1, // second row + 0, 0}; +}; + +template +class NestedListTest : public testing::Test { + public: + MultiLevelTestData test_data_; + ConverterType converter_; +}; + +template +struct RepDefLevelConverter { + using ListLengthType = ListType; + ListLengthType* ComputeListInfo(const MultiLevelTestData& test_data, + LevelInfo level_info, ValidityBitmapInputOutput* output, + ListType* lengths) { + ConvertDefRepLevelsToList(test_data.def_levels_.data(), test_data.rep_levels_.data(), + test_data.def_levels_.size(), level_info, output, lengths); + return lengths + output->values_read; + } +}; + +using ConverterTypes = + ::testing::Types, + RepDefLevelConverter>; +TYPED_TEST_CASE(NestedListTest, ConverterTypes); + +TYPED_TEST(NestedListTest, OuterMostTest) { + // [null, [[1 , null, 3], []], []], + // [[[]], [[], [1, 2]], null, [[3]]], + // null, + // [] + // -> 4 outer most lists (len(3), len(4), null, len(0)) + LevelInfo level_info; + level_info.rep_level = 1; + level_info.def_level = 2; + + std::vector lengths(5, -1); + uint64_t validity_output; + ValidityBitmapInputOutput validity_io; + validity_io.valid_bits = reinterpret_cast(&validity_output); + typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( + this->test_data_, level_info, &validity_io, lengths.data()); + + EXPECT_THAT(next_position, lengths.data() + 4); + EXPECT_THAT(lengths, testing::ElementsAre(-1, 3, 4, 0, 0)); + + EXPECT_EQ(validity_io.values_read, 4); + EXPECT_EQ(validity_io.null_count, 1); + EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/4), "1101"); +} + +TYPED_TEST(NestedListTest, MiddleListTest) { + // [null, [[1 , null, 3], []], []], + // [[[]], [[], [1, 2]], null, [[3]]], + // null, + // [] + // -> middle lists (null, len(2), len(0), + // len(1), len(2), null, len(1), + // N/A, + // N/A + LevelInfo level_info; + level_info.rep_level = 2; + level_info.def_level = 4; + level_info.repeated_ancestor_def_level = 2; + + std::vector lengths(8, -1); + uint64_t validity_output; + ValidityBitmapInputOutput validity_io; + validity_io.valid_bits = reinterpret_cast(&validity_output); + typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( + this->test_data_, level_info, &validity_io, lengths.data()); + + EXPECT_THAT(next_position, lengths.data() + 7); + EXPECT_THAT(lengths, testing::ElementsAre(-1, 0, 2, 0, 1, 2, 0, 1)); + + EXPECT_EQ(validity_io.values_read, 7); + EXPECT_EQ(validity_io.null_count, 2); + EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/7), "0111101"); +} + +TYPED_TEST(NestedListTest, InnerMostListTest) { + // [null, [[1, null, 3], []], []], + // [[[]], [[], [1, 2]], null, [[3]]], + // null, + // [] + // -> 4 inner lists (N/A, [len(3), len(0)], N/A + // len(0), [len(0), len(2)], N/A, len(1), + // N/A, + // N/A + LevelInfo level_info; + level_info.rep_level = 3; + level_info.def_level = 6; + level_info.repeated_ancestor_def_level = 4; + + std::vector lengths(7, -1); + uint64_t validity_output; + ValidityBitmapInputOutput validity_io; + validity_io.valid_bits = reinterpret_cast(&validity_output); + typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( + this->test_data_, level_info, &validity_io, lengths.data()); + + EXPECT_THAT(next_position, lengths.data() + 6); + EXPECT_THAT(lengths, testing::ElementsAre(-1, 3, 0, 0, 0, 2, 1)); + + EXPECT_EQ(validity_io.values_read, 6); + EXPECT_EQ(validity_io.null_count, 0); + EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/6), "111111"); +} + +TYPED_TEST(NestedListTest, SimpleLongList) { + LevelInfo level_info; + level_info.rep_level = 1; + level_info.def_level = 2; + level_info.repeated_ancestor_def_level = 0; + + // No empty lists. + this->test_data_.def_levels_ = std::vector(65 * 9, 2); + this->test_data_.rep_levels_.clear(); + for (int x = 0; x < 65; x++) { + this->test_data_.rep_levels_.push_back(0); + this->test_data_.rep_levels_.insert(this->test_data_.rep_levels_.end(), 8, + /*rep_level=*/1); + } + + std::vector lengths(66, -1); + std::vector expected_lengths(66, 9); + expected_lengths[0] = -1; + std::vector validity_output(9, 0); + ValidityBitmapInputOutput validity_io; + validity_io.valid_bits = validity_output.data(); + typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( + this->test_data_, level_info, &validity_io, lengths.data()); + + EXPECT_THAT(next_position, lengths.data() + 65); + EXPECT_THAT(lengths, testing::ElementsAreArray(expected_lengths)); + + EXPECT_EQ(validity_io.values_read, 65); + EXPECT_EQ(validity_io.null_count, 0); + EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/65), + "11111111 " + "11111111 " + "11111111 " + "11111111 " + "11111111 " + "11111111 " + "11111111 " + "11111111 " + "1"); +} + +TEST(RunBasedExtract, BasicTest) { + EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), 0), 0); + EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), ~uint64_t{0}), + arrow::BitUtil::ToLittleEndian(0xFF)); + + EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF00FF), + arrow::BitUtil::ToLittleEndian(0xAAAA)), + arrow::BitUtil::ToLittleEndian(0x000F)); + EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF0AFF), + arrow::BitUtil::ToLittleEndian(0xAFAA)), + arrow::BitUtil::ToLittleEndian(0x00AF)); + EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFFAAFF), + arrow::BitUtil::ToLittleEndian(0xAFAA)), + arrow::BitUtil::ToLittleEndian(0x03AF)); +} + } // namespace internal } // namespace parquet From e611fcb20f2b7fdbc99b0a17d50eafcf40f9902e Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Sun, 13 Sep 2020 05:57:32 +0000 Subject: [PATCH 02/23] add checks and try to fix windows --- cpp/src/parquet/arrow/reader.cc | 5 +++-- cpp/src/parquet/level_conversion_inc.h | 6 +++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 2da90ab1e18..fde6830eda4 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -446,7 +446,6 @@ class LeafReader : public ColumnReaderImpl { std::shared_ptr<::arrow::ChunkedArray>* out) final { *out = out_; return Status::OK(); - ; } const std::shared_ptr field() override { return field_; } @@ -525,7 +524,9 @@ class ListReader : public ColumnReaderImpl { // + 1 because arrays lengths are values + 1 std::fill(length_data + validity_io.values_read + 1, length_data + number_of_slots + 1, 0); - if (validity_io.valid_bits != nullptr) { + if (validity_io.valid_bits != nullptr && + BitUtil::BytesForBits(validity_io.values_read) < + BitUtil::BytesForBits(number_of_slots)) { std::fill(validity_io.valid_bits + BitUtil::BytesForBits(validity_io.values_read), validity_io.valid_bits + BitUtil::BytesForBits(number_of_slots), 0); } diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h index 11c901b2c16..537d974ba58 100644 --- a/cpp/src/parquet/level_conversion_inc.h +++ b/cpp/src/parquet/level_conversion_inc.h @@ -21,8 +21,12 @@ #include #include #if defined(ARROW_HAVE_BMI2) +#if defined(_MSC_VER) +#include +#else #include -#endif +#endif // _MSC_VER +#endif // ARROW_HAVE_BMI2 #include "arrow/util/bit_run_reader.h" #include "arrow/util/bit_util.h" From 2c57ab67865b4001110fb713ff9aba27fc1210ad Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Sun, 13 Sep 2020 22:01:09 +0000 Subject: [PATCH 03/23] Use VaidityInputOuput every place it can be. this will ensure invalid struct encodings don't segfault --- cpp/src/parquet/arrow/reader.cc | 42 ++++++++-------- cpp/src/parquet/column_reader.cc | 27 ++++++++--- cpp/src/parquet/level_conversion.cc | 44 ++++++++++------- cpp/src/parquet/level_conversion.h | 21 ++++---- cpp/src/parquet/level_conversion_bmi2.cc | 11 ++--- cpp/src/parquet/level_conversion_inc.h | 38 ++++++++++----- cpp/src/parquet/level_conversion_test.cc | 61 +++++++++++++----------- 7 files changed, 144 insertions(+), 100 deletions(-) diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index fde6830eda4..e2e346d35b5 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -498,6 +498,7 @@ class ListReader : public ColumnReaderImpl { RETURN_NOT_OK(item_reader_->GetRepLevels(&rep_levels, &num_levels)); std::shared_ptr validity_buffer; ::parquet::internal::ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = number_of_slots; if (field_->nullable()) { ARROW_ASSIGN_OR_RAISE(validity_buffer, AllocateBitmap(number_of_slots, ctx_->pool)); @@ -648,8 +649,10 @@ Status StructReader::BuildArray(int64_t records_to_read, std::shared_ptr* out) { std::vector> children_array_data; std::shared_ptr null_bitmap; - int64_t null_count = 0; - int64_t values_read = records_to_read; + + ::parquet::internal::ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = records_to_read; + validity_io.values_read = records_to_read; BEGIN_PARQUET_CATCH_EXCEPTIONS const int16_t* def_levels; @@ -658,46 +661,45 @@ Status StructReader::BuildArray(int64_t records_to_read, if (has_repeated_child_) { ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); + validity_io.valid_bits = null_bitmap->mutable_data(); RETURN_NOT_OK(children_.front()->GetDefLevels(&def_levels, &num_levels)); RETURN_NOT_OK(children_.front()->GetRepLevels(&rep_levels, &num_levels)); - - ::parquet::internal::ValidityBitmapInputOutput validity_io; - validity_io.valid_bits = null_bitmap->mutable_data(); ConvertDefRepLevelsToBitmap(def_levels, rep_levels, num_levels, level_info_, &validity_io); - null_count = validity_io.null_count; - values_read = validity_io.values_read; } else if (filtered_field_->nullable()) { ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); + validity_io.valid_bits = null_bitmap->mutable_data(); RETURN_NOT_OK(children_.front()->GetDefLevels(&def_levels, &num_levels)); - DefinitionLevelsToBitmap(def_levels, num_levels, level_info_, &values_read, - &null_count, null_bitmap->mutable_data(), - /*valid_bits_offset=*/0); - // Ensure all values are initialized + DefinitionLevelsToBitmap(def_levels, num_levels, level_info_, &validity_io); } - if (BitUtil::BytesForBits(values_read) < BitUtil::BytesForBits(records_to_read)) { - std::fill(null_bitmap->mutable_data() + BitUtil::BytesForBits(values_read), - null_bitmap->mutable_data() + BitUtil::BytesForBits(records_to_read), 0); + + // Ensure all values are initialized. + if (null_bitmap && BitUtil::BytesForBits(validity_io.values_read) < + BitUtil::BytesForBits(records_to_read)) { + std::fill( + null_bitmap->mutable_data() + BitUtil::BytesForBits(validity_io.values_read), + null_bitmap->mutable_data() + BitUtil::BytesForBits(records_to_read), 0); } END_PARQUET_CATCH_EXCEPTIONS // Gather children arrays and def levels for (auto& child : children_) { std::shared_ptr field; - RETURN_NOT_OK(child->BuildArray(values_read, &field)); + RETURN_NOT_OK(child->BuildArray(validity_io.values_read, &field)); ARROW_ASSIGN_OR_RAISE(std::shared_ptr array_data, ChunksToSingle(*field)); children_array_data.push_back(std::move(array_data)); } if (!filtered_field_->nullable() && !has_repeated_child_) { - values_read = children_array_data.front()->length; + validity_io.values_read = children_array_data.front()->length; } std::vector> buffers{ - null_count > 0 ? null_bitmap : std::shared_ptr()}; - auto data = std::make_shared(filtered_field_->type(), - /*length=*/values_read, std::move(buffers), - std::move(children_array_data)); + validity_io.null_count > 0 ? null_bitmap : std::shared_ptr()}; + auto data = + std::make_shared(filtered_field_->type(), + /*length=*/validity_io.values_read, std::move(buffers), + std::move(children_array_data)); std::shared_ptr result = ::arrow::MakeArray(data); RETURN_NOT_OK(result->Validate()); diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 6b5450312b3..94f4851ed63 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -930,8 +930,17 @@ int64_t TypedColumnReaderImpl::ReadBatchSpaced( info.repeated_ancestor_def_level = this->max_def_level_ - 1; info.def_level = this->max_def_level_; info.rep_level = this->max_rep_level_; - internal::DefinitionLevelsToBitmap(def_levels, num_def_levels, info, values_read, - &null_count, valid_bits, valid_bits_offset); + internal::ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = num_def_levels; + validity_io.valid_bits = valid_bits; + validity_io.valid_bits_offset = valid_bits_offset; + validity_io.null_count = null_count; + validity_io.values_read = *values_read; + + internal::DefinitionLevelsToBitmap(def_levels, num_def_levels, info, &validity_io); + null_count = validity_io.null_count; + *values_read = validity_io.values_read; + total_values = this->ReadValuesSpaced(*values_read, values, static_cast(null_count), valid_bits, valid_bits_offset); @@ -1369,14 +1378,18 @@ class TypedRecordReader : public ColumnReaderImplBase, int64_t null_count = 0; if (leaf_info_.HasNullableValues()) { - int64_t values_with_nulls = 0; + ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = levels_position_ - start_levels_position; + validity_io.valid_bits = valid_bits_->mutable_data(); + validity_io.valid_bits_offset = values_written_; + DefinitionLevelsToBitmap(def_levels() + start_levels_position, levels_position_ - start_levels_position, leaf_info_, - &values_with_nulls, &null_count, - valid_bits_->mutable_data(), values_written_); - values_to_read = values_with_nulls - null_count; + &validity_io); + values_to_read = validity_io.values_read - validity_io.null_count; + null_count = validity_io.null_count; DCHECK_GE(values_to_read, 0); - ReadValuesSpaced(values_with_nulls, null_count); + ReadValuesSpaced(validity_io.values_read, null_count); } else { DCHECK_GE(values_to_read, 0); ReadValuesDense(values_to_read); diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index ecea37dd4aa..bbbb16cf0e1 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -39,27 +39,32 @@ using ::arrow::internal::CpuInfo; #if !defined(ARROW_HAVE_RUNTIME_BMI2) void DefinitionLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, int64_t* values_read, - int64_t* null_count, uint8_t* valid_bits, - int64_t valid_bits_offset) { + ValidityBitmapInputOutput* output) { ::arrow::internal::FirstTimeBitmapWriter valid_bits_writer( - valid_bits, - /*start_offset=*/valid_bits_offset, - /*length=*/num_def_levels); + output->valid_bits, + /*start_offset=*/output->valid_bits_offset, + /*length=*/output->num_def_levels); for (int x = 0; x < num_def_levels; x++) { if (def_levels[x] < level_info.repeated_ancestor_def_level) { continue; } + if (ARROW_PREDICT_FALSE(valid_bits_writer.position() > + output->values_read_upper_bound)) { + std::stringstream ss; + ss << "Definition levels exceeded upper bound: " << output->values_read_upper_bound; + throw ParquetException(ss.str()); + } if (def_levels[x] >= level_info.def_level) { valid_bits_writer.Set(); } else { valid_bits_writer.Clear(); - *null_count += 1; + output->null_count += 1; } valid_bits_writer.Next(); } valid_bits_writer.Finish(); - *values_read = valid_bits_writer.position(); - if (*null_count > 0 && level_info.null_slot_usage > 1) { + output->values_read = valid_bits_writer.position(); + if (output->null_count > 0 && level_info.null_slot_usage > 1) { throw ParquetException( "Null values with null_slot_usage > 1 not supported." "(i.e. FixedSizeLists with null values are not supported"); @@ -83,6 +88,16 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels rep_levels[x] > level_info.rep_level) { continue; } + + if (ARROW_PREDICT_FALSE( + (valid_bits_writer != nullptr && + valid_bits_writer->position() > output->values_read_upper_bound) || + (lengths - orig_pos) > output->values_read_upper_bound)) { + std::stringstream ss; + ss << "Definition levels exceeded upper bound: " << output->values_read_upper_bound; + throw ParquetException(ss.str()); + } + if (rep_levels[x] == level_info.rep_level) { // A continuation of an existing list. if (lengths != nullptr) { @@ -127,9 +142,7 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels } // namespace void DefinitionLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, - LevelInfo level_info, int64_t* values_read, - int64_t* null_count, uint8_t* valid_bits, - int64_t valid_bits_offset) { + LevelInfo level_info, ValidityBitmapInputOutput* output) { if (level_info.rep_level > 0) { #if defined(ARROW_HAVE_RUNTIME_BMI2) using FunctionType = decltype(&standard::DefinitionLevelsToBitmapSimd); @@ -137,17 +150,14 @@ void DefinitionLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, CpuInfo::GetInstance()->HasEfficientBmi2() ? DefinitionLevelsToBitmapBmi2WithRepeatedParent : standard::DefinitionLevelsToBitmapSimd; - fn(def_levels, num_def_levels, level_info, values_read, null_count, valid_bits, - valid_bits_offset); + fn(def_levels, num_def_levels, level_info, output); #else - DefinitionLevelsToBitmapScalar(def_levels, num_def_levels, level_info, values_read, - null_count, valid_bits, valid_bits_offset); + DefinitionLevelsToBitmapScalar(def_levels, num_def_levels, level_info, output); #endif } else { standard::DefinitionLevelsToBitmapSimd( - def_levels, num_def_levels, level_info, values_read, null_count, valid_bits, - valid_bits_offset); + def_levels, num_def_levels, level_info, output); } } diff --git a/cpp/src/parquet/level_conversion.h b/cpp/src/parquet/level_conversion.h index c133f632c95..9534ae10e7e 100644 --- a/cpp/src/parquet/level_conversion.h +++ b/cpp/src/parquet/level_conversion.h @@ -137,16 +137,13 @@ struct PARQUET_EXPORT LevelInfo { } }; -/// Converts def_levels to validity bitmaps for non-list arrays. -/// TODO: use input/output parameter below instead of individual return variables. -void PARQUET_EXPORT DefinitionLevelsToBitmap(const int16_t* def_levels, - int64_t num_def_levels, LevelInfo level_info, - int64_t* values_read, int64_t* null_count, - uint8_t* valid_bits, - int64_t valid_bits_offset); - /// Input/Output structure for reconstructed validity bitmaps. struct PARQUET_EXPORT ValidityBitmapInputOutput { + /// The maximum number of values_read expected (actual + /// values read must be less than or equal to this value. + /// If this number is exceeded methods will throw a + /// ParquetException. + int64_t values_read_upper_bound = 0; /// The number of values added to the bitmap. int64_t values_read = 0; /// The number of nulls encountered. @@ -158,6 +155,11 @@ struct PARQUET_EXPORT ValidityBitmapInputOutput { int64_t valid_bits_offset = 0; }; +/// Converts def_levels to validity bitmaps for non-list arrays. +void PARQUET_EXPORT DefinitionLevelsToBitmap(const int16_t* def_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output); + /// Reconstructs a validity bitmap and list lengths for a ListArray based on /// def/rep levels. void PARQUET_EXPORT ConvertDefRepLevelsToList( @@ -177,8 +179,7 @@ uint64_t RunBasedExtract(uint64_t bitmap, uint64_t selection); #if defined(ARROW_HAVE_RUNTIME_BMI2) void PARQUET_EXPORT DefinitionLevelsToBitmapBmi2WithRepeatedParent( const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, - int64_t* values_read, int64_t* null_count, uint8_t* valid_bits, - int64_t valid_bits_offset); + ValidityBitmapInputOutput* output); #endif } // namespace internal diff --git a/cpp/src/parquet/level_conversion_bmi2.cc b/cpp/src/parquet/level_conversion_bmi2.cc index 1c9a6d6c692..fa291091a30 100644 --- a/cpp/src/parquet/level_conversion_bmi2.cc +++ b/cpp/src/parquet/level_conversion_bmi2.cc @@ -22,13 +22,12 @@ namespace parquet { namespace internal { -void DefinitionLevelsToBitmapBmi2WithRepeatedParent( - const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, - int64_t* values_read, int64_t* null_count, uint8_t* valid_bits, - int64_t valid_bits_offset) { +void DefinitionLevelsToBitmapBmi2WithRepeatedParent(const int16_t* def_levels, + int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output) { bmi2::DefinitionLevelsToBitmapSimd( - def_levels, num_def_levels, level_info, values_read, null_count, valid_bits, - valid_bits_offset); + def_levels, num_def_levels, level_info, output); } } // namespace internal diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h index 537d974ba58..86fb5b7b22e 100644 --- a/cpp/src/parquet/level_conversion_inc.h +++ b/cpp/src/parquet/level_conversion_inc.h @@ -83,7 +83,7 @@ inline uint64_t ExtractBits(uint64_t bitmap, uint64_t select_bitmap) { template int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_size, - LevelInfo level_info, + int64_t upper_bound_remaining, LevelInfo level_info, ::arrow::internal::FirstTimeBitmapWriter* writer) { // Greater than level_info.def_level - 1 implies >= the def_level uint64_t defined_bitmap = @@ -96,9 +96,19 @@ int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t b uint64_t present_bitmap = internal::GreaterThanBitmap( def_levels, batch_size, level_info.repeated_ancestor_def_level - 1); uint64_t selected_bits = ExtractBits(defined_bitmap, present_bitmap); - writer->AppendWord(selected_bits, ::arrow::BitUtil::PopCount(present_bitmap)); + int64_t selected_count = ::arrow::BitUtil::PopCount(present_bitmap); + if (selected_count > upper_bound_remaining) { + throw ParquetException("Values read exceeded upper bound"); + } + writer->AppendWord(selected_bits, selected_count); return ::arrow::BitUtil::PopCount(selected_bits); } else { + if (batch_size > upper_bound_remaining) { + std::stringstream ss; + ss << "Values read exceeded upper bound"; + throw ParquetException(ss.str()); + } + writer->AppendWord(defined_bitmap, batch_size); return ::arrow::BitUtil::PopCount(defined_bitmap); } @@ -106,26 +116,28 @@ int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t b template void DefinitionLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_levels, - LevelInfo level_info, int64_t* values_read, - int64_t* null_count, uint8_t* valid_bits, - int64_t valid_bits_offset) { + LevelInfo level_info, + ValidityBitmapInputOutput* output) { constexpr int64_t kBitMaskSize = 64; - ::arrow::internal::FirstTimeBitmapWriter writer(valid_bits, - /*start_offset=*/valid_bits_offset, - /*length=*/num_def_levels); + ::arrow::internal::FirstTimeBitmapWriter writer( + output->valid_bits, + /*start_offset=*/output->valid_bits_offset, + /*length=*/num_def_levels); int64_t set_count = 0; - *values_read = 0; + output->values_read = 0; + int64_t values_read_remaining = output->values_read_upper_bound; while (num_def_levels > kBitMaskSize) { set_count += DefinitionLevelsBatchToBitmap( - def_levels, kBitMaskSize, level_info, &writer); + def_levels, kBitMaskSize, values_read_remaining, level_info, &writer); def_levels += kBitMaskSize; num_def_levels -= kBitMaskSize; + values_read_remaining = output->values_read_upper_bound - writer.position(); } set_count += DefinitionLevelsBatchToBitmap( - def_levels, num_def_levels, level_info, &writer); + def_levels, num_def_levels, values_read_remaining, level_info, &writer); - *values_read = writer.position(); - *null_count += *values_read - set_count; + output->values_read = writer.position(); + output->null_count += output->values_read - set_count; writer.Finish(); } diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index c967bcf26d7..89e0b01a6a6 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -52,22 +52,22 @@ TEST(TestColumnReader, DefinitionLevelsToBitmap) { level_info.def_level = 3; level_info.rep_level = 1; - int64_t values_read = -1; - int64_t null_count = 0; - internal::DefinitionLevelsToBitmap(def_levels.data(), 9, level_info, &values_read, - &null_count, valid_bits.data(), - 0 /* valid_bits_offset */); - ASSERT_EQ(9, values_read); - ASSERT_EQ(1, null_count); + ValidityBitmapInputOutput io; + io.values_read_upper_bound = def_levels.size(); + io.values_read = -1; + io.valid_bits = valid_bits.data(); + + internal::DefinitionLevelsToBitmap(def_levels.data(), 9, level_info, &io); + ASSERT_EQ(9, io.values_read); + ASSERT_EQ(1, io.null_count); // Call again with 0 definition levels, make sure that valid_bits is unmodified const uint8_t current_byte = valid_bits[1]; - null_count = 0; - internal::DefinitionLevelsToBitmap(def_levels.data(), 0, level_info, &values_read, - &null_count, valid_bits.data(), - 9 /* valid_bits_offset */); - ASSERT_EQ(0, values_read); - ASSERT_EQ(0, null_count); + io.null_count = 0; + internal::DefinitionLevelsToBitmap(def_levels.data(), 0, level_info, &io); + + ASSERT_EQ(0, io.values_read); + ASSERT_EQ(0, io.null_count); ASSERT_EQ(current_byte, valid_bits[1]); } @@ -82,15 +82,15 @@ TEST(TestColumnReader, DefinitionLevelsToBitmapPowerOfTwo) { level_info.rep_level = 1; level_info.def_level = 3; - int64_t values_read = -1; - int64_t null_count = 0; + ValidityBitmapInputOutput io; + io.values_read_upper_bound = def_levels.size(); + io.values_read = -1; + io.valid_bits = valid_bits.data(); // Read the latter half of the validity bitmap - internal::DefinitionLevelsToBitmap(def_levels.data() + 4, 4, level_info, &values_read, - &null_count, valid_bits.data(), - 4 /* valid_bits_offset */); - ASSERT_EQ(4, values_read); - ASSERT_EQ(0, null_count); + internal::DefinitionLevelsToBitmap(def_levels.data() + 4, 4, level_info, &io); + ASSERT_EQ(4, io.values_read); + ASSERT_EQ(0, io.null_count); } #if defined(ARROW_LITTLE_ENDIAN) @@ -113,8 +113,13 @@ TEST(GreaterThanBitmap, GeneratesExpectedBitmasks) { TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { std::vector validity_bitmap(/*count*/ 8, 0); - int64_t null_count = 5; - int64_t values_read = 1; + + ValidityBitmapInputOutput io; + io.values_read_upper_bound = 64; + io.values_read = 1; + io.null_count = 5; + io.valid_bits = validity_bitmap.data(); + io.valid_bits_offset = 1; LevelInfo level_info; level_info.repeated_ancestor_def_level = 1; @@ -122,16 +127,14 @@ TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { level_info.rep_level = 1; // All zeros should be ignored, ones should be unset in the bitmp and 2 should be set. std::vector def_levels = {0, 0, 0, 2, 2, 1, 0, 2}; - DefinitionLevelsToBitmap(def_levels.data(), def_levels.size(), level_info, &values_read, - &null_count, validity_bitmap.data(), - /*valid_bits_offset=*/1); + DefinitionLevelsToBitmap(def_levels.data(), def_levels.size(), level_info, &io); EXPECT_EQ(BitmapToString(validity_bitmap, /*bit_count=*/8), "01101000"); for (size_t x = 1; x < validity_bitmap.size(); x++) { EXPECT_EQ(validity_bitmap[x], 0) << "index: " << x; } - EXPECT_EQ(null_count, /*5 + 1 =*/6); - EXPECT_EQ(values_read, 4); // value should get overwritten. + EXPECT_EQ(io.null_count, /*5 + 1 =*/6); + EXPECT_EQ(io.values_read, 4); // value should get overwritten. } template @@ -207,6 +210,7 @@ TYPED_TEST(NestedListTest, OuterMostTest) { std::vector lengths(5, -1); uint64_t validity_output; ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = 4; validity_io.valid_bits = reinterpret_cast(&validity_output); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, lengths.data()); @@ -236,6 +240,7 @@ TYPED_TEST(NestedListTest, MiddleListTest) { std::vector lengths(8, -1); uint64_t validity_output; ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = 7; validity_io.valid_bits = reinterpret_cast(&validity_output); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, lengths.data()); @@ -265,6 +270,7 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { std::vector lengths(7, -1); uint64_t validity_output; ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = 6; validity_io.valid_bits = reinterpret_cast(&validity_output); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, lengths.data()); @@ -297,6 +303,7 @@ TYPED_TEST(NestedListTest, SimpleLongList) { expected_lengths[0] = -1; std::vector validity_output(9, 0); ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = 65; validity_io.valid_bits = validity_output.data(); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, lengths.data()); From 30882b2ab0d655a6885b37e44523ed1b13fd5107 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Sun, 13 Sep 2020 22:14:02 +0000 Subject: [PATCH 04/23] use cumulative sums in level_decoder --- cpp/src/parquet/arrow/reader.cc | 3 -- cpp/src/parquet/level_conversion.cc | 4 ++- cpp/src/parquet/level_conversion_test.cc | 40 +++++++----------------- 3 files changed, 14 insertions(+), 33 deletions(-) diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index e2e346d35b5..6390dd5ebc6 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -516,9 +516,6 @@ class ListReader : public ColumnReaderImpl { ::parquet::internal::ConvertDefRepLevelsToList( def_levels, rep_levels, num_levels, level_info_, &validity_io, length_data); END_PARQUET_CATCH_EXCEPTIONS - for (int x = 1; x <= validity_io.values_read; x++) { - length_data[x] += length_data[x - 1]; - } // We might return less then the requested slot (i.e. reaching an end of a file) // ensure we've set all the bits here. if (validity_io.values_read < number_of_slots) { diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index bbbb16cf0e1..b19119c67cf 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -108,7 +108,9 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels // filtered out above). if (lengths != nullptr) { ++lengths; - *lengths = (def_levels[x] >= level_info.def_level) ? 1 : 0; + // Use cumulative lengths because this is what the more common + // Arrow list types expect. + *lengths = ((def_levels[x] >= level_info.def_level) ? 1 : 0) + *(lengths-1); } if (valid_bits_writer != nullptr) { diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index 89e0b01a6a6..b23da6285db 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -137,26 +137,6 @@ TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { EXPECT_EQ(io.values_read, 4); // value should get overwritten. } -template -void DefRepLevelsToListLengths(const int16_t* def_levels, const int16_t* rep_levels, - int64_t num_def_levels, LevelInfo level_info, - LengthType* lengths) { - for (int x = 0; x < num_def_levels; x++) { - if (rep_levels[x] < level_info.rep_level) { - // A value less than the current rep_level indicates either a start of a - // new list or an empty list. - if (def_levels[x] >= level_info.repeated_ancestor_def_level) { - ++lengths; - *lengths = def_levels[x] >= level_info.def_level ? 1 : 0; - } - } else { - // The current list length only increases when the rep level is equal - // to the current one (greater rep levels belong to the next length list. - lengths += rep_levels[x] == level_info.rep_level ? 1 : 0; - } - } -} - class MultiLevelTestData { public: // Triply nested list values borrow from write_path @@ -207,7 +187,7 @@ TYPED_TEST(NestedListTest, OuterMostTest) { level_info.rep_level = 1; level_info.def_level = 2; - std::vector lengths(5, -1); + std::vector lengths(5, 0); uint64_t validity_output; ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 4; @@ -216,7 +196,7 @@ TYPED_TEST(NestedListTest, OuterMostTest) { this->test_data_, level_info, &validity_io, lengths.data()); EXPECT_THAT(next_position, lengths.data() + 4); - EXPECT_THAT(lengths, testing::ElementsAre(-1, 3, 4, 0, 0)); + EXPECT_THAT(lengths, testing::ElementsAre(0, 3, 7, 7, 7)); EXPECT_EQ(validity_io.values_read, 4); EXPECT_EQ(validity_io.null_count, 1); @@ -237,7 +217,7 @@ TYPED_TEST(NestedListTest, MiddleListTest) { level_info.def_level = 4; level_info.repeated_ancestor_def_level = 2; - std::vector lengths(8, -1); + std::vector lengths(8, 0); uint64_t validity_output; ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 7; @@ -246,7 +226,7 @@ TYPED_TEST(NestedListTest, MiddleListTest) { this->test_data_, level_info, &validity_io, lengths.data()); EXPECT_THAT(next_position, lengths.data() + 7); - EXPECT_THAT(lengths, testing::ElementsAre(-1, 0, 2, 0, 1, 2, 0, 1)); + EXPECT_THAT(lengths, testing::ElementsAre(0, 0, 2, 2, 3, 5, 5, 6)); EXPECT_EQ(validity_io.values_read, 7); EXPECT_EQ(validity_io.null_count, 2); @@ -267,7 +247,7 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { level_info.def_level = 6; level_info.repeated_ancestor_def_level = 4; - std::vector lengths(7, -1); + std::vector lengths(7, 0); uint64_t validity_output; ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 6; @@ -276,7 +256,7 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { this->test_data_, level_info, &validity_io, lengths.data()); EXPECT_THAT(next_position, lengths.data() + 6); - EXPECT_THAT(lengths, testing::ElementsAre(-1, 3, 0, 0, 0, 2, 1)); + EXPECT_THAT(lengths, testing::ElementsAre(0, 3, 3, 3, 3, 5, 6)); EXPECT_EQ(validity_io.values_read, 6); EXPECT_EQ(validity_io.null_count, 0); @@ -298,9 +278,11 @@ TYPED_TEST(NestedListTest, SimpleLongList) { /*rep_level=*/1); } - std::vector lengths(66, -1); - std::vector expected_lengths(66, 9); - expected_lengths[0] = -1; + std::vector lengths(66, 0); + std::vector expected_lengths(66, 0); + for (size_t x = 1; x < expected_lengths.size(); x++) { + expected_lengths[x] = x * 9; + } std::vector validity_output(9, 0); ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 65; From 4e73ebc8521953deb511a8be61df6de831cd0e79 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Sun, 13 Sep 2020 22:19:46 +0000 Subject: [PATCH 05/23] ninja format and fix scalar version --- cpp/src/parquet/level_conversion.cc | 10 +++++----- cpp/src/parquet/level_conversion_test.cc | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index b19119c67cf..2ca49a1a1f7 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -38,12 +38,12 @@ using ::arrow::internal::CpuInfo; #if !defined(ARROW_HAVE_RUNTIME_BMI2) void DefinitionLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, - LevelInfo level_info, int64_t* values_read, + LevelInfo level_info, ValidityBitmapInputOutput* output) { ::arrow::internal::FirstTimeBitmapWriter valid_bits_writer( output->valid_bits, /*start_offset=*/output->valid_bits_offset, - /*length=*/output->num_def_levels); + /*length=*/num_def_levels); for (int x = 0; x < num_def_levels; x++) { if (def_levels[x] < level_info.repeated_ancestor_def_level) { continue; @@ -108,9 +108,9 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels // filtered out above). if (lengths != nullptr) { ++lengths; - // Use cumulative lengths because this is what the more common - // Arrow list types expect. - *lengths = ((def_levels[x] >= level_info.def_level) ? 1 : 0) + *(lengths-1); + // Use cumulative lengths because this is what the more common + // Arrow list types expect. + *lengths = ((def_levels[x] >= level_info.def_level) ? 1 : 0) + *(lengths - 1); } if (valid_bits_writer != nullptr) { diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index b23da6285db..0e91886e68d 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -281,7 +281,7 @@ TYPED_TEST(NestedListTest, SimpleLongList) { std::vector lengths(66, 0); std::vector expected_lengths(66, 0); for (size_t x = 1; x < expected_lengths.size(); x++) { - expected_lengths[x] = x * 9; + expected_lengths[x] = x * 9; } std::vector validity_output(9, 0); ValidityBitmapInputOutput validity_io; From aca2bed9f4c461fc5163205e74ef09c230f91e88 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Sun, 13 Sep 2020 22:39:43 +0000 Subject: [PATCH 06/23] use non repeated child if it is available --- cpp/src/parquet/arrow/reader.cc | 52 ++++++++++++++++++--------------- 1 file changed, 28 insertions(+), 24 deletions(-) diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 6390dd5ebc6..a954dcf2a1a 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -561,22 +561,29 @@ class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { explicit StructReader(std::shared_ptr ctx, std::shared_ptr filtered_field, ::parquet::internal::LevelInfo level_info, - std::vector>&& children) + std::vector> children) : ctx_(std::move(ctx)), filtered_field_(std::move(filtered_field)), level_info_(level_info), children_(std::move(children)) { - has_repeated_child_ = IsOrHasRepeatedChild(); - } - - bool IsOrHasRepeatedChild() const final { - if (children_.empty()) { - return false; + // There could be a mix of children some might be repeated some might not be. + // If possible use one that isn't since that will be guaranteed to have the least + // number of rep/def levels. + auto result = std::find_if(children_.begin(), children_.end(), + [](const std::unique_ptr& child) { + return !child->IsOrHasRepeatedChild(); + }); + if (result != children_.end()) { + def_rep_level_child_ = result->get(); + has_repeated_child_ = false; + } else if (!children_.empty()) { + def_rep_level_child_ = children_.front().get(); + has_repeated_child_ = true; } - // Needs to be kept consistent with how rep/def levels are chosen. - return children_.front()->IsOrHasRepeatedChild(); } + bool IsOrHasRepeatedChild() const final { return has_repeated_child_; } + Status LoadBatch(int64_t records_to_read) override { for (const std::unique_ptr& reader : children_) { RETURN_NOT_OK(reader->LoadBatch(records_to_read)); @@ -593,6 +600,7 @@ class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { const std::shared_ptr filtered_field_; const ::parquet::internal::LevelInfo level_info_; const std::vector> children_; + ColumnReaderImpl* def_rep_level_child_ = nullptr; bool has_repeated_child_; }; @@ -604,13 +612,10 @@ Status StructReader::GetDefLevels(const int16_t** data, int64_t* length) { } // This method should only be called when this struct or one of its parents - // are optional/repeated. Meaning all children must have rep/def levels associated - // with them. Furthermore, it shouldn't matter which definition levels we - // choose at this point, because callers sitting above this on the call graph - // should all have identical levels for the purposes of decoding (technically - // we could optimize by choosing a child with the least number of rep/def levels - // but we can leave that for future work). - RETURN_NOT_OK(children_[0]->GetDefLevels(data, length)); + // are optional/repeated or it has a repeated child. + // Meaning all children must have rep/def levels associated + // with them. + RETURN_NOT_OK(def_rep_level_child_->GetDefLevels(data, length)); if (*data == nullptr) { // Only happens if there are actually 0 rows available. @@ -628,11 +633,10 @@ Status StructReader::GetRepLevels(const int16_t** data, int64_t* length) { } // This method should only be called when this struct or one of its parents - // are optional/repeated. Meaning all children must have rep/def levels associated - // with them. Furthermore, it shouldn't matter which repetition levels we - // choose at this point, because callers sitting above this on the call graph - // should all have identical levels for the purposes of decoding them. - RETURN_NOT_OK(children_[0]->GetRepLevels(data, length)); + // are optional/repeated or it has repeated child. + // Meaning all children must have rep/def levels associated + // with them. + RETURN_NOT_OK(def_rep_level_child_->GetRepLevels(data, length)); if (data == nullptr) { // Only happens if there are actually 0 rows available. @@ -659,14 +663,14 @@ Status StructReader::BuildArray(int64_t records_to_read, if (has_repeated_child_) { ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); validity_io.valid_bits = null_bitmap->mutable_data(); - RETURN_NOT_OK(children_.front()->GetDefLevels(&def_levels, &num_levels)); - RETURN_NOT_OK(children_.front()->GetRepLevels(&rep_levels, &num_levels)); + RETURN_NOT_OK(GetDefLevels(&def_levels, &num_levels)); + RETURN_NOT_OK(GetRepLevels(&rep_levels, &num_levels)); ConvertDefRepLevelsToBitmap(def_levels, rep_levels, num_levels, level_info_, &validity_io); } else if (filtered_field_->nullable()) { ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); validity_io.valid_bits = null_bitmap->mutable_data(); - RETURN_NOT_OK(children_.front()->GetDefLevels(&def_levels, &num_levels)); + RETURN_NOT_OK(GetDefLevels(&def_levels, &num_levels)); DefinitionLevelsToBitmap(def_levels, num_levels, level_info_, &validity_io); } From c0a31dece140fc0c965d7c5db4816d3172e8a0b6 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Mon, 14 Sep 2020 04:07:59 +0000 Subject: [PATCH 07/23] try to fix windows and check for overflow --- cpp/src/parquet/level_conversion.cc | 11 +++++++- cpp/src/parquet/level_conversion_inc.h | 2 +- cpp/src/parquet/level_conversion_test.cc | 32 ++++++++++++++++++++++++ 3 files changed, 43 insertions(+), 2 deletions(-) diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index 2ca49a1a1f7..b0187fc26d4 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -101,6 +101,9 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels if (rep_levels[x] == level_info.rep_level) { // A continuation of an existing list. if (lengths != nullptr) { + if (ARROW_PREDICT_FALSE(*lengths == std::numeric_limits::max())) { + throw ParquetException("List index overflow."); + } *lengths += 1; } } else { @@ -110,7 +113,13 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels ++lengths; // Use cumulative lengths because this is what the more common // Arrow list types expect. - *lengths = ((def_levels[x] >= level_info.def_level) ? 1 : 0) + *(lengths - 1); + *lengths = *(lengths - 1); + if (def_levels[x] >= level_info.def_level) { + if (ARROW_PREDICT_FALSE(*lengths == std::numeric_limits::max())) { + throw ParquetException("List index overflow."); + } + *lengths += 1; + } } if (valid_bits_writer != nullptr) { diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h index 86fb5b7b22e..b24fe64fa21 100644 --- a/cpp/src/parquet/level_conversion_inc.h +++ b/cpp/src/parquet/level_conversion_inc.h @@ -74,7 +74,7 @@ inline uint64_t RunBasedExtractImpl(uint64_t bitmap, uint64_t select_bitmap) { } inline uint64_t ExtractBits(uint64_t bitmap, uint64_t select_bitmap) { -#if defined(ARROW_HAVE_BMI2) +#if defined(ARROW_HAVE_BMI2) && !defined(__MINGW32__) return _pext_u64(bitmap, select_bitmap); #else return RunBasedExtractImpl(bitmap, select_bitmap); diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index 0e91886e68d..36bf1a6ec7e 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -17,6 +17,7 @@ #include "parquet/level_conversion.h" #include "parquet/level_comparison.h" +#include "parquet/test_util.h" #include #include @@ -307,6 +308,37 @@ TYPED_TEST(NestedListTest, SimpleLongList) { "1"); } +TYPED_TEST(NestedListTest, TestOverflow) { + LevelInfo level_info; + level_info.rep_level = 1; + level_info.def_level = 2; + level_info.repeated_ancestor_def_level = 0; + + // No empty lists. + this->test_data_.def_levels_ = std::vector{2}; + this->test_data_.rep_levels_ = std::vector{0}; + + std::vector lengths( + 2, std::numeric_limits::max()); + + std::vector validity_output(1, 0); + ValidityBitmapInputOutput validity_io; + validity_io.values_read_upper_bound = 1; + validity_io.valid_bits = validity_output.data(); + ASSERT_THROW(this->converter_.ComputeListInfo(this->test_data_, level_info, + &validity_io, lengths.data()), + ParquetException); + + // Same thing should happen if the list already existed. + this->test_data_.rep_levels_ = std::vector{1}; + ASSERT_THROW(this->converter_.ComputeListInfo(this->test_data_, level_info, + &validity_io, lengths.data()), + ParquetException); + + // Should be OK because it shouldn't increment. + this->test_data_.def_levels_ = std::vector{0}; +} + TEST(RunBasedExtract, BasicTest) { EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), 0), 0); EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), ~uint64_t{0}), From 0517131fe93ba0ba48a7522f4c311f86a759f331 Mon Sep 17 00:00:00 2001 From: emkornfield Date: Sun, 13 Sep 2020 23:10:03 -0700 Subject: [PATCH 08/23] Update level_conversion.h Add export to RunBasedExtract --- cpp/src/parquet/level_conversion.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/parquet/level_conversion.h b/cpp/src/parquet/level_conversion.h index 9534ae10e7e..df7a68729c6 100644 --- a/cpp/src/parquet/level_conversion.h +++ b/cpp/src/parquet/level_conversion.h @@ -174,7 +174,7 @@ void PARQUET_EXPORT ConvertDefRepLevelsToBitmap(const int16_t* def_levels, LevelInfo level_info, ValidityBitmapInputOutput* output); -uint64_t RunBasedExtract(uint64_t bitmap, uint64_t selection); +uint64_t PARQUET_EXPORT RunBasedExtract(uint64_t bitmap, uint64_t selection); #if defined(ARROW_HAVE_RUNTIME_BMI2) void PARQUET_EXPORT DefinitionLevelsToBitmapBmi2WithRepeatedParent( From d0a305f73a98da91cdf026cef100d6d169463836 Mon Sep 17 00:00:00 2001 From: emkornfield Date: Sun, 13 Sep 2020 23:11:16 -0700 Subject: [PATCH 09/23] Update level_comparison.h add PARQUET_EXPORT to GreaterThanBitmap --- cpp/src/parquet/level_comparison.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/parquet/level_comparison.h b/cpp/src/parquet/level_comparison.h index 92e5f0d3cb3..0921d6b771a 100644 --- a/cpp/src/parquet/level_comparison.h +++ b/cpp/src/parquet/level_comparison.h @@ -51,7 +51,7 @@ inline uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, /// Builds a bitmap where each set bit indicates the corresponding level is greater /// than rhs. -uint64_t GreaterThanBitmap(const int16_t* levels, int64_t num_levels, int16_t rhs); +uint64_t PARQUET_EXPORT GreaterThanBitmap(const int16_t* levels, int64_t num_levels, int16_t rhs); #if defined(ARROW_HAVE_RUNTIME_AVX2) uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_t rhs); From cf14106e4f2c1e856ded5f3718a3fc114ffb5690 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Mon, 14 Sep 2020 06:26:49 +0000 Subject: [PATCH 10/23] add include and format --- cpp/src/parquet/level_comparison.h | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/src/parquet/level_comparison.h b/cpp/src/parquet/level_comparison.h index 0921d6b771a..af9d05e9a69 100644 --- a/cpp/src/parquet/level_comparison.h +++ b/cpp/src/parquet/level_comparison.h @@ -20,6 +20,7 @@ #include #include "arrow/util/bit_util.h" +#include "parquet/platform.h" namespace parquet { namespace internal { @@ -51,7 +52,8 @@ inline uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, /// Builds a bitmap where each set bit indicates the corresponding level is greater /// than rhs. -uint64_t PARQUET_EXPORT GreaterThanBitmap(const int16_t* levels, int64_t num_levels, int16_t rhs); +uint64_t PARQUET_EXPORT GreaterThanBitmap(const int16_t* levels, int64_t num_levels, + int16_t rhs); #if defined(ARROW_HAVE_RUNTIME_AVX2) uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_t rhs); From f42d1cd498fa74713b305b625fd5c856d78a3317 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Tue, 15 Sep 2020 03:32:24 +0000 Subject: [PATCH 11/23] Revert "remove v1 test" This reverts commit d19ebc02fec16bc363ba610833ee76d8e9b02668. --- python/pyarrow/tests/test_parquet.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/pyarrow/tests/test_parquet.py b/python/pyarrow/tests/test_parquet.py index 81d6efcc6be..cbc8ba026b0 100644 --- a/python/pyarrow/tests/test_parquet.py +++ b/python/pyarrow/tests/test_parquet.py @@ -692,6 +692,14 @@ def test_pandas_can_write_nested_data(tempdir): # This succeeds under V2 _write_table(arrow_table, imos) + # Under V1 it fails. + with pytest.raises(ValueError): + import os + os.environ['ARROW_PARQUET_WRITER_ENGINE'] = 'V1' + imos = pa.BufferOutputStream() + _write_table(arrow_table, imos) + del os.environ['ARROW_PARQUET_WRITER_ENGINE'] + @pytest.mark.pandas @parametrize_legacy_dataset From 2061d6d2f654671d6c67827caf7c2982cedd07fd Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Tue, 15 Sep 2020 05:44:33 +0000 Subject: [PATCH 12/23] first round feedback --- cpp/src/arrow/util/simd.h | 1 + cpp/src/parquet/arrow/reader.cc | 9 +- .../arrow/reconstruct_internal_test.cc | 2 +- cpp/src/parquet/column_reader.cc | 8 +- cpp/src/parquet/level_comparison.cc | 13 +- cpp/src/parquet/level_comparison.h | 53 -------- cpp/src/parquet/level_comparison_avx2.cc | 7 +- cpp/src/parquet/level_conversion.cc | 118 ++++++++++-------- cpp/src/parquet/level_conversion.h | 89 +++++++------ cpp/src/parquet/level_conversion_benchmark.cc | 8 +- cpp/src/parquet/level_conversion_bmi2.cc | 15 ++- cpp/src/parquet/level_conversion_inc.h | 35 +++--- cpp/src/parquet/level_conversion_test.cc | 92 +++++++------- 13 files changed, 218 insertions(+), 232 deletions(-) diff --git a/cpp/src/arrow/util/simd.h b/cpp/src/arrow/util/simd.h index 84c93a825cf..be3c41f631d 100644 --- a/cpp/src/arrow/util/simd.h +++ b/cpp/src/arrow/util/simd.h @@ -31,6 +31,7 @@ #if defined(ARROW_HAVE_AVX2) || defined(ARROW_HAVE_AVX512) #include +#include #elif defined(ARROW_HAVE_SSE4_2) #include #endif diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index a954dcf2a1a..02e37903db5 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -513,8 +513,8 @@ class ListReader : public ColumnReaderImpl { IndexType* length_data = reinterpret_cast(lengths_buffer->mutable_data()); std::fill(length_data, length_data + 2, 0); BEGIN_PARQUET_CATCH_EXCEPTIONS - ::parquet::internal::ConvertDefRepLevelsToList( - def_levels, rep_levels, num_levels, level_info_, &validity_io, length_data); + ::parquet::internal::DefRepLevelsToList(def_levels, rep_levels, num_levels, + level_info_, &validity_io, length_data); END_PARQUET_CATCH_EXCEPTIONS // We might return less then the requested slot (i.e. reaching an end of a file) // ensure we've set all the bits here. @@ -665,13 +665,12 @@ Status StructReader::BuildArray(int64_t records_to_read, validity_io.valid_bits = null_bitmap->mutable_data(); RETURN_NOT_OK(GetDefLevels(&def_levels, &num_levels)); RETURN_NOT_OK(GetRepLevels(&rep_levels, &num_levels)); - ConvertDefRepLevelsToBitmap(def_levels, rep_levels, num_levels, level_info_, - &validity_io); + DefRepLevelsToBitmap(def_levels, rep_levels, num_levels, level_info_, &validity_io); } else if (filtered_field_->nullable()) { ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); validity_io.valid_bits = null_bitmap->mutable_data(); RETURN_NOT_OK(GetDefLevels(&def_levels, &num_levels)); - DefinitionLevelsToBitmap(def_levels, num_levels, level_info_, &validity_io); + DefLevelsToBitmap(def_levels, num_levels, level_info_, &validity_io); } // Ensure all values are initialized. diff --git a/cpp/src/parquet/arrow/reconstruct_internal_test.cc b/cpp/src/parquet/arrow/reconstruct_internal_test.cc index f3850244ceb..495b69f9eab 100644 --- a/cpp/src/parquet/arrow/reconstruct_internal_test.cc +++ b/cpp/src/parquet/arrow/reconstruct_internal_test.cc @@ -878,7 +878,7 @@ TEST_F(TestReconstructColumn, TwoLevelListOptional) { TEST_F(TestReconstructColumn, NestedList1) { // Arrow schema: struct(a: list(int32 not null) not null) not null SetParquetSchema(GroupNode::Make( - "a", Repetition::REQUIRED, // this + "a", Repetition::REQUIRED, {GroupNode::Make( "p", Repetition::REQUIRED, {GroupNode::Make("list", Repetition::REPEATED, diff --git a/cpp/src/parquet/column_reader.cc b/cpp/src/parquet/column_reader.cc index 94f4851ed63..2d64422fa4b 100644 --- a/cpp/src/parquet/column_reader.cc +++ b/cpp/src/parquet/column_reader.cc @@ -937,7 +937,7 @@ int64_t TypedColumnReaderImpl::ReadBatchSpaced( validity_io.null_count = null_count; validity_io.values_read = *values_read; - internal::DefinitionLevelsToBitmap(def_levels, num_def_levels, info, &validity_io); + internal::DefLevelsToBitmap(def_levels, num_def_levels, info, &validity_io); null_count = validity_io.null_count; *values_read = validity_io.values_read; @@ -1383,9 +1383,9 @@ class TypedRecordReader : public ColumnReaderImplBase, validity_io.valid_bits = valid_bits_->mutable_data(); validity_io.valid_bits_offset = values_written_; - DefinitionLevelsToBitmap(def_levels() + start_levels_position, - levels_position_ - start_levels_position, leaf_info_, - &validity_io); + DefLevelsToBitmap(def_levels() + start_levels_position, + levels_position_ - start_levels_position, leaf_info_, + &validity_io); values_to_read = validity_io.values_read - validity_io.null_count; null_count = validity_io.null_count; DCHECK_GE(values_to_read, 0); diff --git a/cpp/src/parquet/level_comparison.cc b/cpp/src/parquet/level_comparison.cc index 2b68fef42c1..30614ae61fb 100644 --- a/cpp/src/parquet/level_comparison.cc +++ b/cpp/src/parquet/level_comparison.cc @@ -15,20 +15,31 @@ // specific language governing permissions and limitations // under the License. -#define IMPL_NAMESPACE standard #include "parquet/level_comparison.h" +#define PARQUET_IMPL_NAMESPACE standard +#include "parquet/level_comparison_inc.h" +#undef PARQUET_IMPL_NAMESPACE + #include #include "arrow/util/dispatch.h" namespace parquet { namespace internal { + +#if defined(ARROW_HAVE_RUNTIME_AVX2) +MinMax FindMinMaxAvx2(const int16_t* levels, int64_t num_levels); +uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_t rhs); +#endif + namespace { using ::arrow::internal::DispatchLevel; using ::arrow::internal::DynamicDispatch; +// defined in level_comparison_avx2.cc + struct GreaterThanDynamicFunction { using FunctionType = decltype(&GreaterThanBitmap); diff --git a/cpp/src/parquet/level_comparison.h b/cpp/src/parquet/level_comparison.h index af9d05e9a69..38e7ef8e2ec 100644 --- a/cpp/src/parquet/level_comparison.h +++ b/cpp/src/parquet/level_comparison.h @@ -19,46 +19,16 @@ #include #include -#include "arrow/util/bit_util.h" #include "parquet/platform.h" namespace parquet { namespace internal { -// These APIs are likely to be revised as part of ARROW-8494 to reduce duplicate code. -// They currently represent minimal functionality for vectorized computation of definition -// levels. - -/// Builds a bitmap by applying predicate to the level vector provided. -/// -/// \param[in] levels Rep or def level array. -/// \param[in] num_levels The number of levels to process (must be [0, 64]) -/// \param[in] predicate The predicate to apply (must have the signature `bool -/// predicate(int16_t)`. -/// \returns The bitmap using least significant "bit" ordering. -/// -/// N.B. Correct byte ordering is dependent on little-endian architectures. -/// -template -inline uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, - Predicate predicate) { - // Both clang and GCC can vectorize this automatically with SSE4/AVX2. - uint64_t mask = 0; - for (int x = 0; x < num_levels; x++) { - mask |= static_cast(predicate(levels[x]) ? 1 : 0) << x; - } - return ::arrow::BitUtil::ToLittleEndian(mask); -} - /// Builds a bitmap where each set bit indicates the corresponding level is greater /// than rhs. uint64_t PARQUET_EXPORT GreaterThanBitmap(const int16_t* levels, int64_t num_levels, int16_t rhs); -#if defined(ARROW_HAVE_RUNTIME_AVX2) -uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_t rhs); -#endif - struct MinMax { int16_t min; int16_t max; @@ -66,28 +36,5 @@ struct MinMax { MinMax FindMinMax(const int16_t* levels, int64_t num_levels); -#if defined(ARROW_HAVE_RUNTIME_AVX2) -MinMax FindMinMaxAvx2(const int16_t* levels, int64_t num_levels); -#endif - -// Used to make sure ODR rule isn't violated. -namespace IMPL_NAMESPACE { -inline MinMax FindMinMaxImpl(const int16_t* levels, int64_t num_levels) { - MinMax out{std::numeric_limits::max(), std::numeric_limits::min()}; - for (int x = 0; x < num_levels; x++) { - out.min = std::min(levels[x], out.min); - out.max = std::max(levels[x], out.max); - } - return out; -} - -inline uint64_t GreaterThanBitmapImpl(const int16_t* levels, int64_t num_levels, - int16_t rhs) { - return LevelsToBitmap(levels, num_levels, [rhs](int16_t value) { return value > rhs; }); -} - -} // namespace IMPL_NAMESPACE - } // namespace internal - } // namespace parquet diff --git a/cpp/src/parquet/level_comparison_avx2.cc b/cpp/src/parquet/level_comparison_avx2.cc index 6cd6d0cc1f7..b33eb2e2953 100644 --- a/cpp/src/parquet/level_comparison_avx2.cc +++ b/cpp/src/parquet/level_comparison_avx2.cc @@ -15,9 +15,9 @@ // specific language governing permissions and limitations // under the License. -#define IMPL_NAMESPACE avx2 -#include "parquet/level_comparison.h" -#undef IMPL_NAMESPACE +#define PARQUET_IMPL_NAMESPACE avx2 +#include "parquet/level_comparison_inc.h" +#undef PARQUET_IMPL_NAMESPACE namespace parquet { namespace internal { @@ -29,5 +29,6 @@ uint64_t GreaterThanBitmapAvx2(const int16_t* levels, int64_t num_levels, int16_ MinMax FindMinMaxAvx2(const int16_t* levels, int64_t num_levels) { return avx2::FindMinMaxImpl(levels, num_levels); } + } // namespace internal } // namespace parquet diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index b0187fc26d4..fdeaad6e4f1 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -24,11 +24,11 @@ #include "arrow/util/cpu_info.h" #include "arrow/util/logging.h" #include "parquet/exception.h" -#include "parquet/level_comparison.h" -#define BMI_RUNTIME_VERSION standard +#define PARQUET_IMPL_NAMESPACE standard +#include "parquet/level_comparison.h" #include "parquet/level_conversion_inc.h" -#undef BMI_RUNTIME_VERSION +#undef PARQUET_IMPL_NAMESPACE namespace parquet { namespace internal { @@ -37,14 +37,15 @@ namespace { using ::arrow::internal::CpuInfo; #if !defined(ARROW_HAVE_RUNTIME_BMI2) -void DefinitionLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, - LevelInfo level_info, - ValidityBitmapInputOutput* output) { +void DefLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, ValidityBitmapInputOutput* output) { ::arrow::internal::FirstTimeBitmapWriter valid_bits_writer( output->valid_bits, /*start_offset=*/output->valid_bits_offset, /*length=*/num_def_levels); for (int x = 0; x < num_def_levels; x++) { + // This indicates that a parent repeated element has zero + // length so the def level is not applicable to this column. if (def_levels[x] < level_info.repeated_ancestor_def_level) { continue; } @@ -72,18 +73,18 @@ void DefinitionLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_l } #endif -template +template void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels, int64_t num_def_levels, LevelInfo level_info, - ValidityBitmapInputOutput* output, LengthType* lengths) { - LengthType* orig_pos = lengths; + ValidityBitmapInputOutput* output, OffsetType* offsets) { + OffsetType* orig_pos = offsets; std::unique_ptr<::arrow::internal::FirstTimeBitmapWriter> valid_bits_writer; if (output->valid_bits) { valid_bits_writer.reset(new ::arrow::internal::FirstTimeBitmapWriter( output->valid_bits, output->valid_bits_offset, num_def_levels)); } for (int x = 0; x < num_def_levels; x++) { - // Skip items that belong to empty ancenstor lists and futher nested lists. + // Skip items that belong to empty or null ancestor lists and further nested lists. if (def_levels[x] < level_info.repeated_ancestor_def_level || rep_levels[x] > level_info.rep_level) { continue; @@ -92,7 +93,7 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels if (ARROW_PREDICT_FALSE( (valid_bits_writer != nullptr && valid_bits_writer->position() > output->values_read_upper_bound) || - (lengths - orig_pos) > output->values_read_upper_bound)) { + (offsets - orig_pos) > output->values_read_upper_bound)) { std::stringstream ss; ss << "Definition levels exceeded upper bound: " << output->values_read_upper_bound; throw ParquetException(ss.str()); @@ -100,25 +101,30 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels if (rep_levels[x] == level_info.rep_level) { // A continuation of an existing list. - if (lengths != nullptr) { - if (ARROW_PREDICT_FALSE(*lengths == std::numeric_limits::max())) { + // offsets can be null for structs with repeated children (we don't need to know + // offsets until we get to the children). + if (offsets != nullptr) { + if (ARROW_PREDICT_FALSE(*offsets == std::numeric_limits::max())) { throw ParquetException("List index overflow."); } - *lengths += 1; + *offsets += 1; } } else { - // current_rep < list rep_level i.e. start of a list (ancenstor empty lists are + // current_rep < list rep_level i.e. start of a list (ancestor empty lists are // filtered out above). - if (lengths != nullptr) { - ++lengths; - // Use cumulative lengths because this is what the more common - // Arrow list types expect. - *lengths = *(lengths - 1); + // offsets can be null for structs with repeated children (we don't need to know + // offsets until we get to the children). + if (offsets != nullptr) { + ++offsets; + // Use cumulative offsets because variable size lists are more common then + // fixed size lists so it should be cheaper to make these cumulative and + // subtract when validating fixed size lists. + *offsets = *(offsets - 1); if (def_levels[x] >= level_info.def_level) { - if (ARROW_PREDICT_FALSE(*lengths == std::numeric_limits::max())) { + if (ARROW_PREDICT_FALSE(*offsets == std::numeric_limits::max())) { throw ParquetException("List index overflow."); } - *lengths += 1; + *offsets += 1; } } @@ -138,8 +144,8 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels if (valid_bits_writer != nullptr) { valid_bits_writer->Finish(); } - if (lengths != nullptr) { - output->values_read = lengths - orig_pos; + if (offsets != nullptr) { + output->values_read = offsets - orig_pos; } else if (valid_bits_writer != nullptr) { output->values_read = valid_bits_writer->position(); } @@ -152,56 +158,64 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels } // namespace -void DefinitionLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, - LevelInfo level_info, ValidityBitmapInputOutput* output) { +#if defined(ARROW_HAVE_RUNTIME_BMI2) +// defined in level_conversion_bmi2.cc for dynamic dispatch. +void DefLevelsToBitmapBmi2WithRepeatedParent(const int16_t* def_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output); +#endif + +void DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, ValidityBitmapInputOutput* output) { + // It is simpler to rely on rep_level here until PARQUET-1899 is done and the code + // is deleted in a follow-up release. if (level_info.rep_level > 0) { #if defined(ARROW_HAVE_RUNTIME_BMI2) - using FunctionType = decltype(&standard::DefinitionLevelsToBitmapSimd); + using FunctionType = decltype(&standard::DefLevelsToBitmapSimd); static FunctionType fn = CpuInfo::GetInstance()->HasEfficientBmi2() - ? DefinitionLevelsToBitmapBmi2WithRepeatedParent - : standard::DefinitionLevelsToBitmapSimd; + ? DefLevelsToBitmapBmi2WithRepeatedParent + : standard::DefLevelsToBitmapSimd; fn(def_levels, num_def_levels, level_info, output); #else - DefinitionLevelsToBitmapScalar(def_levels, num_def_levels, level_info, output); - + // This indicates we are likely on a big-endian platformat which don't have a + // pext function and the current implementation of BitRunReader is always linear + // in bits, so this method will likely be the fastest alternative. + DefLevelsToBitmapScalar(def_levels, num_def_levels, level_info, output); #endif } else { - standard::DefinitionLevelsToBitmapSimd( + standard::DefLevelsToBitmapSimd( def_levels, num_def_levels, level_info, output); } } -uint64_t RunBasedExtract(uint64_t bitmap, uint64_t select_bitmap) { +uint64_t TestOnlyRunBasedExtract(uint64_t bitmap, uint64_t select_bitmap) { return standard::RunBasedExtractImpl(bitmap, select_bitmap); } -void ConvertDefRepLevelsToList(const int16_t* def_levels, const int16_t* rep_levels, - int64_t num_def_levels, LevelInfo level_info, - ValidityBitmapInputOutput* output, - ::arrow::util::variant lengths) { - if (arrow::util::holds_alternative(lengths)) { - auto int32_lengths = ::arrow::util::get(lengths); - DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, - output, int32_lengths); - } else if (arrow::util::holds_alternative(lengths)) { - auto int64_lengths = ::arrow::util::get(lengths); - DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, - output, int64_lengths); - } else { - throw ParquetException("Unrecognized variant"); - } +void DefRepLevelsToList(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output, int32_t* offsets) { + DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, + output, offsets); +} + +void DefRepLevelsToList(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output, int64_t* offsets) { + DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, + output, offsets); } -void ConvertDefRepLevelsToBitmap(const int16_t* def_levels, const int16_t* rep_levels, - int64_t num_def_levels, LevelInfo level_info, - ValidityBitmapInputOutput* output) { +void DefRepLevelsToBitmap(const int16_t* def_levels, const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output) { // DefReplevelsToListInfo assumes it for the actual list method and this // method is for parent structs, so we need to bump def and ref level. level_info.rep_level += 1; level_info.def_level += 1; DefRepLevelsToListInfo(def_levels, rep_levels, num_def_levels, level_info, - output, /*lengths=*/nullptr); + output, /*offsets=*/nullptr); } } // namespace internal diff --git a/cpp/src/parquet/level_conversion.h b/cpp/src/parquet/level_conversion.h index df7a68729c6..570f7dde822 100644 --- a/cpp/src/parquet/level_conversion.h +++ b/cpp/src/parquet/level_conversion.h @@ -20,8 +20,6 @@ #include #include "arrow/util/bitmap.h" -#include "arrow/util/optional.h" -#include "arrow/util/variant.h" #include "parquet/platform.h" #include "parquet/schema.h" @@ -137,50 +135,65 @@ struct PARQUET_EXPORT LevelInfo { } }; -/// Input/Output structure for reconstructed validity bitmaps. +// Input/Output structure for reconstructed validity bitmaps. struct PARQUET_EXPORT ValidityBitmapInputOutput { - /// The maximum number of values_read expected (actual - /// values read must be less than or equal to this value. - /// If this number is exceeded methods will throw a - /// ParquetException. + // Input only. + // The maximum number of values_read expected (actual + // values read must be less than or equal to this value. + // If this number is exceeded methods will throw a + // ParquetException. Exceeding this limit indicates + // either a corrupt or incorrectly written file. int64_t values_read_upper_bound = 0; - /// The number of values added to the bitmap. + // Output only. The number of values added to the encountered + // (this is logicallyt he count of the number of elements + // for an Arrow array). int64_t values_read = 0; - /// The number of nulls encountered. + // Input/Output. The number of nulls encountered. int64_t null_count = 0; - // The validity bitmp to populate. Can only be null - // for DefRepLevelsToListInfo (if all that is needed is list lengths). + // Output only. The validity bitmap to populate. May be be null only + // for DefRepLevelsToListInfo (if all that is needed is list offsets). uint8_t* valid_bits = NULLPTR; - /// Input only, offset into valid_bits to start at. + // Input only, offset into valid_bits to start at. int64_t valid_bits_offset = 0; }; -/// Converts def_levels to validity bitmaps for non-list arrays. -void PARQUET_EXPORT DefinitionLevelsToBitmap(const int16_t* def_levels, - int64_t num_def_levels, LevelInfo level_info, - ValidityBitmapInputOutput* output); - -/// Reconstructs a validity bitmap and list lengths for a ListArray based on -/// def/rep levels. -void PARQUET_EXPORT ConvertDefRepLevelsToList( - const int16_t* def_levels, const int16_t* rep_levels, int64_t num_def_levels, - LevelInfo level_info, ValidityBitmapInputOutput* output, - ::arrow::util::variant lengths); - -/// Reconstructs a validity bitmap for a struct that has nested children. -void PARQUET_EXPORT ConvertDefRepLevelsToBitmap(const int16_t* def_levels, - const int16_t* rep_levels, - int64_t num_def_levels, - LevelInfo level_info, - ValidityBitmapInputOutput* output); - -uint64_t PARQUET_EXPORT RunBasedExtract(uint64_t bitmap, uint64_t selection); - -#if defined(ARROW_HAVE_RUNTIME_BMI2) -void PARQUET_EXPORT DefinitionLevelsToBitmapBmi2WithRepeatedParent( - const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, - ValidityBitmapInputOutput* output); -#endif +// Converts def_levels to validity bitmaps for non-list arrays and structs that have +// at least one member that is not a list and has no list descendents. +// For lists use DefRepLevelsToList and structs where all descendants contain +// a list use DefRepLevelsToBitmap. +void PARQUET_EXPORT DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output); + +// Reconstructs a validity bitmap and list offsets for a list arrays based on +// def/rep levels. The first element of offsets will not be modified if rep_levels +// starts with a new list. The first element of offsets will be used when calculating +// the next offset. See documentation onf DefLevelsToBitmap for when to use this +// method vs the other ones in this file for reconstruction. +// +// Offsets must be size to 1 + values_read_upper_bound. +void PARQUET_EXPORT DefRepLevelsToList(const int16_t* def_levels, + const int16_t* rep_levels, int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output, + int32_t* offsets); +void PARQUET_EXPORT DefRepLevelsToList(const int16_t* def_levels, + const int16_t* rep_levels, int64_t num_def_levels, + LevelInfo level_info, + ValidityBitmapInputOutput* output, + int64_t* offsets); + +// Reconstructs a validity bitmap for a struct every member is a list or has +// a list descendant. See documentation on DefLevelsToBitmap for when more +// details on this method compared to the other ones defined above. +void PARQUET_EXPORT DefRepLevelsToBitmap(const int16_t* def_levels, + const int16_t* rep_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output); + +// This is exposed to ensure we can properly test a software simulated pext function +// (i.e. it isn't hidden by runtime dispatch). +uint64_t PARQUET_EXPORT TestOnlyRunBasedExtract(uint64_t bitmap, uint64_t selection); } // namespace internal } // namespace parquet diff --git a/cpp/src/parquet/level_conversion_benchmark.cc b/cpp/src/parquet/level_conversion_benchmark.cc index dccc922b53c..953d668b78c 100644 --- a/cpp/src/parquet/level_conversion_benchmark.cc +++ b/cpp/src/parquet/level_conversion_benchmark.cc @@ -41,10 +41,12 @@ std::vector RunDefinitionLevelsToBitmap(const std::vector& def parquet::internal::LevelInfo info; info.def_level = kHasRepeatedElements; info.repeated_ancestor_def_level = kPresentDefLevel; + parquet::internal::ValidityBitmapInputOutput validity_io; + io.values_read_upper_bound = def_levels.size(); + io.valid_bits = bitmap.data(); for (auto _ : *state) { - parquet::internal::DefinitionLevelsToBitmap( - def_levels.data(), def_levels.size(), info, &values_read, &null_count, - bitmap.data(), /*valid_bits_offset=*/(rep++ % 8) * def_levels.size()); + parquet::internal::DefinitionLevelsToBitmap(def_levels.data(), def_levels.size(), + info, &validity_io); } state->SetBytesProcessed(int64_t(state->iterations()) * def_levels.size()); return bitmap; diff --git a/cpp/src/parquet/level_conversion_bmi2.cc b/cpp/src/parquet/level_conversion_bmi2.cc index fa291091a30..274d54e503c 100644 --- a/cpp/src/parquet/level_conversion_bmi2.cc +++ b/cpp/src/parquet/level_conversion_bmi2.cc @@ -16,18 +16,17 @@ // under the License. #include "parquet/level_conversion.h" -#define BMI_RUNTIME_VERSION bmi2 +#define PARQUET_IMPL_NAMESPACE bmi2 #include "parquet/level_conversion_inc.h" -#undef BMI_RUNTIME_VERSION +#undef PARQUET_IMPL_NAMESPACE namespace parquet { namespace internal { -void DefinitionLevelsToBitmapBmi2WithRepeatedParent(const int16_t* def_levels, - int64_t num_def_levels, - LevelInfo level_info, - ValidityBitmapInputOutput* output) { - bmi2::DefinitionLevelsToBitmapSimd( - def_levels, num_def_levels, level_info, output); +void DefLevelsToBitmapBmi2WithRepeatedParent(const int16_t* def_levels, + int64_t num_def_levels, LevelInfo level_info, + ValidityBitmapInputOutput* output) { + bmi2::DefLevelsToBitmapSimd(def_levels, num_def_levels, + level_info, output); } } // namespace internal diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h index b24fe64fa21..4c24004bbf1 100644 --- a/cpp/src/parquet/level_conversion_inc.h +++ b/cpp/src/parquet/level_conversion_inc.h @@ -20,26 +20,20 @@ #include #include -#if defined(ARROW_HAVE_BMI2) -#if defined(_MSC_VER) -#include -#else -#include -#endif // _MSC_VER -#endif // ARROW_HAVE_BMI2 #include "arrow/util/bit_run_reader.h" #include "arrow/util/bit_util.h" #include "arrow/util/logging.h" +#include "arrow/util/simd.h" #include "parquet/exception.h" #include "parquet/level_comparison.h" namespace parquet { namespace internal { -namespace BMI_RUNTIME_VERSION { - -using ::arrow::internal::BitRun; -using ::arrow::internal::BitRunReader; +#ifndef PARQUET_IMPL_NAMESPACE +#error "PARQUET_IMPL_NAMESPACE must be defined" +#endif +namespace PARQUET_IMPL_NAMESPACE { /// Algorithm to simulate pext using BitRunReader for cases where all bits /// not set or set. @@ -82,9 +76,9 @@ inline uint64_t ExtractBits(uint64_t bitmap, uint64_t select_bitmap) { } template -int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_size, - int64_t upper_bound_remaining, LevelInfo level_info, - ::arrow::internal::FirstTimeBitmapWriter* writer) { +int64_t DefLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_size, + int64_t upper_bound_remaining, LevelInfo level_info, + ::arrow::internal::FirstTimeBitmapWriter* writer) { // Greater than level_info.def_level - 1 implies >= the def_level uint64_t defined_bitmap = internal::GreaterThanBitmap(def_levels, batch_size, level_info.def_level - 1); @@ -92,7 +86,7 @@ int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t b DCHECK_LE(batch_size, 64); if (has_repeated_parent) { // Greater than level_info.repeated_ancestor_def_level - 1 implies >= the - // repeated_ancenstor_def_level + // repeated_ancestor_def_level uint64_t present_bitmap = internal::GreaterThanBitmap( def_levels, batch_size, level_info.repeated_ancestor_def_level - 1); uint64_t selected_bits = ExtractBits(defined_bitmap, present_bitmap); @@ -115,9 +109,8 @@ int64_t DefinitionLevelsBatchToBitmap(const int16_t* def_levels, const int64_t b } template -void DefinitionLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_levels, - LevelInfo level_info, - ValidityBitmapInputOutput* output) { +void DefLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_levels, + LevelInfo level_info, ValidityBitmapInputOutput* output) { constexpr int64_t kBitMaskSize = 64; ::arrow::internal::FirstTimeBitmapWriter writer( output->valid_bits, @@ -127,13 +120,13 @@ void DefinitionLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_lev output->values_read = 0; int64_t values_read_remaining = output->values_read_upper_bound; while (num_def_levels > kBitMaskSize) { - set_count += DefinitionLevelsBatchToBitmap( + set_count += DefLevelsBatchToBitmap( def_levels, kBitMaskSize, values_read_remaining, level_info, &writer); def_levels += kBitMaskSize; num_def_levels -= kBitMaskSize; values_read_remaining = output->values_read_upper_bound - writer.position(); } - set_count += DefinitionLevelsBatchToBitmap( + set_count += DefLevelsBatchToBitmap( def_levels, num_def_levels, values_read_remaining, level_info, &writer); output->values_read = writer.position(); @@ -141,6 +134,6 @@ void DefinitionLevelsToBitmapSimd(const int16_t* def_levels, int64_t num_def_lev writer.Finish(); } -} // namespace BMI_RUNTIME_VERSION +} // namespace PARQUET_IMPL_NAMESPACE } // namespace internal } // namespace parquet diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index 36bf1a6ec7e..bb8c5cbfcc2 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -16,7 +16,10 @@ // under the License. #include "parquet/level_conversion.h" + +#define PARQUET_IMPL_NAMESPACE test #include "parquet/level_comparison.h" +#undef PARQUET_IMPL_NAMESPACE #include "parquet/test_util.h" #include @@ -43,7 +46,7 @@ std::string BitmapToString(const std::vector& bitmap, int64_t bit_count return BitmapToString(bitmap.data(), bit_count); } -TEST(TestColumnReader, DefinitionLevelsToBitmap) { +TEST(TestColumnReader, DefLevelsToBitmap) { // Bugs in this function were exposed in ARROW-3930 std::vector def_levels = {3, 3, 3, 2, 3, 3, 3, 3, 3}; @@ -58,21 +61,21 @@ TEST(TestColumnReader, DefinitionLevelsToBitmap) { io.values_read = -1; io.valid_bits = valid_bits.data(); - internal::DefinitionLevelsToBitmap(def_levels.data(), 9, level_info, &io); + internal::DefLevelsToBitmap(def_levels.data(), 9, level_info, &io); ASSERT_EQ(9, io.values_read); ASSERT_EQ(1, io.null_count); // Call again with 0 definition levels, make sure that valid_bits is unmodified const uint8_t current_byte = valid_bits[1]; io.null_count = 0; - internal::DefinitionLevelsToBitmap(def_levels.data(), 0, level_info, &io); + internal::DefLevelsToBitmap(def_levels.data(), 0, level_info, &io); ASSERT_EQ(0, io.values_read); ASSERT_EQ(0, io.null_count); ASSERT_EQ(current_byte, valid_bits[1]); } -TEST(TestColumnReader, DefinitionLevelsToBitmapPowerOfTwo) { +TEST(TestColumnReader, DefLevelsToBitmapPowerOfTwo) { // PARQUET-1623: Invalid memory access when decoding a valid bits vector that has a // length equal to a power of two and also using a non-zero valid_bits_offset. This // should not fail when run with ASAN or valgrind. @@ -89,7 +92,7 @@ TEST(TestColumnReader, DefinitionLevelsToBitmapPowerOfTwo) { io.valid_bits = valid_bits.data(); // Read the latter half of the validity bitmap - internal::DefinitionLevelsToBitmap(def_levels.data() + 4, 4, level_info, &io); + internal::DefLevelsToBitmap(def_levels.data() + 4, 4, level_info, &io); ASSERT_EQ(4, io.values_read); ASSERT_EQ(0, io.null_count); } @@ -112,7 +115,7 @@ TEST(GreaterThanBitmap, GeneratesExpectedBitmasks) { } #endif -TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { +TEST(DefLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { std::vector validity_bitmap(/*count*/ 8, 0); ValidityBitmapInputOutput io; @@ -128,7 +131,7 @@ TEST(DefinitionLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { level_info.rep_level = 1; // All zeros should be ignored, ones should be unset in the bitmp and 2 should be set. std::vector def_levels = {0, 0, 0, 2, 2, 1, 0, 2}; - DefinitionLevelsToBitmap(def_levels.data(), def_levels.size(), level_info, &io); + DefLevelsToBitmap(def_levels.data(), def_levels.size(), level_info, &io); EXPECT_EQ(BitmapToString(validity_bitmap, /*bit_count=*/8), "01101000"); for (size_t x = 1; x < validity_bitmap.size(); x++) { @@ -166,10 +169,10 @@ struct RepDefLevelConverter { using ListLengthType = ListType; ListLengthType* ComputeListInfo(const MultiLevelTestData& test_data, LevelInfo level_info, ValidityBitmapInputOutput* output, - ListType* lengths) { - ConvertDefRepLevelsToList(test_data.def_levels_.data(), test_data.rep_levels_.data(), - test_data.def_levels_.size(), level_info, output, lengths); - return lengths + output->values_read; + ListType* offsets) { + DefRepLevelsToList(test_data.def_levels_.data(), test_data.rep_levels_.data(), + test_data.def_levels_.size(), level_info, output, offsets); + return offsets + output->values_read; } }; @@ -188,16 +191,16 @@ TYPED_TEST(NestedListTest, OuterMostTest) { level_info.rep_level = 1; level_info.def_level = 2; - std::vector lengths(5, 0); + std::vector offsets(5, 0); uint64_t validity_output; ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 4; validity_io.valid_bits = reinterpret_cast(&validity_output); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, lengths.data()); + this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, lengths.data() + 4); - EXPECT_THAT(lengths, testing::ElementsAre(0, 3, 7, 7, 7)); + EXPECT_THAT(next_position, offsets.data() + 4); + EXPECT_THAT(offsets, testing::ElementsAre(0, 3, 7, 7, 7)); EXPECT_EQ(validity_io.values_read, 4); EXPECT_EQ(validity_io.null_count, 1); @@ -218,16 +221,16 @@ TYPED_TEST(NestedListTest, MiddleListTest) { level_info.def_level = 4; level_info.repeated_ancestor_def_level = 2; - std::vector lengths(8, 0); + std::vector offsets(8, 0); uint64_t validity_output; ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 7; validity_io.valid_bits = reinterpret_cast(&validity_output); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, lengths.data()); + this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, lengths.data() + 7); - EXPECT_THAT(lengths, testing::ElementsAre(0, 0, 2, 2, 3, 5, 5, 6)); + EXPECT_THAT(next_position, offsets.data() + 7); + EXPECT_THAT(offsets, testing::ElementsAre(0, 0, 2, 2, 3, 5, 5, 6)); EXPECT_EQ(validity_io.values_read, 7); EXPECT_EQ(validity_io.null_count, 2); @@ -239,7 +242,7 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { // [[[]], [[], [1, 2]], null, [[3]]], // null, // [] - // -> 4 inner lists (N/A, [len(3), len(0)], N/A + // -> 6 inner lists (N/A, [len(3), len(0)], N/A // len(0), [len(0), len(2)], N/A, len(1), // N/A, // N/A @@ -248,16 +251,16 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { level_info.def_level = 6; level_info.repeated_ancestor_def_level = 4; - std::vector lengths(7, 0); + std::vector offsets(7, 0); uint64_t validity_output; ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 6; validity_io.valid_bits = reinterpret_cast(&validity_output); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, lengths.data()); + this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, lengths.data() + 6); - EXPECT_THAT(lengths, testing::ElementsAre(0, 3, 3, 3, 3, 5, 6)); + EXPECT_THAT(next_position, offsets.data() + 6); + EXPECT_THAT(offsets, testing::ElementsAre(0, 3, 3, 3, 3, 5, 6)); EXPECT_EQ(validity_io.values_read, 6); EXPECT_EQ(validity_io.null_count, 0); @@ -279,20 +282,20 @@ TYPED_TEST(NestedListTest, SimpleLongList) { /*rep_level=*/1); } - std::vector lengths(66, 0); - std::vector expected_lengths(66, 0); - for (size_t x = 1; x < expected_lengths.size(); x++) { - expected_lengths[x] = x * 9; + std::vector offsets(66, 0); + std::vector expected_offsets(66, 0); + for (size_t x = 1; x < expected_offsets.size(); x++) { + expected_offsets[x] = x * 9; } std::vector validity_output(9, 0); ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = 65; validity_io.valid_bits = validity_output.data(); typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, lengths.data()); + this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, lengths.data() + 65); - EXPECT_THAT(lengths, testing::ElementsAreArray(expected_lengths)); + EXPECT_THAT(next_position, offsets.data() + 65); + EXPECT_THAT(offsets, testing::ElementsAreArray(expected_offsets)); EXPECT_EQ(validity_io.values_read, 65); EXPECT_EQ(validity_io.null_count, 0); @@ -318,7 +321,7 @@ TYPED_TEST(NestedListTest, TestOverflow) { this->test_data_.def_levels_ = std::vector{2}; this->test_data_.rep_levels_ = std::vector{0}; - std::vector lengths( + std::vector offsets( 2, std::numeric_limits::max()); std::vector validity_output(1, 0); @@ -326,32 +329,35 @@ TYPED_TEST(NestedListTest, TestOverflow) { validity_io.values_read_upper_bound = 1; validity_io.valid_bits = validity_output.data(); ASSERT_THROW(this->converter_.ComputeListInfo(this->test_data_, level_info, - &validity_io, lengths.data()), + &validity_io, offsets.data()), ParquetException); // Same thing should happen if the list already existed. this->test_data_.rep_levels_ = std::vector{1}; ASSERT_THROW(this->converter_.ComputeListInfo(this->test_data_, level_info, - &validity_io, lengths.data()), + &validity_io, offsets.data()), ParquetException); // Should be OK because it shouldn't increment. this->test_data_.def_levels_ = std::vector{0}; + this->test_data_.rep_levels_ = std::vector{0}; + this->converter_.ComputeListInfo(this->test_data_, level_info, &validity_io, + offsets.data()); } -TEST(RunBasedExtract, BasicTest) { - EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), 0), 0); - EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), ~uint64_t{0}), +TEST(TestOnlyRunBasedExtract, BasicTest) { + EXPECT_EQ(TestOnlyRunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), 0), 0); + EXPECT_EQ(TestOnlyRunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF), ~uint64_t{0}), arrow::BitUtil::ToLittleEndian(0xFF)); - EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF00FF), - arrow::BitUtil::ToLittleEndian(0xAAAA)), + EXPECT_EQ(TestOnlyRunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF00FF), + arrow::BitUtil::ToLittleEndian(0xAAAA)), arrow::BitUtil::ToLittleEndian(0x000F)); - EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF0AFF), - arrow::BitUtil::ToLittleEndian(0xAFAA)), + EXPECT_EQ(TestOnlyRunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFF0AFF), + arrow::BitUtil::ToLittleEndian(0xAFAA)), arrow::BitUtil::ToLittleEndian(0x00AF)); - EXPECT_EQ(RunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFFAAFF), - arrow::BitUtil::ToLittleEndian(0xAFAA)), + EXPECT_EQ(TestOnlyRunBasedExtract(arrow::BitUtil::ToLittleEndian(0xFFAAFF), + arrow::BitUtil::ToLittleEndian(0xAFAA)), arrow::BitUtil::ToLittleEndian(0x03AF)); } From 73f6aeefc2135df2700f5d41687731c928e629a6 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Tue, 15 Sep 2020 05:45:49 +0000 Subject: [PATCH 13/23] remove comment --- cpp/src/parquet/arrow/arrow_reader_writer_test.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index a23f6e559a2..188bc8c178a 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -2360,7 +2360,6 @@ TEST(ArrowReadWrite, SingleColumnNullableStruct) { 3); } -// Disabled until implementation can be finished. TEST(TestArrowReadWrite, CanonicalNestedRoundTrip) { auto doc_id = field("DocId", ::arrow::int64(), /*nullable=*/false); auto links = field( From d0667e44e4af0ffb334bb7b3a2a006f27962f25c Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Tue, 15 Sep 2020 06:17:10 +0000 Subject: [PATCH 14/23] more comments --- cpp/src/arrow/util/simd.h | 5 +- cpp/src/parquet/level_comparison_inc.h | 64 +++++++++++++++++++ cpp/src/parquet/level_conversion.cc | 23 +++---- cpp/src/parquet/level_conversion_benchmark.cc | 11 ++-- cpp/src/parquet/level_conversion_test.cc | 2 - 5 files changed, 84 insertions(+), 21 deletions(-) create mode 100644 cpp/src/parquet/level_comparison_inc.h diff --git a/cpp/src/arrow/util/simd.h b/cpp/src/arrow/util/simd.h index be3c41f631d..259641dd456 100644 --- a/cpp/src/arrow/util/simd.h +++ b/cpp/src/arrow/util/simd.h @@ -29,9 +29,12 @@ #else // gcc/clang (possibly others) +#if defined(ARROW_HAVE_BMI2) +#include +#endif + #if defined(ARROW_HAVE_AVX2) || defined(ARROW_HAVE_AVX512) #include -#include #elif defined(ARROW_HAVE_SSE4_2) #include #endif diff --git a/cpp/src/parquet/level_comparison_inc.h b/cpp/src/parquet/level_comparison_inc.h new file mode 100644 index 00000000000..f4cf7ab48e7 --- /dev/null +++ b/cpp/src/parquet/level_comparison_inc.h @@ -0,0 +1,64 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. +#pragma once + +#include "arrow/util/bit_util.h" +#include "parquet/level_comparison.h" + +// Used to make sure ODR rule isn't violated. +#ifndef PARQUET_IMPL_NAMESPACE +#error "PARQUET_IMPL_NAMESPACE must be defined" +#endif +namespace parquet { +namespace internal { +namespace PARQUET_IMPL_NAMESPACE { +/// Builds a bitmap by applying predicate to the level vector provided. +/// +/// \param[in] levels Rep or def level array. +/// \param[in] num_levels The number of levels to process (must be [0, 64]) +/// \param[in] predicate The predicate to apply (must have the signature `bool +/// predicate(int16_t)`. +/// \returns The bitmap using least significant "bit" ordering. +/// +template +inline uint64_t LevelsToBitmap(const int16_t* levels, int64_t num_levels, + Predicate predicate) { + // Both clang and GCC can vectorize this automatically with SSE4/AVX2. + uint64_t mask = 0; + for (int x = 0; x < num_levels; x++) { + mask |= static_cast(predicate(levels[x]) ? 1 : 0) << x; + } + return ::arrow::BitUtil::ToLittleEndian(mask); +} + +inline MinMax FindMinMaxImpl(const int16_t* levels, int64_t num_levels) { + MinMax out{std::numeric_limits::max(), std::numeric_limits::min()}; + for (int x = 0; x < num_levels; x++) { + out.min = std::min(levels[x], out.min); + out.max = std::max(levels[x], out.max); + } + return out; +} + +inline uint64_t GreaterThanBitmapImpl(const int16_t* levels, int64_t num_levels, + int16_t rhs) { + return LevelsToBitmap(levels, num_levels, [rhs](int16_t value) { return value > rhs; }); +} + +} // namespace PARQUET_IMPL_NAMESPACE +} // namespace internal +} // namespace parquet diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index fdeaad6e4f1..9868f58ec4a 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -25,8 +25,8 @@ #include "arrow/util/logging.h" #include "parquet/exception.h" -#define PARQUET_IMPL_NAMESPACE standard #include "parquet/level_comparison.h" +#define PARQUET_IMPL_NAMESPACE standard #include "parquet/level_conversion_inc.h" #undef PARQUET_IMPL_NAMESPACE @@ -49,7 +49,7 @@ void DefLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, if (def_levels[x] < level_info.repeated_ancestor_def_level) { continue; } - if (ARROW_PREDICT_FALSE(valid_bits_writer.position() > + if (ARROW_PREDICT_FALSE(valid_bits_writer.position() >= output->values_read_upper_bound)) { std::stringstream ss; ss << "Definition levels exceeded upper bound: " << output->values_read_upper_bound; @@ -90,15 +90,6 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels continue; } - if (ARROW_PREDICT_FALSE( - (valid_bits_writer != nullptr && - valid_bits_writer->position() > output->values_read_upper_bound) || - (offsets - orig_pos) > output->values_read_upper_bound)) { - std::stringstream ss; - ss << "Definition levels exceeded upper bound: " << output->values_read_upper_bound; - throw ParquetException(ss.str()); - } - if (rep_levels[x] == level_info.rep_level) { // A continuation of an existing list. // offsets can be null for structs with repeated children (we don't need to know @@ -110,6 +101,16 @@ void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels *offsets += 1; } } else { + if (ARROW_PREDICT_FALSE( + (valid_bits_writer != nullptr && + valid_bits_writer->position() >= output->values_read_upper_bound) || + (offsets - orig_pos) >= output->values_read_upper_bound)) { + std::stringstream ss; + ss << "Definition levels exceeded upper bound: " + << output->values_read_upper_bound; + throw ParquetException(ss.str()); + } + // current_rep < list rep_level i.e. start of a list (ancestor empty lists are // filtered out above). // offsets can be null for structs with repeated children (we don't need to know diff --git a/cpp/src/parquet/level_conversion_benchmark.cc b/cpp/src/parquet/level_conversion_benchmark.cc index 953d668b78c..759b68fac58 100644 --- a/cpp/src/parquet/level_conversion_benchmark.cc +++ b/cpp/src/parquet/level_conversion_benchmark.cc @@ -34,19 +34,16 @@ constexpr int16_t kHasRepeatedElements = 1; std::vector RunDefinitionLevelsToBitmap(const std::vector& def_levels, ::benchmark::State* state) { - int64_t values_read = 0; - int64_t null_count = 0; std::vector bitmap(/*count=*/def_levels.size(), 0); - int rep = 0; parquet::internal::LevelInfo info; info.def_level = kHasRepeatedElements; info.repeated_ancestor_def_level = kPresentDefLevel; parquet::internal::ValidityBitmapInputOutput validity_io; - io.values_read_upper_bound = def_levels.size(); - io.valid_bits = bitmap.data(); + validity_io.values_read_upper_bound = def_levels.size(); + validity_io.valid_bits = bitmap.data(); for (auto _ : *state) { - parquet::internal::DefinitionLevelsToBitmap(def_levels.data(), def_levels.size(), - info, &validity_io); + parquet::internal::DefLevelsToBitmap(def_levels.data(), def_levels.size(), info, + &validity_io); } state->SetBytesProcessed(int64_t(state->iterations()) * def_levels.size()); return bitmap; diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index bb8c5cbfcc2..1ae3f91413c 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -17,9 +17,7 @@ #include "parquet/level_conversion.h" -#define PARQUET_IMPL_NAMESPACE test #include "parquet/level_comparison.h" -#undef PARQUET_IMPL_NAMESPACE #include "parquet/test_util.h" #include From d3c0a8ae8838efbf9426391de182b43773196048 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Tue, 15 Sep 2020 06:20:41 +0000 Subject: [PATCH 15/23] fix expect_that --- cpp/src/parquet/level_conversion_test.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index 1ae3f91413c..8a7954ede66 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -197,7 +197,7 @@ TYPED_TEST(NestedListTest, OuterMostTest) { typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, offsets.data() + 4); + EXPECT_EQ(next_position, offsets.data() + 4); EXPECT_THAT(offsets, testing::ElementsAre(0, 3, 7, 7, 7)); EXPECT_EQ(validity_io.values_read, 4); @@ -227,7 +227,7 @@ TYPED_TEST(NestedListTest, MiddleListTest) { typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, offsets.data() + 7); + EXPECT_EQ(next_position, offsets.data() + 7); EXPECT_THAT(offsets, testing::ElementsAre(0, 0, 2, 2, 3, 5, 5, 6)); EXPECT_EQ(validity_io.values_read, 7); @@ -257,7 +257,7 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, offsets.data() + 6); + EXPECT_EQ(next_position, offsets.data() + 6); EXPECT_THAT(offsets, testing::ElementsAre(0, 3, 3, 3, 3, 5, 6)); EXPECT_EQ(validity_io.values_read, 6); @@ -292,7 +292,7 @@ TYPED_TEST(NestedListTest, SimpleLongList) { typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( this->test_data_, level_info, &validity_io, offsets.data()); - EXPECT_THAT(next_position, offsets.data() + 65); + EXPECT_EQ(next_position, offsets.data() + 65); EXPECT_THAT(offsets, testing::ElementsAreArray(expected_offsets)); EXPECT_EQ(validity_io.values_read, 65); From f4d082d0e5958a1f8d977b388a082acb6fbfc348 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Wed, 16 Sep 2020 03:34:19 +0000 Subject: [PATCH 16/23] address more comments --- cpp/cmake_modules/SetupCxxFlags.cmake | 2 +- cpp/src/parquet/CMakeLists.txt | 10 +++++++--- cpp/src/parquet/level_conversion.cc | 7 ++++++- cpp/src/parquet/level_conversion.h | 3 +-- cpp/src/parquet/level_conversion_benchmark.cc | 1 + cpp/src/parquet/level_conversion_inc.h | 1 + cpp/src/parquet/level_conversion_test.cc | 3 +++ 7 files changed, 20 insertions(+), 7 deletions(-) diff --git a/cpp/cmake_modules/SetupCxxFlags.cmake b/cpp/cmake_modules/SetupCxxFlags.cmake index 96e79c69fc4..1606199c406 100644 --- a/cpp/cmake_modules/SetupCxxFlags.cmake +++ b/cpp/cmake_modules/SetupCxxFlags.cmake @@ -50,7 +50,7 @@ if(ARROW_CPU_FLAG STREQUAL "x86") # skylake-avx512 consists of AVX512F,AVX512BW,AVX512VL,AVX512CD,AVX512DQ set(ARROW_AVX512_FLAG "-march=skylake-avx512 -mbmi2") # Append the avx2/avx512 subset option also, fix issue ARROW-9877 for homebrew-cpp - set(ARROW_AVX2_FLAG "${ARROW_AVX2_FLAG} -mavx2 -mbmi2") + set(ARROW_AVX2_FLAG "${ARROW_AVX2_FLAG} -mavx2") set(ARROW_AVX512_FLAG "${ARROW_AVX512_FLAG} -mavx512f -mavx512cd -mavx512vl -mavx512dq -mavx512bw") check_cxx_compiler_flag(${ARROW_SSE4_2_FLAG} CXX_SUPPORTS_SSE4_2) diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index 3722a229338..a137ec775c2 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -207,13 +207,17 @@ if(CXX_SUPPORTS_AVX2) # AVX2 is used as a proxy for BMI2. list(APPEND PARQUET_SRCS level_comparison_avx2.cc level_conversion_bmi2.cc) set_source_files_properties(level_comparison_avx2.cc - level_conversion_bmi2.cc PROPERTIES SKIP_PRECOMPILE_HEADERS ON COMPILE_FLAGS - "${ARROW_AVX2_FLAG} -DARROW_HAVE_BMI2") - + "${ARROW_AVX2_FLAG}") + set_source_files_properties(level_conversion_bmi2.cc + PROPERTIES + SKIP_PRECOMPILE_HEADERS + ON + COMPILE_FLAGS + "${ARROW_AVX2_FLAG} -DARROW_HAVE_BMI2 -mbmi2") endif() if(PARQUET_REQUIRE_ENCRYPTION) diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index 9868f58ec4a..fa282c27178 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -178,13 +178,18 @@ void DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, ? DefLevelsToBitmapBmi2WithRepeatedParent : standard::DefLevelsToBitmapSimd; fn(def_levels, num_def_levels, level_info, output); +#elif ARROW_LITTLE_ENDIAN + standard::DefLevelsToBitmapSimd(def_levels, num_def_levels, level_info, output); #else - // This indicates we are likely on a big-endian platformat which don't have a + // Big-endian platforms don't have a // pext function and the current implementation of BitRunReader is always linear // in bits, so this method will likely be the fastest alternative. DefLevelsToBitmapScalar(def_levels, num_def_levels, level_info, output); #endif } else { + // SIMD here applies to all platforms because the only operation that + // happens is def_levels->bitmap which should have good SIMD options + // on all platforms. standard::DefLevelsToBitmapSimd( def_levels, num_def_levels, level_info, output); } diff --git a/cpp/src/parquet/level_conversion.h b/cpp/src/parquet/level_conversion.h index 570f7dde822..c664cbae4cb 100644 --- a/cpp/src/parquet/level_conversion.h +++ b/cpp/src/parquet/level_conversion.h @@ -19,7 +19,6 @@ #include -#include "arrow/util/bitmap.h" #include "parquet/platform.h" #include "parquet/schema.h" @@ -171,7 +170,7 @@ void PARQUET_EXPORT DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def // the next offset. See documentation onf DefLevelsToBitmap for when to use this // method vs the other ones in this file for reconstruction. // -// Offsets must be size to 1 + values_read_upper_bound. +// Offsets must be sized to 1 + values_read_upper_bound. void PARQUET_EXPORT DefRepLevelsToList(const int16_t* def_levels, const int16_t* rep_levels, int64_t num_def_levels, LevelInfo level_info, diff --git a/cpp/src/parquet/level_conversion_benchmark.cc b/cpp/src/parquet/level_conversion_benchmark.cc index 759b68fac58..f9e91c4820f 100644 --- a/cpp/src/parquet/level_conversion_benchmark.cc +++ b/cpp/src/parquet/level_conversion_benchmark.cc @@ -38,6 +38,7 @@ std::vector RunDefinitionLevelsToBitmap(const std::vector& def parquet::internal::LevelInfo info; info.def_level = kHasRepeatedElements; info.repeated_ancestor_def_level = kPresentDefLevel; + info.rep_level = 1; parquet::internal::ValidityBitmapInputOutput validity_io; validity_io.values_read_upper_bound = def_levels.size(); validity_io.valid_bits = bitmap.data(); diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h index 4c24004bbf1..c7b8c9204f9 100644 --- a/cpp/src/parquet/level_conversion_inc.h +++ b/cpp/src/parquet/level_conversion_inc.h @@ -68,6 +68,7 @@ inline uint64_t RunBasedExtractImpl(uint64_t bitmap, uint64_t select_bitmap) { } inline uint64_t ExtractBits(uint64_t bitmap, uint64_t select_bitmap) { +// MING32 doesn't support 64-bit pext. #if defined(ARROW_HAVE_BMI2) && !defined(__MINGW32__) return _pext_u64(bitmap, select_bitmap); #else diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index 8a7954ede66..ccbf4f6056d 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -319,6 +319,9 @@ TYPED_TEST(NestedListTest, TestOverflow) { this->test_data_.def_levels_ = std::vector{2}; this->test_data_.rep_levels_ = std::vector{0}; + // Offsets is populated as the cumulative sum of all elements, + // so populating the offsets[0] with max-value impacts the + // other values populated. std::vector offsets( 2, std::numeric_limits::max()); From a3b52a4cd5522b5ed1324ab420ff8c2c6da6c83c Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Wed, 16 Sep 2020 04:37:18 +0000 Subject: [PATCH 17/23] formta --- cpp/src/arrow/buffer.cc | 2 +- cpp/src/parquet/arrow/reader.cc | 111 ++++++++++------------- cpp/src/parquet/level_conversion.cc | 3 +- cpp/src/parquet/level_conversion_test.cc | 2 +- 4 files changed, 54 insertions(+), 64 deletions(-) diff --git a/cpp/src/arrow/buffer.cc b/cpp/src/arrow/buffer.cc index 17884db9476..9215d9ab544 100644 --- a/cpp/src/arrow/buffer.cc +++ b/cpp/src/arrow/buffer.cc @@ -206,7 +206,7 @@ class PoolBuffer : public ResizableBuffer { } Status Resize(const int64_t new_size, bool shrink_to_fit = true) override { - if (new_size < 0) { + if (ARROW_PREDICT_FALSE(new_size < 0)) { return Status::Invalid("Negative buffer resize: ", new_size); } if (mutable_data_ && shrink_to_fit && new_size <= size_) { diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index 02e37903db5..8b3bfcb5704 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -103,12 +103,16 @@ class ColumnReaderImpl : public ColumnReader { ::arrow::Status NextBatch(int64_t batch_size, std::shared_ptr<::arrow::ChunkedArray>* out) final { RETURN_NOT_OK(LoadBatch(batch_size)); - return BuildArray(batch_size, out); + RETURN_NOT_OK(BuildArray(batch_size, out)); + for (int x = 0; x < (*out)->num_chunks(); x++) { + RETURN_NOT_OK((*out)->chunk(x)->Validate()); + } + return Status::OK(); } virtual ::arrow::Status LoadBatch(int64_t num_records) = 0; - virtual ::arrow::Status BuildArray(int64_t number_of_slots, + virtual ::arrow::Status BuildArray(int64_t length_upper_bound, std::shared_ptr<::arrow::ChunkedArray>* out) = 0; virtual bool IsOrHasRepeatedChild() const = 0; }; @@ -442,7 +446,7 @@ class LeafReader : public ColumnReaderImpl { END_PARQUET_CATCH_EXCEPTIONS } - ::arrow::Status BuildArray(int64_t expected_array_entries, + ::arrow::Status BuildArray(int64_t length_upper_bound, std::shared_ptr<::arrow::ChunkedArray>* out) final { *out = out_; return Status::OK(); @@ -489,60 +493,55 @@ class ListReader : public ColumnReaderImpl { return item_reader_->LoadBatch(number_of_records); } - Status BuildArray(int64_t number_of_slots, + Status BuildArray(int64_t length_upper_bound, std::shared_ptr* out) override { const int16_t* def_levels; const int16_t* rep_levels; int64_t num_levels; RETURN_NOT_OK(item_reader_->GetDefLevels(&def_levels, &num_levels)); RETURN_NOT_OK(item_reader_->GetRepLevels(&rep_levels, &num_levels)); - std::shared_ptr validity_buffer; + std::shared_ptr validity_buffer; ::parquet::internal::ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = number_of_slots; + validity_io.values_read_upper_bound = length_upper_bound; if (field_->nullable()) { - ARROW_ASSIGN_OR_RAISE(validity_buffer, AllocateBitmap(number_of_slots, ctx_->pool)); - + ARROW_ASSIGN_OR_RAISE( + validity_buffer, + AllocateResizableBuffer(BitUtil::BytesForBits(length_upper_bound), ctx_->pool)); validity_io.valid_bits = validity_buffer->mutable_data(); } ARROW_ASSIGN_OR_RAISE( - std::shared_ptr lengths_buffer, - AllocateBuffer(sizeof(IndexType) * std::max(int64_t{2}, number_of_slots + 1), - ctx_->pool)); - // ensure zero initialization in case we have reached a zero length list (and + std::shared_ptr offsets_buffer, + AllocateResizableBuffer( + sizeof(IndexType) * std::max(int64_t{1}, length_upper_bound + 1), + ctx_->pool)); + // Ensure zero initialization in case we have reached a zero length list (and // because first entry is always zero). - IndexType* length_data = reinterpret_cast(lengths_buffer->mutable_data()); - std::fill(length_data, length_data + 2, 0); + IndexType* offset_data = reinterpret_cast(offsets_buffer->mutable_data()); + offset_data[0] = 0; BEGIN_PARQUET_CATCH_EXCEPTIONS ::parquet::internal::DefRepLevelsToList(def_levels, rep_levels, num_levels, - level_info_, &validity_io, length_data); + level_info_, &validity_io, offset_data); END_PARQUET_CATCH_EXCEPTIONS - // We might return less then the requested slot (i.e. reaching an end of a file) - // ensure we've set all the bits here. - if (validity_io.values_read < number_of_slots) { - // + 1 because arrays lengths are values + 1 - std::fill(length_data + validity_io.values_read + 1, - length_data + number_of_slots + 1, 0); - if (validity_io.valid_bits != nullptr && - BitUtil::BytesForBits(validity_io.values_read) < - BitUtil::BytesForBits(number_of_slots)) { - std::fill(validity_io.valid_bits + BitUtil::BytesForBits(validity_io.values_read), - validity_io.valid_bits + BitUtil::BytesForBits(number_of_slots), 0); - } - } - RETURN_NOT_OK(item_reader_->BuildArray(length_data[validity_io.values_read], out)); + RETURN_NOT_OK(item_reader_->BuildArray(offset_data[validity_io.values_read], out)); + + // Resize to actual number of elements returned. + RETURN_NOT_OK( + offsets_buffer->Resize((validity_io.values_read + 1) * sizeof(IndexType))); + if (validity_buffer != nullptr) { + RETURN_NOT_OK( + validity_buffer->Resize(BitUtil::BytesForBits(validity_io.values_read))); + } ARROW_ASSIGN_OR_RAISE(std::shared_ptr item_chunk, ChunksToSingle(**out)); std::vector> buffers{ - validity_io.null_count > 0 ? validity_buffer : std::shared_ptr(), - lengths_buffer}; + validity_io.null_count > 0 ? validity_buffer : nullptr, offsets_buffer}; auto data = std::make_shared( field_->type(), /*length=*/validity_io.values_read, std::move(buffers), std::vector>{item_chunk}, validity_io.null_count); std::shared_ptr result = ::arrow::MakeArray(data); - RETURN_NOT_OK(result->Validate()); *out = std::make_shared(result); return Status::OK(); } @@ -568,7 +567,7 @@ class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { children_(std::move(children)) { // There could be a mix of children some might be repeated some might not be. // If possible use one that isn't since that will be guaranteed to have the least - // number of rep/def levels. + // number of levels to reconstruct a nullable bitmap. auto result = std::find_if(children_.begin(), children_.end(), [](const std::unique_ptr& child) { return !child->IsOrHasRepeatedChild(); @@ -590,7 +589,8 @@ class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { } return Status::OK(); } - Status BuildArray(int64_t records_to_read, std::shared_ptr* out) override; + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override; Status GetDefLevels(const int16_t** data, int64_t* length) override; Status GetRepLevels(const int16_t** data, int64_t* length) override; const std::shared_ptr field() override { return filtered_field_; } @@ -608,7 +608,7 @@ Status StructReader::GetDefLevels(const int16_t** data, int64_t* length) { *data = nullptr; if (children_.size() == 0) { *length = 0; - return Status::Invalid("StructReader had no childre"); + return Status::Invalid("StructReader had no children"); } // This method should only be called when this struct or one of its parents @@ -616,12 +616,6 @@ Status StructReader::GetDefLevels(const int16_t** data, int64_t* length) { // Meaning all children must have rep/def levels associated // with them. RETURN_NOT_OK(def_rep_level_child_->GetDefLevels(data, length)); - - if (*data == nullptr) { - // Only happens if there are actually 0 rows available. - *data = nullptr; - *length = 0; - } return Status::OK(); } @@ -637,23 +631,18 @@ Status StructReader::GetRepLevels(const int16_t** data, int64_t* length) { // Meaning all children must have rep/def levels associated // with them. RETURN_NOT_OK(def_rep_level_child_->GetRepLevels(data, length)); - - if (data == nullptr) { - // Only happens if there are actually 0 rows available. - *data = nullptr; - *length = 0; - } return Status::OK(); } -Status StructReader::BuildArray(int64_t records_to_read, +Status StructReader::BuildArray(int64_t length_upper_bound, std::shared_ptr* out) { std::vector> children_array_data; - std::shared_ptr null_bitmap; + std::shared_ptr null_bitmap; ::parquet::internal::ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = records_to_read; - validity_io.values_read = records_to_read; + validity_io.values_read_upper_bound = length_upper_bound; + // This simplifies accounting below. + validity_io.values_read = length_upper_bound; BEGIN_PARQUET_CATCH_EXCEPTIONS const int16_t* def_levels; @@ -661,24 +650,25 @@ Status StructReader::BuildArray(int64_t records_to_read, int64_t num_levels; if (has_repeated_child_) { - ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); + ARROW_ASSIGN_OR_RAISE( + null_bitmap, + AllocateResizableBuffer(BitUtil::BytesForBits(length_upper_bound), ctx_->pool)); validity_io.valid_bits = null_bitmap->mutable_data(); RETURN_NOT_OK(GetDefLevels(&def_levels, &num_levels)); RETURN_NOT_OK(GetRepLevels(&rep_levels, &num_levels)); DefRepLevelsToBitmap(def_levels, rep_levels, num_levels, level_info_, &validity_io); } else if (filtered_field_->nullable()) { - ARROW_ASSIGN_OR_RAISE(null_bitmap, AllocateBitmap(records_to_read, ctx_->pool)); + ARROW_ASSIGN_OR_RAISE( + null_bitmap, + AllocateResizableBuffer(BitUtil::BytesForBits(length_upper_bound), ctx_->pool)); validity_io.valid_bits = null_bitmap->mutable_data(); RETURN_NOT_OK(GetDefLevels(&def_levels, &num_levels)); DefLevelsToBitmap(def_levels, num_levels, level_info_, &validity_io); } // Ensure all values are initialized. - if (null_bitmap && BitUtil::BytesForBits(validity_io.values_read) < - BitUtil::BytesForBits(records_to_read)) { - std::fill( - null_bitmap->mutable_data() + BitUtil::BytesForBits(validity_io.values_read), - null_bitmap->mutable_data() + BitUtil::BytesForBits(records_to_read), 0); + if (null_bitmap) { + RETURN_NOT_OK(null_bitmap->Resize(BitUtil::BytesForBits(validity_io.values_read))); } END_PARQUET_CATCH_EXCEPTIONS @@ -694,14 +684,13 @@ Status StructReader::BuildArray(int64_t records_to_read, validity_io.values_read = children_array_data.front()->length; } - std::vector> buffers{ - validity_io.null_count > 0 ? null_bitmap : std::shared_ptr()}; + std::vector> buffers{validity_io.null_count > 0 ? null_bitmap + : nullptr}; auto data = std::make_shared(filtered_field_->type(), /*length=*/validity_io.values_read, std::move(buffers), std::move(children_array_data)); std::shared_ptr result = ::arrow::MakeArray(data); - RETURN_NOT_OK(result->Validate()); *out = std::make_shared(result); return Status::OK(); diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index fa282c27178..2279b3f3211 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -179,7 +179,8 @@ void DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, : standard::DefLevelsToBitmapSimd; fn(def_levels, num_def_levels, level_info, output); #elif ARROW_LITTLE_ENDIAN - standard::DefLevelsToBitmapSimd(def_levels, num_def_levels, level_info, output); + standard::DefLevelsToBitmapSimd( + def_levels, num_def_levels, level_info, output); #else // Big-endian platforms don't have a // pext function and the current implementation of BitRunReader is always linear diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index ccbf4f6056d..d31a5c16ab2 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -320,7 +320,7 @@ TYPED_TEST(NestedListTest, TestOverflow) { this->test_data_.rep_levels_ = std::vector{0}; // Offsets is populated as the cumulative sum of all elements, - // so populating the offsets[0] with max-value impacts the + // so populating the offsets[0] with max-value impacts the // other values populated. std::vector offsets( 2, std::numeric_limits::max()); From 4aaec3221f83e4181da0828e395960e2bd58c96e Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Wed, 16 Sep 2020 05:22:57 +0000 Subject: [PATCH 18/23] refactor conversion test --- cpp/src/parquet/level_conversion_test.cc | 184 +++++++++++------------ 1 file changed, 91 insertions(+), 93 deletions(-) diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index d31a5c16ab2..35e36411086 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -139,37 +139,59 @@ TEST(DefLevelsToBitmap, WithRepetitionLevelFiltersOutEmptyListValues) { EXPECT_EQ(io.values_read, 4); // value should get overwritten. } -class MultiLevelTestData { +struct MultiLevelTestData { public: + std::vector def_levels; + std::vector rep_levels; +}; + +MultiLevelTestData TriplyNestedList() { // Triply nested list values borrow from write_path // [null, [[1 , null, 3], []], []], // [[[]], [[], [1, 2]], null, [[3]]], // null, // [] - std::vector def_levels_{2, 7, 6, 7, 5, 3, // first row - 5, 5, 7, 7, 2, 7, // second row - 0, // third row - 1}; - std::vector rep_levels_{0, 1, 3, 3, 2, 1, // first row - 0, 1, 2, 3, 1, 1, // second row - 0, 0}; -}; + return MultiLevelTestData{ + /*def_levels=*/std::vector{2, 7, 6, 7, 5, 3, // first row + 5, 5, 7, 7, 2, 7, // second row + 0, // third row + 1}, + /*rep_levels=*/std::vector{0, 1, 3, 3, 2, 1, // first row + 0, 1, 2, 3, 1, 1, // second row + 0, 0}}; +} template class NestedListTest : public testing::Test { public: - MultiLevelTestData test_data_; + void InitForLength(int length) { + this->validity_bits_.clear(); + this->validity_bits_.insert(this->validity_bits_.end(), length, 0); + validity_io_.valid_bits = validity_bits_.data(); + validity_io_.values_read_upper_bound = length; + offsets_.clear(); + offsets_.insert(offsets_.end(), length + 1, 0); + } + + typename ConverterType::OffsetsType* Run(const MultiLevelTestData& test_data, + LevelInfo level_info) { + return this->converter_.ComputeListInfo(test_data, level_info, &validity_io_, + offsets_.data()); + } + ConverterType converter_; + ValidityBitmapInputOutput validity_io_; + std::vector validity_bits_; + std::vector offsets_; }; -template +template struct RepDefLevelConverter { - using ListLengthType = ListType; - ListLengthType* ComputeListInfo(const MultiLevelTestData& test_data, - LevelInfo level_info, ValidityBitmapInputOutput* output, - ListType* offsets) { - DefRepLevelsToList(test_data.def_levels_.data(), test_data.rep_levels_.data(), - test_data.def_levels_.size(), level_info, output, offsets); + using OffsetsType = IndexType; + OffsetsType* ComputeListInfo(const MultiLevelTestData& test_data, LevelInfo level_info, + ValidityBitmapInputOutput* output, IndexType* offsets) { + DefRepLevelsToList(test_data.def_levels.data(), test_data.rep_levels.data(), + test_data.def_levels.size(), level_info, output, offsets); return offsets + output->values_read; } }; @@ -189,20 +211,16 @@ TYPED_TEST(NestedListTest, OuterMostTest) { level_info.rep_level = 1; level_info.def_level = 2; - std::vector offsets(5, 0); - uint64_t validity_output; - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = 4; - validity_io.valid_bits = reinterpret_cast(&validity_output); - typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, offsets.data()); + this->InitForLength(4); + typename TypeParam::OffsetsType* next_position = + this->Run(TriplyNestedList(), level_info); - EXPECT_EQ(next_position, offsets.data() + 4); - EXPECT_THAT(offsets, testing::ElementsAre(0, 3, 7, 7, 7)); + EXPECT_EQ(next_position, this->offsets_.data() + 4); + EXPECT_THAT(this->offsets_, testing::ElementsAre(0, 3, 7, 7, 7)); - EXPECT_EQ(validity_io.values_read, 4); - EXPECT_EQ(validity_io.null_count, 1); - EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/4), "1101"); + EXPECT_EQ(this->validity_io_.values_read, 4); + EXPECT_EQ(this->validity_io_.null_count, 1); + EXPECT_EQ(BitmapToString(this->validity_io_.valid_bits, /*length=*/4), "1101"); } TYPED_TEST(NestedListTest, MiddleListTest) { @@ -219,20 +237,16 @@ TYPED_TEST(NestedListTest, MiddleListTest) { level_info.def_level = 4; level_info.repeated_ancestor_def_level = 2; - std::vector offsets(8, 0); - uint64_t validity_output; - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = 7; - validity_io.valid_bits = reinterpret_cast(&validity_output); - typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, offsets.data()); + this->InitForLength(7); + typename TypeParam::OffsetsType* next_position = + this->Run(TriplyNestedList(), level_info); - EXPECT_EQ(next_position, offsets.data() + 7); - EXPECT_THAT(offsets, testing::ElementsAre(0, 0, 2, 2, 3, 5, 5, 6)); + EXPECT_EQ(next_position, this->offsets_.data() + 7); + EXPECT_THAT(this->offsets_, testing::ElementsAre(0, 0, 2, 2, 3, 5, 5, 6)); - EXPECT_EQ(validity_io.values_read, 7); - EXPECT_EQ(validity_io.null_count, 2); - EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/7), "0111101"); + EXPECT_EQ(this->validity_io_.values_read, 7); + EXPECT_EQ(this->validity_io_.null_count, 2); + EXPECT_EQ(BitmapToString(this->validity_io_.valid_bits, /*length=*/7), "0111101"); } TYPED_TEST(NestedListTest, InnerMostListTest) { @@ -249,20 +263,16 @@ TYPED_TEST(NestedListTest, InnerMostListTest) { level_info.def_level = 6; level_info.repeated_ancestor_def_level = 4; - std::vector offsets(7, 0); - uint64_t validity_output; - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = 6; - validity_io.valid_bits = reinterpret_cast(&validity_output); - typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, offsets.data()); + this->InitForLength(6); + typename TypeParam::OffsetsType* next_position = + this->Run(TriplyNestedList(), level_info); - EXPECT_EQ(next_position, offsets.data() + 6); - EXPECT_THAT(offsets, testing::ElementsAre(0, 3, 3, 3, 3, 5, 6)); + EXPECT_EQ(next_position, this->offsets_.data() + 6); + EXPECT_THAT(this->offsets_, testing::ElementsAre(0, 3, 3, 3, 3, 5, 6)); - EXPECT_EQ(validity_io.values_read, 6); - EXPECT_EQ(validity_io.null_count, 0); - EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/6), "111111"); + EXPECT_EQ(this->validity_io_.values_read, 6); + EXPECT_EQ(this->validity_io_.null_count, 0); + EXPECT_EQ(BitmapToString(this->validity_io_.valid_bits, /*length=*/6), "111111"); } TYPED_TEST(NestedListTest, SimpleLongList) { @@ -271,33 +281,28 @@ TYPED_TEST(NestedListTest, SimpleLongList) { level_info.def_level = 2; level_info.repeated_ancestor_def_level = 0; + MultiLevelTestData test_data; // No empty lists. - this->test_data_.def_levels_ = std::vector(65 * 9, 2); - this->test_data_.rep_levels_.clear(); + test_data.def_levels = std::vector(65 * 9, 2); for (int x = 0; x < 65; x++) { - this->test_data_.rep_levels_.push_back(0); - this->test_data_.rep_levels_.insert(this->test_data_.rep_levels_.end(), 8, - /*rep_level=*/1); + test_data.rep_levels.push_back(0); + test_data.rep_levels.insert(test_data.rep_levels.end(), 8, + /*rep_level=*/1); } - std::vector offsets(66, 0); - std::vector expected_offsets(66, 0); + std::vector expected_offsets(66, 0); for (size_t x = 1; x < expected_offsets.size(); x++) { expected_offsets[x] = x * 9; } - std::vector validity_output(9, 0); - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = 65; - validity_io.valid_bits = validity_output.data(); - typename TypeParam::ListLengthType* next_position = this->converter_.ComputeListInfo( - this->test_data_, level_info, &validity_io, offsets.data()); - - EXPECT_EQ(next_position, offsets.data() + 65); - EXPECT_THAT(offsets, testing::ElementsAreArray(expected_offsets)); - - EXPECT_EQ(validity_io.values_read, 65); - EXPECT_EQ(validity_io.null_count, 0); - EXPECT_EQ(BitmapToString(validity_io.valid_bits, /*length=*/65), + this->InitForLength(65); + typename TypeParam::OffsetsType* next_position = this->Run(test_data, level_info); + + EXPECT_EQ(next_position, this->offsets_.data() + 65); + EXPECT_THAT(this->offsets_, testing::ElementsAreArray(expected_offsets)); + + EXPECT_EQ(this->validity_io_.values_read, 65); + EXPECT_EQ(this->validity_io_.null_count, 0); + EXPECT_EQ(BitmapToString(this->validity_io_.valid_bits, /*length=*/65), "11111111 " "11111111 " "11111111 " @@ -315,35 +320,28 @@ TYPED_TEST(NestedListTest, TestOverflow) { level_info.def_level = 2; level_info.repeated_ancestor_def_level = 0; - // No empty lists. - this->test_data_.def_levels_ = std::vector{2}; - this->test_data_.rep_levels_ = std::vector{0}; + MultiLevelTestData test_data; + test_data.def_levels = std::vector{2}; + test_data.rep_levels = std::vector{0}; + this->InitForLength(2); // Offsets is populated as the cumulative sum of all elements, // so populating the offsets[0] with max-value impacts the // other values populated. - std::vector offsets( - 2, std::numeric_limits::max()); + this->offsets_[0] = std::numeric_limits::max(); + this->offsets_[1] = std::numeric_limits::max(); + ASSERT_THROW(this->Run(test_data, level_info), ParquetException); - std::vector validity_output(1, 0); - ValidityBitmapInputOutput validity_io; - validity_io.values_read_upper_bound = 1; - validity_io.valid_bits = validity_output.data(); - ASSERT_THROW(this->converter_.ComputeListInfo(this->test_data_, level_info, - &validity_io, offsets.data()), - ParquetException); + ASSERT_THROW(this->Run(test_data, level_info), ParquetException); // Same thing should happen if the list already existed. - this->test_data_.rep_levels_ = std::vector{1}; - ASSERT_THROW(this->converter_.ComputeListInfo(this->test_data_, level_info, - &validity_io, offsets.data()), - ParquetException); + test_data.rep_levels = std::vector{1}; + ASSERT_THROW(this->Run(test_data, level_info), ParquetException); // Should be OK because it shouldn't increment. - this->test_data_.def_levels_ = std::vector{0}; - this->test_data_.rep_levels_ = std::vector{0}; - this->converter_.ComputeListInfo(this->test_data_, level_info, &validity_io, - offsets.data()); + test_data.def_levels = std::vector{0}; + test_data.rep_levels = std::vector{0}; + this->Run(test_data, level_info); } TEST(TestOnlyRunBasedExtract, BasicTest) { From 2e08db41815991fc68ade6993b3bf86dfb922b3e Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Wed, 16 Sep 2020 05:36:03 +0000 Subject: [PATCH 19/23] fix guard for scalar function --- cpp/src/parquet/level_conversion.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index 2279b3f3211..ab4aaa6f1ef 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -36,7 +36,7 @@ namespace { using ::arrow::internal::CpuInfo; -#if !defined(ARROW_HAVE_RUNTIME_BMI2) +#if !ARROW_LITTLE_ENDIAN void DefLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, ValidityBitmapInputOutput* output) { ::arrow::internal::FirstTimeBitmapWriter valid_bits_writer( From 4093cc99229708fc9ca498342a7cbdd8bee4c086 Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Wed, 16 Sep 2020 15:39:04 +0000 Subject: [PATCH 20/23] make windows pedantry happier --- cpp/src/parquet/level_conversion_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index 35e36411086..be432c41586 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -292,7 +292,7 @@ TYPED_TEST(NestedListTest, SimpleLongList) { std::vector expected_offsets(66, 0); for (size_t x = 1; x < expected_offsets.size(); x++) { - expected_offsets[x] = x * 9; + expected_offsets[x] = static_cast(x) * 9; } this->InitForLength(65); typename TypeParam::OffsetsType* next_position = this->Run(test_data, level_info); From 4743d80fc3fcaf9b46f8a9eb0b33629f9ce2c45b Mon Sep 17 00:00:00 2001 From: Antoine Pitrou Date: Mon, 21 Sep 2020 16:35:00 +0200 Subject: [PATCH 21/23] Use DefLevelsToBitmapScalar instead of emulated PEXT --- cpp/src/parquet/level_conversion.cc | 17 +++++------------ cpp/src/parquet/level_conversion_inc.h | 4 ++-- cpp/src/parquet/level_conversion_test.cc | 3 ++- 3 files changed, 9 insertions(+), 15 deletions(-) diff --git a/cpp/src/parquet/level_conversion.cc b/cpp/src/parquet/level_conversion.cc index ab4aaa6f1ef..3f57be39e72 100644 --- a/cpp/src/parquet/level_conversion.cc +++ b/cpp/src/parquet/level_conversion.cc @@ -36,7 +36,6 @@ namespace { using ::arrow::internal::CpuInfo; -#if !ARROW_LITTLE_ENDIAN void DefLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, LevelInfo level_info, ValidityBitmapInputOutput* output) { ::arrow::internal::FirstTimeBitmapWriter valid_bits_writer( @@ -71,7 +70,6 @@ void DefLevelsToBitmapScalar(const int16_t* def_levels, int64_t num_def_levels, "(i.e. FixedSizeLists with null values are not supported"); } } -#endif template void DefRepLevelsToListInfo(const int16_t* def_levels, const int16_t* rep_levels, @@ -173,18 +171,13 @@ void DefLevelsToBitmap(const int16_t* def_levels, int64_t num_def_levels, if (level_info.rep_level > 0) { #if defined(ARROW_HAVE_RUNTIME_BMI2) using FunctionType = decltype(&standard::DefLevelsToBitmapSimd); - static FunctionType fn = - CpuInfo::GetInstance()->HasEfficientBmi2() - ? DefLevelsToBitmapBmi2WithRepeatedParent - : standard::DefLevelsToBitmapSimd; + // DefLevelsToBitmapSimd with emulated PEXT would be slow, so use the + // scalar version if BMI2 is unavailable. + static FunctionType fn = CpuInfo::GetInstance()->HasEfficientBmi2() + ? DefLevelsToBitmapBmi2WithRepeatedParent + : DefLevelsToBitmapScalar; fn(def_levels, num_def_levels, level_info, output); -#elif ARROW_LITTLE_ENDIAN - standard::DefLevelsToBitmapSimd( - def_levels, num_def_levels, level_info, output); #else - // Big-endian platforms don't have a - // pext function and the current implementation of BitRunReader is always linear - // in bits, so this method will likely be the fastest alternative. DefLevelsToBitmapScalar(def_levels, num_def_levels, level_info, output); #endif } else { diff --git a/cpp/src/parquet/level_conversion_inc.h b/cpp/src/parquet/level_conversion_inc.h index c7b8c9204f9..c688f748043 100644 --- a/cpp/src/parquet/level_conversion_inc.h +++ b/cpp/src/parquet/level_conversion_inc.h @@ -92,13 +92,13 @@ int64_t DefLevelsBatchToBitmap(const int16_t* def_levels, const int64_t batch_si def_levels, batch_size, level_info.repeated_ancestor_def_level - 1); uint64_t selected_bits = ExtractBits(defined_bitmap, present_bitmap); int64_t selected_count = ::arrow::BitUtil::PopCount(present_bitmap); - if (selected_count > upper_bound_remaining) { + if (ARROW_PREDICT_FALSE(selected_count > upper_bound_remaining)) { throw ParquetException("Values read exceeded upper bound"); } writer->AppendWord(selected_bits, selected_count); return ::arrow::BitUtil::PopCount(selected_bits); } else { - if (batch_size > upper_bound_remaining) { + if (ARROW_PREDICT_FALSE(batch_size > upper_bound_remaining)) { std::stringstream ss; ss << "Values read exceeded upper bound"; throw ParquetException(ss.str()); diff --git a/cpp/src/parquet/level_conversion_test.cc b/cpp/src/parquet/level_conversion_test.cc index be432c41586..b4f2d3ad5d1 100644 --- a/cpp/src/parquet/level_conversion_test.cc +++ b/cpp/src/parquet/level_conversion_test.cc @@ -26,6 +26,7 @@ #include #include +#include "arrow/testing/gtest_compat.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap.h" #include "arrow/util/ubsan.h" @@ -199,7 +200,7 @@ struct RepDefLevelConverter { using ConverterTypes = ::testing::Types, RepDefLevelConverter>; -TYPED_TEST_CASE(NestedListTest, ConverterTypes); +TYPED_TEST_SUITE(NestedListTest, ConverterTypes); TYPED_TEST(NestedListTest, OuterMostTest) { // [null, [[1 , null, 3], []], []], From ce328259eb2c6e9dca24abed4781026a23c7286b Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Mon, 21 Sep 2020 15:31:07 +0000 Subject: [PATCH 22/23] remove residual writer v1 test --- python/pyarrow/tests/test_parquet.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/python/pyarrow/tests/test_parquet.py b/python/pyarrow/tests/test_parquet.py index cbc8ba026b0..81d6efcc6be 100644 --- a/python/pyarrow/tests/test_parquet.py +++ b/python/pyarrow/tests/test_parquet.py @@ -692,14 +692,6 @@ def test_pandas_can_write_nested_data(tempdir): # This succeeds under V2 _write_table(arrow_table, imos) - # Under V1 it fails. - with pytest.raises(ValueError): - import os - os.environ['ARROW_PARQUET_WRITER_ENGINE'] = 'V1' - imos = pa.BufferOutputStream() - _write_table(arrow_table, imos) - del os.environ['ARROW_PARQUET_WRITER_ENGINE'] - @pytest.mark.pandas @parametrize_legacy_dataset From fac386892051e59e8a8347ac583bf7e35d12dbaa Mon Sep 17 00:00:00 2001 From: Micah Kornfield Date: Mon, 21 Sep 2020 15:35:52 +0000 Subject: [PATCH 23/23] add warning comment --- cpp/src/parquet/CMakeLists.txt | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cpp/src/parquet/CMakeLists.txt b/cpp/src/parquet/CMakeLists.txt index a137ec775c2..22ad69219a3 100644 --- a/cpp/src/parquet/CMakeLists.txt +++ b/cpp/src/parquet/CMakeLists.txt @@ -212,6 +212,10 @@ if(CXX_SUPPORTS_AVX2) ON COMPILE_FLAGS "${ARROW_AVX2_FLAG}") + # WARNING: DO NOT BLINDLY COPY THIS CODE FOR OTHER BMI2 USE CASES. + # This code is always guarded by runtime dispatch which verifies + # BMI2 is present. For a very small number of CPUs AVX2 does not + # imply BMI2. set_source_files_properties(level_conversion_bmi2.cc PROPERTIES SKIP_PRECOMPILE_HEADERS