You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by GitBox <gi...@apache.org> on 2020/09/22 02:24:05 UTC

[GitHub] [arrow] wesm commented on a change in pull request #8219: ARROW-9603: [C++] Fix parquet write

wesm commented on a change in pull request #8219:
URL: https://github.com/apache/arrow/pull/8219#discussion_r492407465



##########
File path: cpp/src/parquet/arrow/arrow_reader_writer_test.cc
##########
@@ -2360,6 +2361,49 @@ TEST(ArrowReadWrite, SingleColumnNullableStruct) {
       3);
 }
 
+TEST(ArrowReadWrite, DisagreeingValidityBitmap) {}
+
+TEST(ArrowReadWrite, NestedRequiredField) {
+  auto int_field = ::arrow::field("int_array", ::arrow::int32(), /*nullable=*/false);
+  auto int_array = ::arrow::ArrayFromJSON(int_field->type(), "[0, 1, 2, 3, 4, 5, 7, 8]");
+  auto struct_field =
+      ::arrow::field("root", ::arrow::struct_({int_field}), /*nullable=*/true);
+  std::shared_ptr<Buffer> validity_bitmap;
+  ASSERT_OK_AND_ASSIGN(validity_bitmap, ::arrow::AllocateBitmap(8));
+  validity_bitmap->mutable_data()[0] = 0xCC;
+
+  auto struct_data = std::make_shared<ArrayData>(
+      struct_field->type(), /*length=*/8,
+      std::vector<std::shared_ptr<Buffer>>{validity_bitmap},
+      std::vector<std::shared_ptr<ArrayData>>{int_array->data()});
+  CheckSimpleRoundtrip(
+      ::arrow::Table::Make(
+          ::arrow::schema({struct_field}),
+          {std::make_shared<::arrow::ChunkedArray>(::arrow::MakeArray(struct_data))}),
+      /*row_group_size=*/8);

Review comment:
       Seems like there might be a helper function opportunity if this pattern is repeated in other test functions

##########
File path: cpp/src/parquet/arrow/writer.cc
##########
@@ -134,15 +134,14 @@ class ArrowColumnWriterV2 {
               std::shared_ptr<Array> values_array =
                   result.leaf_array->Slice(range.start, range.Size());
 
-              return column_writer->WriteArrow(result.def_levels, result.rep_levels,
-                                               result.def_rep_level_count, *values_array,
-                                               ctx);
+              PARQUET_CATCH_AND_RETURN(column_writer->WriteArrow(
+                  result.def_levels, result.rep_levels, result.def_rep_level_count,
+                  *values_array, ctx, level_builder->Nested(), result.leaf_is_nullable));

Review comment:
       Since `WriteArrow` returns Status, should we adopt that APIs must either return Status or throw an exception, but not both? (FWIW I regret that we chose to allow exceptions in the Parquet C++ project back in 2016)

##########
File path: cpp/src/parquet/column_writer.cc
##########
@@ -1130,37 +1185,60 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
     return values_to_write;
   }
 
+  void MaybeCalculateValidityBits(const int16_t* def_levels, int64_t batch_size,
+                                  int64_t* out_values_to_write,
+                                  int64_t* out_spaced_values_to_write,
+                                  int64_t* null_count) {
+    if (bits_buffer_ == nullptr) {
+      if (!level_info_.HasNullableValues()) {
+        *out_values_to_write = batch_size;
+        *out_spaced_values_to_write = batch_size;
+        *null_count = 0;
+      } else {
+        for (int x = 0; x < batch_size; x++) {
+          *out_values_to_write += def_levels[x] == level_info_.def_level ? 1 : 0;
+          *out_spaced_values_to_write +=
+              def_levels[x] >= level_info_.repeated_ancestor_def_level ? 1 : 0;
+        }
+        *null_count = *out_values_to_write - *out_spaced_values_to_write;
+      }
+      return;
+    }
+    // Shrink to fit possible causes another allocation, and would only be necessary
+    // on the last batch.
+    int64_t new_bitmap_size = BitUtil::BytesForBits(batch_size);
+    if (new_bitmap_size != bits_buffer_->size()) {
+      PARQUET_THROW_NOT_OK(
+          bits_buffer_->Resize(new_bitmap_size, /*shrink_to_fit=*/false));
+      bits_buffer_->ZeroPadding();
+    }
+    internal::ValidityBitmapInputOutput io;
+    io.valid_bits = bits_buffer_->mutable_data();
+    io.values_read_upper_bound = batch_size;
+    internal::DefLevelsToBitmap(def_levels, batch_size, level_info_, &io);
+    *out_values_to_write = io.values_read - io.null_count;
+    *out_spaced_values_to_write = io.values_read;
+    *null_count = io.null_count;
+  }
+
+  std::shared_ptr<Array> MaybeUpdateArray(std::shared_ptr<Array> array,

Review comment:
       `MaybeReplaceValidity`?

##########
File path: cpp/src/parquet/column_writer.cc
##########
@@ -1009,11 +1046,29 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
 
   Status WriteArrow(const int16_t* def_levels, const int16_t* rep_levels,
                     int64_t num_levels, const ::arrow::Array& array,
-                    ArrowWriteContext* ctx) override {
+                    ArrowWriteContext* ctx, bool nested, bool array_nullable) override {
+    bool leaf_is_not_nullable = !level_info_.HasNullableValues();
+    // Leaf nulls are canonical when there is only a single null element and it is at the
+    // leaf.
+    bool leaf_nulls_are_canonical =
+        (level_info_.def_level == level_info_.repeated_ancestor_def_level + 1) &&
+        array_nullable;
+    bool maybe_has_nulls = nested && !(leaf_is_not_nullable || leaf_nulls_are_canonical);

Review comment:
       The fact that `maybe_has_nulls` is false whenever `nested` is false seems odd

##########
File path: cpp/src/parquet/column_writer.cc
##########
@@ -1130,37 +1185,60 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, public TypedColumnWriter<
     return values_to_write;
   }
 
+  void MaybeCalculateValidityBits(const int16_t* def_levels, int64_t batch_size,
+                                  int64_t* out_values_to_write,
+                                  int64_t* out_spaced_values_to_write,
+                                  int64_t* null_count) {
+    if (bits_buffer_ == nullptr) {
+      if (!level_info_.HasNullableValues()) {
+        *out_values_to_write = batch_size;
+        *out_spaced_values_to_write = batch_size;
+        *null_count = 0;
+      } else {
+        for (int x = 0; x < batch_size; x++) {
+          *out_values_to_write += def_levels[x] == level_info_.def_level ? 1 : 0;
+          *out_spaced_values_to_write +=
+              def_levels[x] >= level_info_.repeated_ancestor_def_level ? 1 : 0;
+        }
+        *null_count = *out_values_to_write - *out_spaced_values_to_write;
+      }
+      return;
+    }
+    // Shrink to fit possible causes another allocation, and would only be necessary
+    // on the last batch.
+    int64_t new_bitmap_size = BitUtil::BytesForBits(batch_size);
+    if (new_bitmap_size != bits_buffer_->size()) {
+      PARQUET_THROW_NOT_OK(
+          bits_buffer_->Resize(new_bitmap_size, /*shrink_to_fit=*/false));
+      bits_buffer_->ZeroPadding();
+    }
+    internal::ValidityBitmapInputOutput io;
+    io.valid_bits = bits_buffer_->mutable_data();
+    io.values_read_upper_bound = batch_size;
+    internal::DefLevelsToBitmap(def_levels, batch_size, level_info_, &io);
+    *out_values_to_write = io.values_read - io.null_count;
+    *out_spaced_values_to_write = io.values_read;
+    *null_count = io.null_count;
+  }
+
+  std::shared_ptr<Array> MaybeUpdateArray(std::shared_ptr<Array> array,
+                                          int64_t new_null_count) {
+    if (bits_buffer_ == nullptr) {
+      return array;
+    }
+    std::vector<std::shared_ptr<Buffer>> buffers = array->data()->buffers;
+    buffers[0] = bits_buffer_;
+    DCHECK(array->num_fields() == 0);
+    return arrow::MakeArray(std::make_shared<ArrayData>(
+        array->type(), array->length(), std::move(buffers), new_null_count));

Review comment:
       Might be useful someday to have a helper function to make an array copy with a particular buffer replaced, I seem to recall a JIRA issue about this

##########
File path: cpp/src/parquet/arrow/arrow_reader_writer_test.cc
##########
@@ -2360,6 +2361,49 @@ TEST(ArrowReadWrite, SingleColumnNullableStruct) {
       3);
 }
 
+TEST(ArrowReadWrite, DisagreeingValidityBitmap) {}
+
+TEST(ArrowReadWrite, NestedRequiredField) {
+  auto int_field = ::arrow::field("int_array", ::arrow::int32(), /*nullable=*/false);
+  auto int_array = ::arrow::ArrayFromJSON(int_field->type(), "[0, 1, 2, 3, 4, 5, 7, 8]");
+  auto struct_field =
+      ::arrow::field("root", ::arrow::struct_({int_field}), /*nullable=*/true);
+  std::shared_ptr<Buffer> validity_bitmap;
+  ASSERT_OK_AND_ASSIGN(validity_bitmap, ::arrow::AllocateBitmap(8));
+  validity_bitmap->mutable_data()[0] = 0xCC;
+
+  auto struct_data = std::make_shared<ArrayData>(
+      struct_field->type(), /*length=*/8,
+      std::vector<std::shared_ptr<Buffer>>{validity_bitmap},
+      std::vector<std::shared_ptr<ArrayData>>{int_array->data()});
+  CheckSimpleRoundtrip(
+      ::arrow::Table::Make(
+          ::arrow::schema({struct_field}),
+          {std::make_shared<::arrow::ChunkedArray>(::arrow::MakeArray(struct_data))}),
+      /*row_group_size=*/8);
+}
+
+TEST(ArrowReadWrite, NestedNullableField) {
+  auto int_field = ::arrow::field("int_array", ::arrow::int32());
+  auto int_array = ::arrow::ArrayFromJSON(int_field->type(), "[0, null, 2, null, 4, 5, null, 8]");
+  auto struct_field =
+      ::arrow::field("root", ::arrow::struct_({int_field}), /*nullable=*/true);
+  std::shared_ptr<Buffer> validity_bitmap;
+  ASSERT_OK_AND_ASSIGN(validity_bitmap, ::arrow::AllocateBitmap(8));
+  validity_bitmap->mutable_data()[0] = 0xCC;
+
+  auto struct_data = std::make_shared<ArrayData>(
+      struct_field->type(), /*length=*/8,
+      std::vector<std::shared_ptr<Buffer>>{validity_bitmap},
+      std::vector<std::shared_ptr<ArrayData>>{int_array->data()});

Review comment:
       You can use `ArrayData::Make` for nicer syntax (don't have to write out `std::vector<std::shared_ptr<Buffer>>`)

##########
File path: cpp/src/parquet/arrow/path_internal.cc
##########
@@ -871,6 +877,8 @@ class MultipathLevelBuilderImpl : public MultipathLevelBuilder {
                      std::move(write_leaf_callback));
   }
 
+  bool Nested() const override { return !data_->child_data.empty(); }

Review comment:
       IsNested?




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org