pitrou commented on code in PR #50927:
URL: https://github.com/apache/arrow/pull/50927#discussion_r3853272604
##########
cpp/src/arrow/ipc/read_write_test.cc:
##########
@@ -71,6 +71,73 @@ using MetadataVector =
std::vector<std::shared_ptr<KeyValueMetadata>>;
namespace test {
+class UnionExtensionArray : public ExtensionArray {
Review Comment:
Can we move these two classes to `arrow/testing/extension_type.h`? They're
probably going to be useful for other tests at some point.
##########
cpp/src/arrow/ipc/read_write_test.cc:
##########
@@ -1954,6 +2023,22 @@ TEST_P(TestFileFormatGeneratorCoalesced, RoundTrip) {
TEST_P(TestStreamFormat, RoundTrip) { TestRoundTripWithOptions(*GetParam()); }
+TEST_F(TestFileFormat, DenseUnionExtensionRoundTrip) {
Review Comment:
Why not extend `kBatchCases` instead of adding dedicated test functions?
##########
cpp/src/arrow/ipc/writer.cc:
##########
@@ -164,18 +170,19 @@ class RecordBatchSerializer {
// In V4, null types have no validity bitmap
// In V5 and later, null and union types have no validity bitmap
- if (internal::HasValidityBitmap(arr.type_id(), options_.metadata_version))
{
+ if (internal::HasValidityBitmap(physical_arr->type_id(),
options_.metadata_version)) {
if (arr.null_count() > 0) {
Review Comment:
This is still using `arr` while switching to `physical_arr` in other places
(e.g. `arr.length()` above vs. `physical_arr->length()` below), with no obvious
rationale. Can we make the code more consistent?
##########
cpp/src/arrow/ipc/read_write_test.cc:
##########
@@ -71,6 +71,73 @@ using MetadataVector =
std::vector<std::shared_ptr<KeyValueMetadata>>;
namespace test {
+class UnionExtensionArray : public ExtensionArray {
+ public:
+ using ExtensionArray::ExtensionArray;
+};
+
+class UnionExtensionType : public ExtensionType {
+ public:
+ UnionExtensionType(std::shared_ptr<DataType> storage_type, std::string
extension_name)
+ : ExtensionType(std::move(storage_type)),
+ extension_name_(std::move(extension_name)) {}
+
+ std::string extension_name() const override { return extension_name_; }
+
+ bool ExtensionEquals(const ExtensionType& other) const override {
+ return other.extension_name() == extension_name_;
+ }
+
+ std::shared_ptr<Array> MakeArray(std::shared_ptr<ArrayData> data) const
override {
+ return std::make_shared<UnionExtensionArray>(std::move(data));
+ }
+
+ Result<std::shared_ptr<DataType>> Deserialize(
+ std::shared_ptr<DataType> storage_type,
+ const std::string& serialized) const override {
+ if (serialized != extension_name_) {
+ return Status::Invalid("Type identifier did not match");
+ }
+ return std::make_shared<UnionExtensionType>(std::move(storage_type),
+ extension_name_);
+ }
+
+ std::string Serialize() const override { return extension_name_; }
+
+ private:
+ std::string extension_name_;
+};
+
+std::shared_ptr<DataType> dense_union_extension_type() {
+ return std::make_shared<UnionExtensionType>(
+ dense_union({field("floats", float64()), field("strings",
large_utf8())}, {0, 1}),
+ "dense-union-extension");
+}
+
+std::shared_ptr<DataType> sparse_union_extension_type() {
+ return std::make_shared<UnionExtensionType>(
+ sparse_union({field("floats", float64()), field("strings",
large_utf8())}, {0, 1}),
+ "sparse-union-extension");
+}
+
+Status MakeDenseUnionExtension(std::shared_ptr<RecordBatch>* out) {
Review Comment:
Perhaps move these two functions to `arrow/ipc/test_common.h`?
--
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.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
