diff --git a/cpp/src/generated/parquet_types.cpp b/cpp/src/generated/parquet_types.cpp index 0ee973f2a2d6..531d9a84e0c3 100644 --- a/cpp/src/generated/parquet_types.cpp +++ b/cpp/src/generated/parquet_types.cpp @@ -335,7 +335,62 @@ int _kFieldRepetitionTypeValues[] = { /** * The field is repeated and can contain 0 or more values */ - FieldRepetitionType::REPEATED + FieldRepetitionType::REPEATED, + /** + * EXPERIMENTAL: The field repeats a fixed number of times per parent value. + * + * A VECTOR field repeats exactly vector_length times for every parent value + * and, unlike REPEATED, does not increase the maximum definition or + * repetition level of its descendants: readers reconstruct vector values + * from the fixed multiplicity declared in the schema instead of decoding + * repetition levels. + * + * Vector fields MUST use the following 3-level structure, mirroring LIST: + * + * group (VECTOR) { + * vector group list [vector_length] { + * element; + * } + * } + * + * - The outer group MUST be REQUIRED or OPTIONAL and carries the + * nullability of the vector value itself. It MUST be annotated with the + * VECTOR logical type and MUST have exactly one child: the + * VECTOR-repeated group. + * - The VECTOR-repeated middle group carries vector_length and MUST have + * exactly one child: the element field. + * - The element field MUST be REQUIRED or OPTIONAL and MAY be a primitive + * type, a group of fields (a struct), or another vector group in the same + * form. For nested vectors a leaf stores the product of the + * vector_length values of its VECTOR ancestors physical values per parent + * record. + * + * Data layout rules: + * - Writers MUST emit exactly vector_length child slots for every parent + * record in which the vector value is present or null; null vector values + * emit vector_length slots carrying the definition level of the null + * vector. Ancestors that make the vector value absent altogether (for + * example an empty or null list) contribute a single level entry with + * their usual definition level and no slots. Column chunk num_values and + * null counts therefore include the padding slots of null vector values. + * - Every slot of an existing vector occupies one position in both level + * streams. The repetition level of the first slot encodes record + * structure as usual; the repetition levels of the remaining + * vector_length - 1 slots MUST be written as the column's maximum + * repetition level and MUST be ignored by readers. Stream consumption is + * definition-driven: after reading a position whose definition level + * indicates an existing vector, a reader consumes that position and the + * following vector_length - 1 positions as one vector value. + * - Writers MUST NOT split the slots of one vector value across data pages; + * pages begin and end on whole-vector boundaries. + * - vector_length MUST be positive: a zero-length vector value would + * contribute no slots at all, leaving row counts and row nullability + * unrepresentable in this layout. Writers MUST represent zero-length + * fixed-size lists with the LIST encoding instead. + * + * Readers that do not understand VECTOR are expected to reject the file. + */ + FieldRepetitionType::VECTOR }; const char* _kFieldRepetitionTypeNames[] = { /** @@ -349,9 +404,64 @@ const char* _kFieldRepetitionTypeNames[] = { /** * The field is repeated and can contain 0 or more values */ - "REPEATED" + "REPEATED", + /** + * EXPERIMENTAL: The field repeats a fixed number of times per parent value. + * + * A VECTOR field repeats exactly vector_length times for every parent value + * and, unlike REPEATED, does not increase the maximum definition or + * repetition level of its descendants: readers reconstruct vector values + * from the fixed multiplicity declared in the schema instead of decoding + * repetition levels. + * + * Vector fields MUST use the following 3-level structure, mirroring LIST: + * + * group (VECTOR) { + * vector group list [vector_length] { + * element; + * } + * } + * + * - The outer group MUST be REQUIRED or OPTIONAL and carries the + * nullability of the vector value itself. It MUST be annotated with the + * VECTOR logical type and MUST have exactly one child: the + * VECTOR-repeated group. + * - The VECTOR-repeated middle group carries vector_length and MUST have + * exactly one child: the element field. + * - The element field MUST be REQUIRED or OPTIONAL and MAY be a primitive + * type, a group of fields (a struct), or another vector group in the same + * form. For nested vectors a leaf stores the product of the + * vector_length values of its VECTOR ancestors physical values per parent + * record. + * + * Data layout rules: + * - Writers MUST emit exactly vector_length child slots for every parent + * record in which the vector value is present or null; null vector values + * emit vector_length slots carrying the definition level of the null + * vector. Ancestors that make the vector value absent altogether (for + * example an empty or null list) contribute a single level entry with + * their usual definition level and no slots. Column chunk num_values and + * null counts therefore include the padding slots of null vector values. + * - Every slot of an existing vector occupies one position in both level + * streams. The repetition level of the first slot encodes record + * structure as usual; the repetition levels of the remaining + * vector_length - 1 slots MUST be written as the column's maximum + * repetition level and MUST be ignored by readers. Stream consumption is + * definition-driven: after reading a position whose definition level + * indicates an existing vector, a reader consumes that position and the + * following vector_length - 1 positions as one vector value. + * - Writers MUST NOT split the slots of one vector value across data pages; + * pages begin and end on whole-vector boundaries. + * - vector_length MUST be positive: a zero-length vector value would + * contribute no slots at all, leaving row counts and row nullability + * unrepresentable in this layout. Writers MUST represent zero-length + * fixed-size lists with the LIST encoding instead. + * + * Readers that do not understand VECTOR are expected to reject the file. + */ + "VECTOR" }; -const std::map _FieldRepetitionType_VALUES_TO_NAMES(::apache::thrift::TEnumIterator(3, _kFieldRepetitionTypeValues, _kFieldRepetitionTypeNames), ::apache::thrift::TEnumIterator(-1, nullptr, nullptr)); +const std::map _FieldRepetitionType_VALUES_TO_NAMES(::apache::thrift::TEnumIterator(4, _kFieldRepetitionTypeValues, _kFieldRepetitionTypeNames), ::apache::thrift::TEnumIterator(-1, nullptr, nullptr)); std::ostream& operator<<(std::ostream& out, const FieldRepetitionType::type& val) { std::map::const_iterator it = _FieldRepetitionType_VALUES_TO_NAMES.find(val); @@ -2345,6 +2455,11 @@ void LogicalType::__set_GEOGRAPHY(const GeographyType& val) { this->GEOGRAPHY = val; __isset.GEOGRAPHY = true; } + +void LogicalType::__set_VECTOR(const VectorType& val) { + this->VECTOR = val; +__isset.VECTOR = true; +} std::ostream& operator<<(std::ostream& out, const LogicalType& obj) { obj.printTo(out); @@ -2371,6 +2486,7 @@ void swap(LogicalType &a, LogicalType &b) { swap(a.VARIANT, b.VARIANT); swap(a.GEOMETRY, b.GEOMETRY); swap(a.GEOGRAPHY, b.GEOGRAPHY); + swap(a.VECTOR, b.VECTOR); swap(a.__isset, b.__isset); } @@ -2444,6 +2560,10 @@ bool LogicalType::operator==(const LogicalType & rhs) const return false; else if (__isset.GEOGRAPHY && !(GEOGRAPHY == rhs.GEOGRAPHY)) return false; + if (__isset.VECTOR != rhs.__isset.VECTOR) + return false; + else if (__isset.VECTOR && !(VECTOR == rhs.VECTOR)) + return false; return true; } @@ -2465,6 +2585,7 @@ LogicalType::LogicalType(const LogicalType& other119) { VARIANT = other119.VARIANT; GEOMETRY = other119.GEOMETRY; GEOGRAPHY = other119.GEOGRAPHY; + VECTOR = other119.VECTOR; __isset = other119.__isset; } LogicalType::LogicalType(LogicalType&& other120) noexcept { @@ -2485,6 +2606,7 @@ LogicalType::LogicalType(LogicalType&& other120) noexcept { VARIANT = std::move(other120.VARIANT); GEOMETRY = std::move(other120.GEOMETRY); GEOGRAPHY = std::move(other120.GEOGRAPHY); + VECTOR = std::move(other120.VECTOR); __isset = other120.__isset; } LogicalType& LogicalType::operator=(const LogicalType& other121) { @@ -2505,6 +2627,7 @@ LogicalType& LogicalType::operator=(const LogicalType& other121) { VARIANT = other121.VARIANT; GEOMETRY = other121.GEOMETRY; GEOGRAPHY = other121.GEOGRAPHY; + VECTOR = other121.VECTOR; __isset = other121.__isset; return *this; } @@ -2526,6 +2649,7 @@ LogicalType& LogicalType::operator=(LogicalType&& other122) noexcept { VARIANT = std::move(other122.VARIANT); GEOMETRY = std::move(other122.GEOMETRY); GEOGRAPHY = std::move(other122.GEOGRAPHY); + VECTOR = std::move(other122.VECTOR); __isset = other122.__isset; return *this; } @@ -2549,6 +2673,7 @@ void LogicalType::printTo(std::ostream& out) const { out << ", " << "VARIANT="; (__isset.VARIANT ? (out << to_string(VARIANT)) : (out << "")); out << ", " << "GEOMETRY="; (__isset.GEOMETRY ? (out << to_string(GEOMETRY)) : (out << "")); out << ", " << "GEOGRAPHY="; (__isset.GEOGRAPHY ? (out << to_string(GEOGRAPHY)) : (out << "")); + out << ", " << "VECTOR="; (__isset.VECTOR ? (out << to_string(VECTOR)) : (out << "")); out << ")"; } @@ -2565,7 +2690,8 @@ SchemaElement::SchemaElement() noexcept converted_type(static_cast(0)), scale(0), precision(0), - field_id(0) { + field_id(0), + vector_length(0) { } void SchemaElement::__set_type(const Type::type val) { @@ -2616,6 +2742,11 @@ void SchemaElement::__set_logicalType(const LogicalType& val) { this->logicalType = val; __isset.logicalType = true; } + +void SchemaElement::__set_vector_length(const int32_t val) { + this->vector_length = val; +__isset.vector_length = true; +} std::ostream& operator<<(std::ostream& out, const SchemaElement& obj) { obj.printTo(out); @@ -2635,6 +2766,7 @@ void swap(SchemaElement &a, SchemaElement &b) { swap(a.precision, b.precision); swap(a.field_id, b.field_id); swap(a.logicalType, b.logicalType); + swap(a.vector_length, b.vector_length); swap(a.__isset, b.__isset); } @@ -2678,6 +2810,10 @@ bool SchemaElement::operator==(const SchemaElement & rhs) const return false; else if (__isset.logicalType && !(logicalType == rhs.logicalType)) return false; + if (__isset.vector_length != rhs.__isset.vector_length) + return false; + else if (__isset.vector_length && !(vector_length == rhs.vector_length)) + return false; return true; } @@ -2692,6 +2828,7 @@ SchemaElement::SchemaElement(const SchemaElement& other126) { precision = other126.precision; field_id = other126.field_id; logicalType = other126.logicalType; + vector_length = other126.vector_length; __isset = other126.__isset; } SchemaElement::SchemaElement(SchemaElement&& other127) noexcept { @@ -2705,6 +2842,7 @@ SchemaElement::SchemaElement(SchemaElement&& other127) noexcept { precision = other127.precision; field_id = other127.field_id; logicalType = std::move(other127.logicalType); + vector_length = other127.vector_length; __isset = other127.__isset; } SchemaElement& SchemaElement::operator=(const SchemaElement& other128) { @@ -2718,6 +2856,7 @@ SchemaElement& SchemaElement::operator=(const SchemaElement& other128) { precision = other128.precision; field_id = other128.field_id; logicalType = other128.logicalType; + vector_length = other128.vector_length; __isset = other128.__isset; return *this; } @@ -2732,6 +2871,7 @@ SchemaElement& SchemaElement::operator=(SchemaElement&& other129) noexcept { precision = other129.precision; field_id = other129.field_id; logicalType = std::move(other129.logicalType); + vector_length = other129.vector_length; __isset = other129.__isset; return *this; } @@ -2748,6 +2888,7 @@ void SchemaElement::printTo(std::ostream& out) const { out << ", " << "precision="; (__isset.precision ? (out << to_string(precision)) : (out << "")); out << ", " << "field_id="; (__isset.field_id ? (out << to_string(field_id)) : (out << "")); out << ", " << "logicalType="; (__isset.logicalType ? (out << to_string(logicalType)) : (out << "")); + out << ", " << "vector_length="; (__isset.vector_length ? (out << to_string(vector_length)) : (out << "")); out << ")"; } diff --git a/cpp/src/generated/parquet_types.h b/cpp/src/generated/parquet_types.h index 1f1e254f5cf2..a64780651852 100644 --- a/cpp/src/generated/parquet_types.h +++ b/cpp/src/generated/parquet_types.h @@ -202,7 +202,62 @@ struct FieldRepetitionType { /** * The field is repeated and can contain 0 or more values */ - REPEATED = 2 + REPEATED = 2, + /** + * EXPERIMENTAL: The field repeats a fixed number of times per parent value. + * + * A VECTOR field repeats exactly vector_length times for every parent value + * and, unlike REPEATED, does not increase the maximum definition or + * repetition level of its descendants: readers reconstruct vector values + * from the fixed multiplicity declared in the schema instead of decoding + * repetition levels. + * + * Vector fields MUST use the following 3-level structure, mirroring LIST: + * + * group (VECTOR) { + * vector group list [vector_length] { + * element; + * } + * } + * + * - The outer group MUST be REQUIRED or OPTIONAL and carries the + * nullability of the vector value itself. It MUST be annotated with the + * VECTOR logical type and MUST have exactly one child: the + * VECTOR-repeated group. + * - The VECTOR-repeated middle group carries vector_length and MUST have + * exactly one child: the element field. + * - The element field MUST be REQUIRED or OPTIONAL and MAY be a primitive + * type, a group of fields (a struct), or another vector group in the same + * form. For nested vectors a leaf stores the product of the + * vector_length values of its VECTOR ancestors physical values per parent + * record. + * + * Data layout rules: + * - Writers MUST emit exactly vector_length child slots for every parent + * record in which the vector value is present or null; null vector values + * emit vector_length slots carrying the definition level of the null + * vector. Ancestors that make the vector value absent altogether (for + * example an empty or null list) contribute a single level entry with + * their usual definition level and no slots. Column chunk num_values and + * null counts therefore include the padding slots of null vector values. + * - Every slot of an existing vector occupies one position in both level + * streams. The repetition level of the first slot encodes record + * structure as usual; the repetition levels of the remaining + * vector_length - 1 slots MUST be written as the column's maximum + * repetition level and MUST be ignored by readers. Stream consumption is + * definition-driven: after reading a position whose definition level + * indicates an existing vector, a reader consumes that position and the + * following vector_length - 1 positions as one vector value. + * - Writers MUST NOT split the slots of one vector value across data pages; + * pages begin and end on whole-vector boundaries. + * - vector_length MUST be positive: a zero-length vector value would + * contribute no slots at all, leaving row counts and row nullability + * unrepresentable in this layout. Writers MUST represent zero-length + * fixed-size lists with the LIST encoding instead. + * + * Readers that do not understand VECTOR are expected to reject the file. + */ + VECTOR = 3 }; }; @@ -367,6 +422,23 @@ std::ostream& operator<<(std::ostream& out, const BoundaryOrder::type& val); std::string to_string(const BoundaryOrder::type& val); +/** + * Embedded Vector logical type annotation + * + * EXPERIMENTAL: annotates the outer group of a fixed-size vector field; see + * FieldRepetitionType.VECTOR for the full structure and data layout rules. + * The annotated group MUST be REQUIRED or OPTIONAL and MUST have exactly one + * child: a VECTOR-repeated group whose SchemaElement.vector_length carries + * the fixed multiplicity. The annotation is required so that readers can + * distinguish a vector field from a struct-like group that happens to + * contain a VECTOR-repeated child. + * + * Allowed for group nodes only. Defined as a typedef of the (empty) + * NullType struct to keep the generated code small; a format-level proposal + * would define a distinct empty struct. + */ +typedef class NullType VectorType; + class SizeStatistics; class BoundingBox; @@ -1621,7 +1693,7 @@ void swap(GeographyType &a, GeographyType &b); std::ostream& operator<<(std::ostream& out, const GeographyType& obj); typedef struct _LogicalType__isset { - _LogicalType__isset() : STRING(false), MAP(false), LIST(false), ENUM(false), DECIMAL(false), DATE(false), TIME(false), TIMESTAMP(false), INTEGER(false), UNKNOWN(false), JSON(false), BSON(false), UUID(false), FLOAT16(false), VARIANT(false), GEOMETRY(false), GEOGRAPHY(false) {} + _LogicalType__isset() : STRING(false), MAP(false), LIST(false), ENUM(false), DECIMAL(false), DATE(false), TIME(false), TIMESTAMP(false), INTEGER(false), UNKNOWN(false), JSON(false), BSON(false), UUID(false), FLOAT16(false), VARIANT(false), GEOMETRY(false), GEOGRAPHY(false), VECTOR(false) {} bool STRING :1; bool MAP :1; bool LIST :1; @@ -1639,6 +1711,7 @@ typedef struct _LogicalType__isset { bool VARIANT :1; bool GEOMETRY :1; bool GEOGRAPHY :1; + bool VECTOR :1; } _LogicalType__isset; /** @@ -1675,6 +1748,7 @@ class LogicalType { VariantType VARIANT; GeometryType GEOMETRY; GeographyType GEOGRAPHY; + VectorType VECTOR; _LogicalType__isset __isset; @@ -1712,6 +1786,8 @@ class LogicalType { void __set_GEOGRAPHY(const GeographyType& val); + void __set_VECTOR(const VectorType& val); + bool operator == (const LogicalType & rhs) const; bool operator != (const LogicalType &rhs) const { return !(*this == rhs); @@ -1732,7 +1808,7 @@ void swap(LogicalType &a, LogicalType &b); std::ostream& operator<<(std::ostream& out, const LogicalType& obj); typedef struct _SchemaElement__isset { - _SchemaElement__isset() : type(false), type_length(false), repetition_type(false), num_children(false), converted_type(false), scale(false), precision(false), field_id(false), logicalType(false) {} + _SchemaElement__isset() : type(false), type_length(false), repetition_type(false), num_children(false), converted_type(false), scale(false), precision(false), field_id(false), logicalType(false), vector_length(false) {} bool type :1; bool type_length :1; bool repetition_type :1; @@ -1742,6 +1818,7 @@ typedef struct _SchemaElement__isset { bool precision :1; bool field_id :1; bool logicalType :1; + bool vector_length :1; } _SchemaElement__isset; /** @@ -1820,6 +1897,17 @@ class SchemaElement { * for some logical types to ensure forward-compatibility in format v1. */ LogicalType logicalType; + /** + * EXPERIMENTAL: The fixed number of times the field repeats per parent + * value. + * + * MUST be set, and positive, when repetition_type is VECTOR; MUST NOT be + * set otherwise. Zero-length vectors are not representable as VECTOR and + * use the LIST encoding (see FieldRepetitionType.VECTOR). For nested + * VECTOR fields the number of physical leaf values per parent record is + * the product of vector_length over the leaf's VECTOR ancestors. + */ + int32_t vector_length; _SchemaElement__isset __isset; @@ -1843,6 +1931,8 @@ class SchemaElement { void __set_logicalType(const LogicalType& val); + void __set_vector_length(const int32_t val); + bool operator == (const SchemaElement & rhs) const; bool operator != (const SchemaElement &rhs) const { return !(*this == rhs); diff --git a/cpp/src/generated/parquet_types.tcc b/cpp/src/generated/parquet_types.tcc index 78e3e2549394..c28f214d450b 100644 --- a/cpp/src/generated/parquet_types.tcc +++ b/cpp/src/generated/parquet_types.tcc @@ -1783,6 +1783,14 @@ uint32_t LogicalType::read(Protocol_* iprot) { xfer += iprot->skip(ftype); } break; + case 19: + if (ftype == ::apache::thrift::protocol::T_STRUCT) { + xfer += this->VECTOR.read(iprot); + this->__isset.VECTOR = true; + } else { + xfer += iprot->skip(ftype); + } + break; default: xfer += iprot->skip(ftype); break; @@ -1886,6 +1894,11 @@ uint32_t LogicalType::write(Protocol_* oprot) const { xfer += this->GEOGRAPHY.write(oprot); xfer += oprot->writeFieldEnd(); } + if (this->__isset.VECTOR) { + xfer += oprot->writeFieldBegin("VECTOR", ::apache::thrift::protocol::T_STRUCT, 19); + xfer += this->VECTOR.write(oprot); + xfer += oprot->writeFieldEnd(); + } xfer += oprot->writeFieldStop(); xfer += oprot->writeStructEnd(); return xfer; @@ -2000,6 +2013,14 @@ uint32_t SchemaElement::read(Protocol_* iprot) { xfer += iprot->skip(ftype); } break; + case 12: + if (ftype == ::apache::thrift::protocol::T_I32) { + xfer += iprot->readI32(this->vector_length); + this->__isset.vector_length = true; + } else { + xfer += iprot->skip(ftype); + } + break; default: xfer += iprot->skip(ftype); break; @@ -2069,6 +2090,11 @@ uint32_t SchemaElement::write(Protocol_* oprot) const { xfer += this->logicalType.write(oprot); xfer += oprot->writeFieldEnd(); } + if (this->__isset.vector_length) { + xfer += oprot->writeFieldBegin("vector_length", ::apache::thrift::protocol::T_I32, 12); + xfer += oprot->writeI32(this->vector_length); + xfer += oprot->writeFieldEnd(); + } xfer += oprot->writeFieldStop(); xfer += oprot->writeStructEnd(); return xfer; diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index d29458bf226b..232b55d8d8d4 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -24,10 +24,13 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" +#include #include #include #include #include +#include +#include #include #include @@ -64,6 +67,8 @@ #include "parquet/api/reader.h" #include "parquet/api/writer.h" +#include "parquet/bloom_filter.h" +#include "parquet/bloom_filter_reader.h" #include "parquet/arrow/fuzz_internal.h" #include "parquet/arrow/reader.h" @@ -73,6 +78,7 @@ #include "parquet/arrow/writer.h" #include "parquet/column_writer.h" #include "parquet/file_writer.h" +#include "parquet/page_index.h" #include "parquet/properties.h" #include "parquet/test_util.h" #include "parquet/types.h" @@ -421,6 +427,19 @@ void WriteTableToBuffer(const std::shared_ptr& table, int64_t row_group_s *out, WriteTableToBuffer(table, row_group_size, write_props, arrow_properties)); } +std::shared_ptr VectorWriterProperties() { + auto builder = WriterProperties::Builder(); + return builder.enable_write_page_index()->encoding(Encoding::PLAIN)->build(); +} + +std::shared_ptr VectorByteStreamSplitWriterProperties() { + auto builder = WriterProperties::Builder(); + return builder.disable_dictionary() + ->enable_write_page_index() + ->encoding(Encoding::BYTE_STREAM_SPLIT) + ->build(); +} + void DoRoundtrip(const std::shared_ptr
& table, int64_t row_group_size, std::shared_ptr
* out, const std::shared_ptr<::parquet::WriterProperties>& writer_properties = @@ -442,6 +461,62 @@ void DoRoundtrip(const std::shared_ptr
& table, int64_t row_group_size, ASSERT_OK_AND_ASSIGN(*out, reader->ReadTable()); } +std::shared_ptr<::arrow::DataType> MakeVectorFixedSizeListType( + const std::shared_ptr<::arrow::DataType>& item_type, bool element_nullable = false) { + return ::arrow::fixed_size_list(::arrow::field("item", item_type, element_nullable), + /*size=*/3); +} + +std::shared_ptr<::arrow::DataType> VectorFixedSizeListType( + bool element_nullable = false) { + return MakeVectorFixedSizeListType(::arrow::int16(), element_nullable); +} + +std::shared_ptr<::arrow::DataType> VectorFloatFixedSizeListType( + bool element_nullable = false) { + return MakeVectorFixedSizeListType(::arrow::float32(), element_nullable); +} + +std::shared_ptr
MakeVectorFixedSizeListTable( + const std::vector>& chunks, bool nullable = true, + bool element_nullable = false) { + auto type = VectorFixedSizeListType(element_nullable); + auto field = ::arrow::field("root", type, nullable); + auto column = std::make_shared(chunks, type); + return ::arrow::Table::Make(::arrow::schema({field}), {column}); +} + +std::shared_ptr
MakeVectorFixedSizeListTable(std::string_view json, + bool nullable = true, + bool element_nullable = false) { + auto type = VectorFixedSizeListType(element_nullable); + return MakeVectorFixedSizeListTable({::arrow::ArrayFromJSON(type, std::string(json))}, + nullable, element_nullable); +} + +std::shared_ptr
MakeVectorFixedSizeListTable( + const std::shared_ptr<::arrow::DataType>& item_type, std::string_view json, + bool nullable = true, bool element_nullable = false) { + auto type = MakeVectorFixedSizeListType(item_type, element_nullable); + auto field = ::arrow::field("root", type, nullable); + auto array = ::arrow::ArrayFromJSON(type, std::string(json)); + return ::arrow::Table::Make(::arrow::schema({field}), {array}); +} + +void CheckVectorFixedSizeListRoundtrip( + const std::shared_ptr
& table, int64_t row_group_size, + const ArrowReaderProperties& arrow_reader_properties = + default_arrow_reader_properties()) { + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + std::shared_ptr
result; + ASSERT_NO_FATAL_FAILURE(DoRoundtrip(table, row_group_size, &result, + VectorWriterProperties(), builder.build(), + arrow_reader_properties)); + ::arrow::AssertSchemaEqual(*table->schema(), *result->schema(), false); + ::arrow::AssertTablesEqual(*table, *result, false); +} + void CheckConfiguredRoundtrip( const std::shared_ptr
& input_table, const std::shared_ptr
& expected_table = nullptr, @@ -551,6 +626,18 @@ void CheckSimpleRoundtrip(const std::shared_ptr
& table, int64_t row_group ::arrow::AssertTablesEqual(*table, *result, false); } +// Roundtrip for tables whose FixedSizeList columns fall back to the standard +// LIST encoding even though vector encoding is enabled. store_schema() is +// needed to restore the FixedSizeList type from the LIST encoding on read. +void CheckVectorFallsBackToListRoundtrip(const std::shared_ptr
& table, + int64_t row_group_size) { + auto arrow_writer_properties = ArrowWriterProperties::Builder() + .store_schema() + ->enable_experimental_vector_encoding() + ->build(); + CheckSimpleRoundtrip(table, row_group_size, arrow_writer_properties); +} + static std::shared_ptr MakeSimpleSchema(const DataType& type, Repetition::type repetition) { int32_t byte_width = -1; @@ -3325,22 +3412,800 @@ TEST(ArrowReadWrite, LargeList) { } TEST(ArrowReadWrite, FixedSizeList) { - using ::arrow::field; - using ::arrow::fixed_size_list; - using ::arrow::struct_; - - auto type = fixed_size_list(::arrow::int16(), /*size=*/3); - - const char* json = R"([ + auto table = MakeVectorFixedSizeListTable(R"([ [1, 2, 3], [4, 5, 6], - [7, 8, 9]])"; - auto array = ::arrow::ArrayFromJSON(type, json); - auto table = ::arrow::Table::Make(::arrow::schema({field("root", type)}), {array}); + [7, 8, 9]])", + /*nullable=*/true, + /*element_nullable=*/true); auto props_store_schema = ArrowWriterProperties::Builder().store_schema()->build(); CheckSimpleRoundtrip(table, 2, props_store_schema); } +TEST(ArrowWriteOnly, FixedSizeListVectorSchemaRequired) { + auto table = MakeVectorFixedSizeListTable(R"([ + [1, 2, 3], + [4, 5, 6], + [7, 8, 9]])", + /*nullable=*/false); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/2, + VectorWriterProperties(), builder.build())); + + auto reader = ParquetFileReader::Open(std::make_shared(buffer)); + const auto* schema = reader->metadata()->schema(); + const auto* root = schema->group_node()->field(0).get(); + ASSERT_TRUE(root->is_group()); + const auto* root_group = static_cast(root); + ASSERT_EQ(root_group->repetition(), Repetition::REQUIRED); + ASSERT_TRUE(root_group->logical_type() != nullptr && + root_group->logical_type()->is_vector()); + ASSERT_EQ(root_group->field_count(), 1); + const auto* vector_node = root_group->field(0).get(); + ASSERT_TRUE(vector_node->is_vector()); + ASSERT_TRUE(vector_node->is_group()); + ASSERT_EQ(vector_node->vector_length(), 3); + const auto* vector_group = static_cast(vector_node); + ASSERT_EQ(vector_group->field_count(), 1); + ASSERT_TRUE(vector_group->field(0)->is_primitive()); + ASSERT_TRUE(vector_group->field(0)->is_required()); + ASSERT_EQ(reader->metadata()->RowGroup(0)->ColumnChunk(0)->num_values(), 6); + ASSERT_EQ(reader->metadata()->RowGroup(1)->ColumnChunk(0)->num_values(), 3); +} + +TEST(ArrowWriteOnly, FixedSizeListVectorSchemaNullable) { + auto table = MakeVectorFixedSizeListTable(R"([ + [1, 2, 3], + null, + [7, 8, 9]])"); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/3, + VectorWriterProperties(), builder.build())); + + auto reader = ParquetFileReader::Open(std::make_shared(buffer)); + const auto* schema = reader->metadata()->schema(); + const auto* root = schema->group_node()->field(0).get(); + ASSERT_TRUE(root->is_group()); + const auto* root_group = static_cast(root); + ASSERT_EQ(root_group->repetition(), Repetition::OPTIONAL); + ASSERT_TRUE(root_group->logical_type() != nullptr && + root_group->logical_type()->is_vector()); + ASSERT_EQ(root_group->field_count(), 1); + const auto* vector_node = root_group->field(0).get(); + ASSERT_TRUE(vector_node->is_vector()); + ASSERT_TRUE(vector_node->is_group()); + ASSERT_EQ(vector_node->vector_length(), 3); + const auto* vector_group = static_cast(vector_node); + ASSERT_EQ(vector_group->field_count(), 1); + ASSERT_TRUE(vector_group->field(0)->is_primitive()); + ASSERT_TRUE(vector_group->field(0)->is_required()); + // Option B uses fixed-stride nullable parent layout: every parent row, + // including a null vector row, contributes vector_length child slots. + ASSERT_EQ(reader->metadata()->RowGroup(0)->ColumnChunk(0)->num_values(), 9); +} + +TEST(ArrowWriteOnly, FixedSizeListVectorNullRowsEmitFixedWidthNullSlots) { + auto table = MakeVectorFixedSizeListTable(R"([ + null, + null, + null])"); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/3, + VectorWriterProperties(), builder.build())); + + auto reader = ParquetFileReader::Open(std::make_shared(buffer)); + ASSERT_EQ(reader->metadata()->RowGroup(0)->num_rows(), 3); + ASSERT_EQ(reader->metadata()->RowGroup(0)->ColumnChunk(0)->num_values(), 9); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableRoundTripNullPatterns) { + const std::vector cases = { + R"([[1, 2, 3], [4, 5, 6], [7, 8, 9]])", + R"([[1, 2, 3], null, [7, 8, 9]])", + R"([null, [4, 5, 6], [7, 8, 9]])", + R"([[1, 2, 3], [4, 5, 6], null])", + R"([[1, 2, 3], null, null, [10, 11, 12], null])", + R"([null, null, null])", + R"([])"}; + + for (const auto& json : cases) { + SCOPED_TRACE(json); + auto table = MakeVectorFixedSizeListTable(json); + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip( + table, std::max(1, table->num_rows()))); + } +} + +TEST(ArrowReadWrite, FixedSizeListVectorStructOfVectorsRoundTrip) { + auto ids_type = ::arrow::fixed_size_list( + ::arrow::field("value", ::arrow::int32(), /*nullable=*/true), 2); + auto scores_type = ::arrow::fixed_size_list(::arrow::float32(), 3); + auto type = ::arrow::struct_({::arrow::field("ids", ids_type, /*nullable=*/true), + ::arrow::field("scores", scores_type, + /*nullable=*/false)}); + auto array = ::arrow::ArrayFromJSON(type, R"([ + {"ids": [10, null], "scores": [0.1, 0.2, 0.3]}, + {"ids": null, "scores": [1.1, 1.2, 1.3]}, + {"ids": [30, 40], "scores": [2.1, 2.2, 2.3]}])"); + auto table = ::arrow::Table::Make( + ::arrow::schema({::arrow::field("root", type, /*nullable=*/false)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableStructAboveVectorFallsBackToList) { + // A null row of a nullable struct above a vector would map one definition + // level to vector_length leaf slots, which the VECTOR level machinery does + // not support; vectors below nullable groups fall back to the LIST encoding. + auto vector_type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::int32(), /*nullable=*/false), 2); + auto struct_type = + ::arrow::struct_({::arrow::field("v", vector_type, /*nullable=*/false)}); + // No null struct rows: writing those through the LIST encoding is an + // upstream NotImplemented ("Lists with non-zero length null components") + // independent of VECTOR. + auto array = ::arrow::ArrayFromJSON(struct_type, R"([ + {"v": [1, 2]}, + {"v": [3, 4]}, + {"v": [5, 6]}])"); + auto table = ::arrow::Table::Make( + ::arrow::schema({::arrow::field("root", struct_type, /*nullable=*/true)}), {array}); + ASSERT_NO_FATAL_FAILURE(CheckVectorFallsBackToListRoundtrip(table, + /*row_group_size=*/3)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableStructUnderListFallsBackToList) { + auto vector_type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::float32(), /*nullable=*/false), 3); + auto struct_type = + ::arrow::struct_({::arrow::field("embedding", vector_type, /*nullable=*/false)}); + auto type = ::arrow::list(::arrow::field("item", struct_type, /*nullable=*/true)); + // No null struct items: writing those through the LIST encoding is an + // upstream NotImplemented ("Lists with non-zero length null components") + // independent of VECTOR. + auto array = ::arrow::ArrayFromJSON(type, R"([ + [{"embedding": [1.0, 2.0, 3.0]}, {"embedding": [4.0, 5.0, 6.0]}], + [{"embedding": [7.0, 8.0, 9.0]}]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + ASSERT_NO_FATAL_FAILURE(CheckVectorFallsBackToListRoundtrip(table, + /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorStructWithListFallsBackToListRoundTrip) { + // The struct element contains a variable-size list, so the column cannot use + // VECTOR encoding and must fall back to the standard LIST encoding. The + // schema conversion and the write-path level generation must agree on that + // fallback, otherwise the writer emits VECTOR-style levels for a LIST + // schema. + auto items_type = ::arrow::list(::arrow::field("item", ::arrow::int32(), + /*nullable=*/false)); + auto struct_type = + ::arrow::struct_({::arrow::field("items", items_type, /*nullable=*/false)}); + auto type = ::arrow::fixed_size_list( + ::arrow::field("element", struct_type, /*nullable=*/false), 2); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [{"items": [1, 2]}, {"items": [3]}], + [{"items": []}, {"items": [4, 5, 6]}]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFallsBackToListRoundtrip(table, + /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorListOfNullableVectorsRoundTrip) { + // Null and present vector items interleaved within and across lists. + auto vector_type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::int32(), /*nullable=*/false), 2); + auto type = ::arrow::list(::arrow::field("item", vector_type, /*nullable=*/true)); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [[1, 2], null, [3, 4]], + [null], + [], + null, + [[5, 6], [7, 8], null, [9, 10]]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, + /*row_group_size=*/3)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNestedVectorsUnderListRoundTrip) { + auto inner = ::arrow::fixed_size_list( + ::arrow::field("value", ::arrow::int32(), /*nullable=*/false), 2); + auto outer = + ::arrow::fixed_size_list(::arrow::field("element", inner, /*nullable=*/false), 2); + auto type = ::arrow::list(::arrow::field("item", outer, /*nullable=*/true)); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [[[1, 2], [3, 4]], null], + [], + null, + [[[5, 6], [7, 8]]]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, + /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorListOfVectorsRoundTrip) { + // FixedSizeList values below a repeated field can use VECTOR encoding; + // empty and null list entries must roundtrip. + auto vector_type = ::arrow::fixed_size_list(::arrow::int32(), 2); + auto type = ::arrow::list(::arrow::field("item", vector_type, /*nullable=*/false)); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [[1, 2], [3, 4]], + [], + null, + [[5, 6]], + [[7, 8], [9, 10], [11, 12]]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, + /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorListOfStructsWithVectorsRoundTrip) { + auto vector_type = ::arrow::fixed_size_list(::arrow::float32(), 3); + auto struct_type = + ::arrow::struct_({::arrow::field("embedding", vector_type, /*nullable=*/false)}); + auto type = ::arrow::list(::arrow::field("item", struct_type, /*nullable=*/false)); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [{"embedding": [1.0, 2.0, 3.0]}, {"embedding": [4.0, 5.0, 6.0]}], + [], + [{"embedding": [7.0, 8.0, 9.0]}], + [{"embedding": [10.0, 11.0, 12.0]}, {"embedding": [13.0, 14.0, 15.0]}]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, + /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorMapOfVectorsRoundTrip) { + auto vector_type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::int32(), /*nullable=*/false), 2); + auto type = ::arrow::map(::arrow::utf8(), ::arrow::field("value", vector_type, + /*nullable=*/false)); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [["a", [1, 2]], ["b", [3, 4]]], + [], + [["c", [5, 6]]]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, + /*row_group_size=*/3)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorStructWithNestedVectorFallsBackToList) { + // A fixed-size list nested inside a struct vector element would give the + // struct's leaves different per-row multiplicities, which the write path + // does not support; the column falls back to the LIST encoding. + auto inner = ::arrow::fixed_size_list( + ::arrow::field("value", ::arrow::int32(), /*nullable=*/false), 2); + auto struct_type = + ::arrow::struct_({::arrow::field("v", inner, /*nullable=*/false), + ::arrow::field("x", ::arrow::int32(), /*nullable=*/false)}); + auto type = ::arrow::fixed_size_list( + ::arrow::field("element", struct_type, /*nullable=*/false), 3); + // No null rows: writing null FixedSizeList rows with nested children via + // the LIST encoding is an upstream NotImplemented independent of VECTOR. + auto array = ::arrow::ArrayFromJSON(type, R"([ + [{"v": [1, 2], "x": 1}, {"v": [3, 4], "x": 2}, {"v": [5, 6], "x": 3}], + [{"v": [7, 8], "x": 4}, {"v": [9, 10], "x": 5}, {"v": [11, 12], "x": 6}]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFallsBackToListRoundtrip(table, + /*row_group_size=*/3)); +} + +TEST(ArrowReadWrite, FixedSizeListZeroLength) { + // Regression test: zero-length fixed-size lists used to be written with the + // definition level of a present element instead of an empty list, producing + // one garbage element per row on read. + auto type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::int32(), /*nullable=*/false), 0); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [], + null, + []])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + auto props_store_schema = ArrowWriterProperties::Builder().store_schema()->build(); + CheckSimpleRoundtrip(table, 3, props_store_schema); +} + +TEST(ArrowReadWrite, FixedSizeListVectorZeroLengthFallsBackToList) { + // A zero-length vector value would contribute no physical leaf slots, so + // neither row counts nor row nullability are representable in the VECTOR + // encoding; zero-length fixed-size lists use the LIST encoding instead. + auto type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::int32(), /*nullable=*/false), 0); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [], + null, + []])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFallsBackToListRoundtrip(table, + /*row_group_size=*/3)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorBoolRoundTrip) { + auto type = ::arrow::fixed_size_list( + ::arrow::field("element", ::arrow::boolean(), /*nullable=*/false), 3); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [true, false, true], + null, + [false, false, true]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/3)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorRecordBatchReaderSmallBatches) { + auto table = MakeVectorFixedSizeListTable(R"([ + [1, 2, 3], + null, + [7, 8, 9], + [10, 11, 12]])"); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/4, + VectorWriterProperties(), builder.build())); + + ArrowReaderProperties reader_properties; + reader_properties.set_batch_size(1); + std::unique_ptr reader; + FileReaderBuilder reader_builder; + ASSERT_OK_NO_THROW(reader_builder.Open(std::make_shared(buffer))); + ASSERT_OK(reader_builder.properties(reader_properties)->Build(&reader)); + std::unique_ptr<::arrow::RecordBatchReader> batch_reader; + ASSERT_OK_AND_ASSIGN(batch_reader, reader->GetRecordBatchReader()); + ASSERT_OK_AND_ASSIGN(auto result, Table::FromRecordBatchReader(batch_reader.get())); + ::arrow::AssertTablesEqual(*table, *result, false); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNestedVectorsRoundTrip) { + auto inner_type = ::arrow::fixed_size_list(::arrow::uint8(), 2); + auto type = + ::arrow::fixed_size_list(::arrow::field("item", inner_type, /*nullable=*/false), 2); + auto array = ::arrow::ArrayFromJSON(type, R"([ + [[1, 2], [3, 4]], + null, + [[5, 6], [7, 8]]])"); + auto table = + ::arrow::Table::Make(::arrow::schema({::arrow::field("root", type)}), {array}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableChunkedReadAcrossRowGroups) { + auto table = MakeVectorFixedSizeListTable(R"([ + [1, 2, 3], + null, + [7, 8, 9], + null, + [13, 14, 15]])"); + + ArrowWriterProperties::Builder writer_builder; + writer_builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN( + auto buffer, WriteTableToBuffer(table, /*row_group_size=*/2, + VectorWriterProperties(), writer_builder.build())); + + ArrowReaderProperties reader_properties; + reader_properties.set_batch_size(2); + FileReaderBuilder reader_builder; + ASSERT_OK(reader_builder.Open(std::make_shared(buffer))); + reader_builder.properties(reader_properties); + std::unique_ptr reader; + ASSERT_OK(reader_builder.Build(&reader)); + ASSERT_OK_AND_ASSIGN(auto rb_reader, reader->GetRecordBatchReader()); + ASSERT_OK_AND_ASSIGN(auto out, Table::FromRecordBatchReader(rb_reader.get())); + + ASSERT_EQ(out->column(0)->num_chunks(), 3); + ASSERT_EQ(out->column(0)->chunk(0)->length(), 2); + ASSERT_EQ(out->column(0)->chunk(1)->length(), 2); + ASSERT_EQ(out->column(0)->chunk(2)->length(), 1); + ::arrow::AssertTablesEqual(*table, *out, false); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableRoundTripSlicedInput) { + auto base = ::arrow::ArrayFromJSON(VectorFixedSizeListType(), R"([ + [100, 101, 102], + [1, 2, 3], + null, + [7, 8, 9], + null, + [200, 201, 202]])"); + auto table = MakeVectorFixedSizeListTable({base->Slice(/*offset=*/1, /*length=*/4)}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableRowsAndElementsRoundTrip) { + auto table = MakeVectorFixedSizeListTable(R"([ + [1, null, 3], + null, + [null, null, null], + [7, 8, null]])", + /*nullable=*/true, + /*element_nullable=*/true); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/4)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNestedRoundTrip) { + auto inner_type = ::arrow::fixed_size_list( + ::arrow::field("item", ::arrow::int16(), /*nullable=*/true), /*list_size=*/2); + auto outer_type = ::arrow::fixed_size_list( + ::arrow::field("item", inner_type, /*nullable=*/true), /*list_size=*/3); + auto table = ::arrow::Table::Make( + ::arrow::schema({::arrow::field("root", outer_type, /*nullable=*/true)}), + {::arrow::ArrayFromJSON(outer_type, R"([ + [[1, 2], null, [5, null]], + null, + [[7, 8], [9, 10], [null, 12]]])")}); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorMixedColumnsRoundTripAcrossRowGroups) { + auto vector_type = + ::arrow::fixed_size_list(::arrow::field("item", ::arrow::int16(), true), + /*list_size=*/3); + auto list_type = ::arrow::list(::arrow::field("item", ::arrow::int32(), true)); + + auto vector_chunks = std::vector>{ + ::arrow::ArrayFromJSON(vector_type, + R"([[1, null, 3], null, [7, 8, null], [10, 11, 12]])"), + ::arrow::ArrayFromJSON(vector_type, + R"([null, [16, null, 18], [19, 20, 21], [22, null, 24]])")}; + auto table = ::arrow::Table::Make( + ::arrow::schema({ + ::arrow::field("id", ::arrow::int32(), false), + ::arrow::field("embedding", vector_type, true), + ::arrow::field("label", ::arrow::utf8(), true), + ::arrow::field("tags", list_type, true), + }), + { + std::make_shared( + ::arrow::ArrayFromJSON(::arrow::int32(), "[0, 1, 2, 3, 4, 5, 6, 7]")), + std::make_shared(std::move(vector_chunks), vector_type), + std::make_shared(::arrow::ArrayFromJSON( + ::arrow::utf8(), R"(["a", null, "c", "d", null, "f", "g", "h"])")), + std::make_shared(::arrow::ArrayFromJSON( + list_type, R"([[1, 2], [], null, [3, null], [4], [5, 6], null, []])")), + }); + + ArrowWriterProperties::Builder writer_builder; + writer_builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN( + auto buffer, WriteTableToBuffer(table, /*row_group_size=*/3, + VectorWriterProperties(), writer_builder.build())); + + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + ASSERT_EQ(parquet_reader->metadata()->num_row_groups(), 3); + for (int i = 0; i < 3; ++i) { + const int64_t expected_rows = i == 2 ? 2 : 3; + EXPECT_EQ(parquet_reader->metadata()->RowGroup(i)->num_rows(), expected_rows); + EXPECT_EQ(parquet_reader->metadata()->RowGroup(i)->ColumnChunk(1)->num_values(), + expected_rows * 3); + } + + ArrowReaderProperties reader_properties; + reader_properties.set_batch_size(2); + FileReaderBuilder reader_builder; + ASSERT_OK(reader_builder.Open(std::make_shared(buffer))); + reader_builder.properties(reader_properties); + std::unique_ptr reader; + ASSERT_OK(reader_builder.Build(&reader)); + ASSERT_OK_AND_ASSIGN(auto rb_reader, reader->GetRecordBatchReader()); + ASSERT_OK_AND_ASSIGN(auto out, Table::FromRecordBatchReader(rb_reader.get())); + + ::arrow::AssertSchemaEqual(*table->schema(), *out->schema(), false); + ::arrow::AssertTablesEqual(*table, *out, false); +} + +std::shared_ptr<::arrow::DataType> VectorFixedSizeListStructType( + bool element_nullable = false, bool field_nullable = false) { + return ::arrow::fixed_size_list( + ::arrow::field( + "item", + ::arrow::struct_({::arrow::field("x", ::arrow::float32(), false), + ::arrow::field("y", ::arrow::int32(), field_nullable)}), + element_nullable), + /*size=*/2); +} + +std::shared_ptr
MakeVectorFixedSizeListStructTable(std::string_view json, + bool nullable = true, + bool element_nullable = false, + bool field_nullable = false) { + auto type = VectorFixedSizeListStructType(element_nullable, field_nullable); + auto field = ::arrow::field("root", type, nullable); + auto array = ::arrow::ArrayFromJSON(type, std::string(json)); + return ::arrow::Table::Make(::arrow::schema({field}), {array}); +} + +std::shared_ptr<::arrow::DataType> VectorFixedSizeListNestedStructType( + bool element_nullable = false) { + return ::arrow::fixed_size_list( + ::arrow::field( + "item", + ::arrow::struct_( + {::arrow::field( + "point", + ::arrow::struct_({::arrow::field("x", ::arrow::float32(), false), + ::arrow::field("y", ::arrow::int32(), true)}), + false), + ::arrow::field("z", ::arrow::int16(), true)}), + element_nullable), + /*size=*/2); +} + +std::shared_ptr
MakeVectorFixedSizeListNestedStructTable( + std::string_view json, bool nullable = true, bool element_nullable = false) { + auto type = VectorFixedSizeListNestedStructType(element_nullable); + auto field = ::arrow::field("root", type, nullable); + auto array = ::arrow::ArrayFromJSON(type, std::string(json)); + return ::arrow::Table::Make(::arrow::schema({field}), {array}); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableStructElementsRoundTrip) { + auto table = MakeVectorFixedSizeListStructTable(R"([ + [{"x": 1.0, "y": 1}, null], + null, + [null, {"x": 6.0, "y": null}], + [{"x": 7.0, "y": 7}, {"x": 8.0, "y": 8}]])", + /*nullable=*/true, + /*element_nullable=*/true, + /*field_nullable=*/true); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/2)); +} + +TEST(ArrowReadWrite, FixedSizeListVectorNullableNestedStructElementsRoundTrip) { + auto table = MakeVectorFixedSizeListNestedStructTable(R"([ + [{"point": {"x": 1.0, "y": 1}, "z": 10}, null], + null, + [null, {"point": {"x": 4.0, "y": 4}, "z": 40}], + [{"point": {"x": 5.0, "y": null}, "z": null}, + {"point": {"x": 6.0, "y": 6}, "z": 60}]])", + /*nullable=*/true, + /*element_nullable=*/true); + + ASSERT_NO_FATAL_FAILURE(CheckVectorFixedSizeListRoundtrip(table, /*row_group_size=*/2)); +} + +TEST(ArrowWriteOnly, FixedSizeListVectorDefaultWriterProperties) { + auto table = MakeVectorFixedSizeListTable(R"([ + [1, 2, 3], + [4, 5, 6]])", + /*nullable=*/false); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/2, + default_writer_properties(), builder.build())); + + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + const auto encodings = + parquet_reader->metadata()->RowGroup(0)->ColumnChunk(0)->encodings(); + EXPECT_NE(std::find(encodings.begin(), encodings.end(), Encoding::RLE_DICTIONARY), + encodings.end()); +} + +TEST(ArrowWriteOnly, FixedSizeListVectorNullableElementStatistics) { + auto table = MakeVectorFixedSizeListTable(::arrow::int32(), R"([ + [1, 2, 3], + null, + [4, null, 6]])", + /*nullable=*/true, + /*element_nullable=*/true); + + auto writer_properties = WriterProperties::Builder() + .disable_dictionary() + ->encoding(Encoding::PLAIN) + ->build(); + ArrowWriterProperties::Builder arrow_builder; + arrow_builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/3, writer_properties, + arrow_builder.build())); + + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + auto column = parquet_reader->metadata()->RowGroup(0)->ColumnChunk(0); + ASSERT_TRUE(column->is_stats_set()); + auto statistics = column->statistics(); + ASSERT_EQ(5, statistics->num_values()); + ASSERT_EQ(4, statistics->null_count()); + + std::shared_ptr min, max; + ASSERT_OK(StatisticsAsScalars(*statistics, &min, &max)); + ASSERT_OK_AND_ASSIGN(auto expected_min, ::arrow::MakeScalar(::arrow::int32(), 1)); + ASSERT_OK_AND_ASSIGN(auto expected_max, ::arrow::MakeScalar(::arrow::int32(), 6)); + ::arrow::AssertScalarsEqual(*expected_min, *min, /*verbose=*/true); + ::arrow::AssertScalarsEqual(*expected_max, *max, /*verbose=*/true); +} + +TEST(ArrowWriteOnly, FixedSizeListVectorPageIndex) { + auto table = MakeVectorFixedSizeListTable(::arrow::int32(), R"([ + [1, 2, 3], + [4, 5, 6], + [7, 8, 9]])", + /*nullable=*/false); + + auto writer_properties = WriterProperties::Builder() + .disable_dictionary() + ->encoding(Encoding::PLAIN) + ->enable_write_page_index() + ->build(); + ArrowWriterProperties::Builder arrow_builder; + arrow_builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/2, writer_properties, + arrow_builder.build())); + + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + auto page_index_reader = parquet_reader->GetPageIndexReader(); + ASSERT_NE(nullptr, page_index_reader); + auto row_group_index_reader = page_index_reader->RowGroup(0); + ASSERT_NE(nullptr, row_group_index_reader); + ASSERT_NE(nullptr, row_group_index_reader->GetColumnIndex(0)); + ASSERT_NE(nullptr, row_group_index_reader->GetOffsetIndex(0)); +} + +TEST(ArrowWriteOnly, FixedSizeListVectorBloomFilter) { + auto table = MakeVectorFixedSizeListTable(::arrow::int32(), R"([ + [1, 2, 3], + [4, 5, 6]])", + /*nullable=*/false); + + BloomFilterOptions bloom_filter_options{10, 0.05}; + auto writer_properties = + WriterProperties::Builder() + .disable_dictionary() + ->encoding(Encoding::PLAIN) + ->enable_bloom_filter("root.list.element", bloom_filter_options) + ->build(); + ArrowWriterProperties::Builder arrow_builder; + arrow_builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN(auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/2, writer_properties, + arrow_builder.build())); + + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + auto& bloom_filter_reader = parquet_reader->GetBloomFilterReader(); + auto row_group_bloom_filter = bloom_filter_reader.RowGroup(0); + ASSERT_NE(nullptr, row_group_bloom_filter); + auto bloom_filter = row_group_bloom_filter->GetColumnBloomFilter(0); + ASSERT_NE(nullptr, bloom_filter); + + int32_t present = 4; + EXPECT_TRUE(bloom_filter->FindHash(bloom_filter->Hash(present))); +} + +TEST(ArrowReadWrite, FixedSizeListVectorContentDefinedChunkingRoundTrip) { + std::stringstream json; + json << "["; + for (int i = 0; i < 40; ++i) { + if (i != 0) { + json << ","; + } + json << "[" << (3 * i + 1) << "," << (3 * i + 2) << "," << (3 * i + 3) << "]"; + } + json << "]"; + auto table = MakeVectorFixedSizeListTable(::arrow::int32(), json.str(), + /*nullable=*/false); + + CdcOptions cdc_options; + cdc_options.min_chunk_size = 64; + cdc_options.max_chunk_size = 128; + auto writer_properties = WriterProperties::Builder() + .enable_content_defined_chunking() + ->content_defined_chunking_options(cdc_options) + ->build(); + ArrowWriterProperties::Builder arrow_builder; + arrow_builder.enable_experimental_vector_encoding(); + + std::shared_ptr
result; + ASSERT_NO_FATAL_FAILURE(DoRoundtrip(table, /*row_group_size=*/40, &result, + writer_properties, arrow_builder.build())); + ::arrow::AssertSchemaEqual(*table->schema(), *result->schema(), false); + ::arrow::AssertTablesEqual(*table, *result, false); +} + +TEST(ArrowReadWrite, FixedSizeListVectorContentDefinedChunkingNullableRoundTrip) { + auto table = MakeVectorFixedSizeListTable(::arrow::int32(), R"([ + [1, 2, 3], + null, + [4, null, 6], + [7, 8, 9]])", + /*nullable=*/true, + /*element_nullable=*/true); + + CdcOptions cdc_options; + cdc_options.min_chunk_size = 64; + cdc_options.max_chunk_size = 128; + auto writer_properties = WriterProperties::Builder() + .disable_dictionary() + ->enable_content_defined_chunking() + ->content_defined_chunking_options(cdc_options) + ->build(); + ArrowWriterProperties::Builder arrow_builder; + arrow_builder.enable_experimental_vector_encoding(); + + std::shared_ptr
result; + ASSERT_NO_FATAL_FAILURE(DoRoundtrip(table, /*row_group_size=*/4, &result, + writer_properties, arrow_builder.build())); + ::arrow::AssertSchemaEqual(*table->schema(), *result->schema(), false); + ::arrow::AssertTablesEqual(*table, *result, false); +} + +TEST(ArrowReadWrite, FixedSizeListVectorByteStreamSplitRoundTrip) { + auto table = MakeVectorFixedSizeListTable(::arrow::float32(), R"([ + [1.0, 2.0, 3.0], + [4.0, 5.0, 6.0], + [7.0, 8.0, 9.0]])", + /*nullable=*/false); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK_AND_ASSIGN( + auto buffer, + WriteTableToBuffer(table, /*row_group_size=*/2, + VectorByteStreamSplitWriterProperties(), builder.build())); + + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + const auto encodings = + parquet_reader->metadata()->RowGroup(0)->ColumnChunk(0)->encodings(); + ASSERT_NE(std::find(encodings.begin(), encodings.end(), Encoding::BYTE_STREAM_SPLIT), + encodings.end()); + + FileReaderBuilder reader_builder; + ASSERT_OK(reader_builder.Open(std::make_shared(buffer))); + std::unique_ptr reader; + ASSERT_OK(reader_builder.Build(&reader)); + std::shared_ptr
result; + ASSERT_OK(reader->ReadTable(&result)); + ::arrow::AssertSchemaEqual(*table->schema(), *result->schema(), false); + ::arrow::AssertTablesEqual(*table, *result, false); +} + +TEST(ArrowReadWrite, FixedSizeListVectorByteStreamSplitNullableRoundTrip) { + auto table = MakeVectorFixedSizeListTable(::arrow::float32(), R"([ + [1.0, 2.0, 3.0], + null, + [7.0, 8.0, 9.0]])"); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + std::shared_ptr
result; + ASSERT_NO_FATAL_FAILURE(DoRoundtrip( + table, 3, &result, VectorByteStreamSplitWriterProperties(), builder.build())); + ::arrow::AssertSchemaEqual(*table->schema(), *result->schema(), false); + ::arrow::AssertTablesEqual(*table, *result, false); +} + TEST(ArrowReadWrite, ListOfStructOfList2) { using ::arrow::field; using ::arrow::list; @@ -4360,8 +5225,7 @@ TEST(TestArrowReaderAdHoc, OldDataPageV2) { GTEST_SKIP() << "ARROW_TEST_DATA not set."; } std::stringstream ss; - ss << c_root << "/" - << "parquet/ARROW-17100.parquet"; + ss << c_root << "/" << "parquet/ARROW-17100.parquet"; std::string path = ss.str(); TryReadDataFile(path); } diff --git a/cpp/src/parquet/arrow/arrow_schema_test.cc b/cpp/src/parquet/arrow/arrow_schema_test.cc index 7d9ecb5e6449..8e9a92a928ae 100644 --- a/cpp/src/parquet/arrow/arrow_schema_test.cc +++ b/cpp/src/parquet/arrow/arrow_schema_test.cc @@ -289,6 +289,138 @@ TEST_F(TestConvertParquetSchema, ParquetAnnotatedFields) { ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(arrow_schema)); } +TEST_F(TestConvertParquetSchema, UnannotatedVectorGroupRejected) { + // The VECTOR logical type annotation is required: a plain group containing a + // VECTOR-repeated child is not silently interpreted as a vector. + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::FLOAT); + auto vector = GroupNode::Make("vector", Repetition::VECTOR, {element}, + /*logical_type=*/nullptr, -1, 3); + std::vector parquet_fields = {GroupNode::Make( + "embedding", Repetition::OPTIONAL, {vector}, /*logical_type=*/nullptr)}; + + ASSERT_RAISES(Invalid, ConvertSchema(parquet_fields)); +} + +TEST_F(TestConvertParquetSchema, VectorFixedSizeListAnnotated) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::FLOAT); + auto vector = GroupNode::Make("vector", Repetition::VECTOR, {element}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + arrow_fields.push_back(::arrow::field( + "embedding", + ::arrow::fixed_size_list(::arrow::field("element", ::arrow::float32(), false), 3), + true)); + + ASSERT_OK(ConvertSchema(parquet_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(::arrow::schema(arrow_fields))); +} + +TEST_F(TestConvertParquetSchema, VectorAnnotatedGroupWithoutVectorChildRejected) { + // A group annotated with the VECTOR logical type must contain a single + // VECTOR-repeated child. + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::FLOAT); + std::vector parquet_fields = {GroupNode::Make( + "embedding", Repetition::OPTIONAL, {element}, LogicalType::Vector())}; + + ASSERT_RAISES(Invalid, ConvertSchema(parquet_fields)); +} + +TEST_F(TestConvertParquetSchema, VectorOnPrimitiveRejected) { + // VECTOR repetition directly on a primitive is not part of the canonical + // 3-level vector structure; vectors must use a VECTOR group whose single + // child is a VECTOR-repeated group containing the element. + auto element = PrimitiveNode::Make("element", Repetition::VECTOR, ParquetType::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, 3); + std::vector parquet_fields = { + GroupNode::Make("embedding", Repetition::OPTIONAL, {element}, + LogicalType::Vector()), + }; + ASSERT_RAISES(Invalid, ConvertSchema(parquet_fields)); + + // Also rejected without the annotation (structural detection). + std::vector unannotated_fields = { + GroupNode::Make("embedding", Repetition::OPTIONAL, {element}, + /*logical_type=*/nullptr), + }; + ASSERT_RAISES(Invalid, ConvertSchema(unannotated_fields)); +} + +TEST_F(TestConvertParquetSchema, VectorFixedSizeListNullableStructElement) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto x = PrimitiveNode::Make("x", Repetition::REQUIRED, ParquetType::FLOAT); + auto y = PrimitiveNode::Make("y", Repetition::OPTIONAL, ParquetType::INT32); + auto item = GroupNode::Make("element", Repetition::OPTIONAL, {x, y}); + auto vector = GroupNode::Make("vector", Repetition::VECTOR, {item}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + arrow_fields.push_back(::arrow::field( + "embedding", + ::arrow::fixed_size_list( + ::arrow::field("element", + ::arrow::struct_({::arrow::field("x", FLOAT, false), + ::arrow::field("y", INT32, true)}), + true), + 3), + true)); + + ASSERT_OK(ConvertSchema(parquet_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(::arrow::schema(arrow_fields))); +} + +TEST_F(TestConvertParquetSchema, VectorFixedSizeListNullableNestedStructElement) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto x = PrimitiveNode::Make("x", Repetition::REQUIRED, ParquetType::FLOAT); + auto y = PrimitiveNode::Make("y", Repetition::OPTIONAL, ParquetType::INT32); + auto point = GroupNode::Make("point", Repetition::REQUIRED, {x, y}); + auto z = PrimitiveNode::Make("z", Repetition::OPTIONAL, LogicalType::Int(16, true), + ParquetType::INT32); + auto item = GroupNode::Make("element", Repetition::OPTIONAL, {point, z}); + auto vector = GroupNode::Make("vector", Repetition::VECTOR, {item}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + arrow_fields.push_back(::arrow::field( + "embedding", + ::arrow::fixed_size_list( + ::arrow::field( + "element", + ::arrow::struct_( + {::arrow::field("point", + ::arrow::struct_({::arrow::field("x", FLOAT, false), + ::arrow::field("y", INT32, true)}), + false), + ::arrow::field("z", ::arrow::int16(), true)}), + true), + 3), + true)); + + ASSERT_OK(ConvertSchema(parquet_fields)); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(::arrow::schema(arrow_fields))); +} + +TEST_F(TestConvertParquetSchema, VectorFixedSizeListStructWithListRejected) { + auto list_element = + PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32); + auto list = GroupNode::Make("list", Repetition::REPEATED, {list_element}); + auto items = + GroupNode::Make("items", Repetition::REQUIRED, {list}, LogicalType::List()); + auto vector = GroupNode::Make("vector", Repetition::VECTOR, {items}, + /*logical_type=*/nullptr, -1, 3); + std::vector parquet_fields = {GroupNode::Make( + "embedding", Repetition::OPTIONAL, {vector}, LogicalType::Vector())}; + + ASSERT_RAISES(NotImplemented, ConvertSchema(parquet_fields)); +} + TEST_F(TestConvertParquetSchema, DuplicateFieldNames) { std::vector parquet_fields; std::vector> arrow_fields; @@ -1746,12 +1878,8 @@ TEST_F(TestConvertArrowSchema, ParquetOtherLists) { auto arrow_list = ::arrow::large_list(arrow_element); arrow_fields.push_back(::arrow::field("my_list", arrow_list, false)); } - // // FixedSizeList[10] (list-like non-null, elements nullable) - // required group my_list (LIST) { - // repeated group list { - // optional binary element (UTF8); - // } - // } + // FixedSizeList defaults to the legacy LIST encoding unless experimental VECTOR + // encoding is explicitly enabled. { auto element = PrimitiveNode::Make("element", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY, ConvertedType::UTF8); @@ -1768,6 +1896,230 @@ TEST_F(TestConvertArrowSchema, ParquetOtherLists) { ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); } +TEST_F(TestConvertArrowSchema, + ParquetFixedSizeListVectorUnsupportedElementFallsBackToList) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto element = PrimitiveNode::Make("element", Repetition::OPTIONAL, + ParquetType::BYTE_ARRAY, ConvertedType::UTF8); + auto list = GroupNode::Make("list", Repetition::REPEATED, {element}); + parquet_fields.push_back( + GroupNode::Make("embedding", Repetition::OPTIONAL, {list}, ConvertedType::LIST)); + + auto arrow_element = ::arrow::field("string", UTF8, true); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 3); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorZeroLengthFallsBackToList) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::FLOAT); + auto list = GroupNode::Make("list", Repetition::REPEATED, {element}); + parquet_fields.push_back( + GroupNode::Make("embedding", Repetition::OPTIONAL, {list}, ConvertedType::LIST)); + + auto arrow_element = ::arrow::field("element", FLOAT, false); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 0); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorUnderList) { + // FixedSizeList values below a repeated field can use VECTOR encoding. + std::vector parquet_fields; + std::vector> arrow_fields; + + auto inner_element = + PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32); + auto vector = GroupNode::Make("list", Repetition::VECTOR, {inner_element}, + /*logical_type=*/nullptr, -1, 2); + auto element = + GroupNode::Make("element", Repetition::REQUIRED, {vector}, LogicalType::Vector()); + auto list = GroupNode::Make("list", Repetition::REPEATED, {element}); + parquet_fields.push_back( + GroupNode::Make("embeddings", Repetition::OPTIONAL, {list}, ConvertedType::LIST)); + + auto vector_type = + ::arrow::fixed_size_list(::arrow::field("element", INT32, /*nullable=*/false), 2); + arrow_fields.push_back(::arrow::field( + "embeddings", ::arrow::list(::arrow::field("element", vector_type, false)), true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorMixedSupportedAndFallback) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::FLOAT); + auto vector = GroupNode::Make("list", Repetition::VECTOR, {element}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + + auto string_element = PrimitiveNode::Make("element", Repetition::OPTIONAL, + ParquetType::BYTE_ARRAY, ConvertedType::UTF8); + auto string_list = GroupNode::Make("list", Repetition::REPEATED, {string_element}); + parquet_fields.push_back(GroupNode::Make("labels", Repetition::OPTIONAL, {string_list}, + ConvertedType::LIST)); + + arrow_fields.push_back(::arrow::field( + "embedding", ::arrow::fixed_size_list(::arrow::field("element", FLOAT, false), 3), + true)); + arrow_fields.push_back(::arrow::field( + "labels", ::arrow::fixed_size_list(::arrow::field("string", UTF8, true), 3), true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorNullableElement) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto item = PrimitiveNode::Make("element", Repetition::OPTIONAL, ParquetType::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, -1); + auto vector = GroupNode::Make("list", Repetition::VECTOR, {item}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + + auto arrow_element = ::arrow::field("element", FLOAT, true); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 3); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorNullableStructElement) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto x = PrimitiveNode::Make("x", Repetition::REQUIRED, ParquetType::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, -1); + auto y = PrimitiveNode::Make("y", Repetition::OPTIONAL, ParquetType::INT32, + ConvertedType::NONE, -1, -1, -1, -1, -1); + auto item = GroupNode::Make("element", Repetition::OPTIONAL, {x, y}); + auto vector = GroupNode::Make("list", Repetition::VECTOR, {item}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + + auto arrow_element = ::arrow::field( + "element", + ::arrow::struct_({::arrow::field("x", FLOAT, false), ::arrow::field("y", INT32)}), + true); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 3); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorNullableNestedStructElement) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto x = PrimitiveNode::Make("x", Repetition::REQUIRED, ParquetType::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, -1); + auto y = PrimitiveNode::Make("y", Repetition::OPTIONAL, ParquetType::INT32, + ConvertedType::NONE, -1, -1, -1, -1, -1); + auto point = GroupNode::Make("point", Repetition::REQUIRED, {x, y}); + auto z = PrimitiveNode::Make("z", Repetition::OPTIONAL, LogicalType::Int(16, true), + ParquetType::INT32, -1, -1, -1); + auto item = GroupNode::Make("element", Repetition::OPTIONAL, {point, z}); + auto vector = GroupNode::Make("list", Repetition::VECTOR, {item}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + + auto arrow_element = ::arrow::field( + "element", + ::arrow::struct_( + {::arrow::field("point", + ::arrow::struct_({::arrow::field("x", FLOAT, false), + ::arrow::field("y", INT32, true)}), + false), + ::arrow::field("z", ::arrow::int16(), true)}), + true); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 3); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVectorStructWithListFallsBackToList) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto item = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::INT32); + auto items_list = GroupNode::Make("list", Repetition::REPEATED, {item}); + auto items = + GroupNode::Make("items", Repetition::REQUIRED, {items_list}, ConvertedType::LIST); + auto element = GroupNode::Make("element", Repetition::REQUIRED, {items}); + auto list = GroupNode::Make("list", Repetition::REPEATED, {element}); + parquet_fields.push_back( + GroupNode::Make("embedding", Repetition::OPTIONAL, {list}, ConvertedType::LIST)); + + auto arrow_element = ::arrow::field( + "element", + ::arrow::struct_({::arrow::field( + "items", ::arrow::list(::arrow::field("element", INT32, false)), false)}), + false); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 3); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + +TEST_F(TestConvertArrowSchema, ParquetFixedSizeListVector) { + std::vector parquet_fields; + std::vector> arrow_fields; + + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, ParquetType::FLOAT); + auto vector = GroupNode::Make("list", Repetition::VECTOR, {element}, + /*logical_type=*/nullptr, -1, 3); + parquet_fields.push_back(GroupNode::Make("embedding", Repetition::OPTIONAL, {vector}, + LogicalType::Vector())); + + auto arrow_element = ::arrow::field("element", FLOAT, false); + auto arrow_list = ::arrow::fixed_size_list(arrow_element, 3); + arrow_fields.push_back(::arrow::field("embedding", arrow_list, true)); + + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + ASSERT_OK(ConvertSchema(arrow_fields, builder.build())); + ASSERT_NO_FATAL_FAILURE(CheckFlatSchema(parquet_fields)); +} + TEST_F(TestConvertArrowSchema, ParquetNestedComplianceEnabledNullable) { std::vector parquet_fields; std::vector> arrow_fields; diff --git a/cpp/src/parquet/arrow/path_internal.cc b/cpp/src/parquet/arrow/path_internal.cc index 002859a5e7d5..6648d4c40b65 100644 --- a/cpp/src/parquet/arrow/path_internal.cc +++ b/cpp/src/parquet/arrow/path_internal.cc @@ -106,6 +106,7 @@ #include "arrow/util/macros.h" #include "arrow/visit_array_inline.h" +#include "parquet/arrow/schema_internal.h" #include "parquet/properties.h" namespace parquet::arrow { @@ -209,6 +210,29 @@ struct PathWriteContext { visited_elements.push_back(range); } + void RecordVectorPostListVisit(const ElementRange& logical_range, + const ElementRange& physical_range) { + if (!visited_elements.empty() && + visited_elements.back().start == logical_range.start && + visited_elements.back().end == logical_range.end) { + visited_elements.back() = physical_range; + return; + } + RecordPostListVisit(physical_range); + } + + // The enclosing list machinery records the logical (vector-index) child + // range before delegating to a VectorNullableNode, which then records the + // physical (slot-index) ranges of its present runs itself. Remove the + // logical range so only physical ranges remain. + void TrimVectorLogicalVisit(const ElementRange& logical_range) { + if (!visited_elements.empty() && + visited_elements.back().start == logical_range.start && + visited_elements.back().end == logical_range.end) { + visited_elements.pop_back(); + } + } + Status last_status; TypedBufferBuilder rep_levels; TypedBufferBuilder def_levels; @@ -450,6 +474,120 @@ struct FixedSizedRangeSelector { int list_size; }; +struct NoLevelTerminalNode { + IterationResult Run(const ElementRange&, PathWriteContext*) { return kDone; } +}; + +IterationResult ExpandLastRepLevels(int64_t logical_count, int32_t multiplier, + int16_t filler_rep_level, PathWriteContext* context); + +class VectorNullableNode { + public: + VectorNullableNode(const uint8_t* null_bitmap, int64_t entry_offset, + int32_t vector_length, int32_t child_slot_multiplier, + int16_t def_level_if_present, int16_t filler_rep_level, + bool child_emits_present_def_levels, + bool child_records_visited_elements, bool expand_rep_levels) + : null_bitmap_(null_bitmap), + entry_offset_(entry_offset), + vector_length_(vector_length), + child_slot_multiplier_(child_slot_multiplier), + valid_bits_reader_(MakeReader(ElementRange{0, 0})), + def_level_if_present_(def_level_if_present), + def_level_if_null_(def_level_if_present - 1), + filler_rep_level_(filler_rep_level), + child_emits_present_def_levels_(child_emits_present_def_levels), + child_records_visited_elements_(child_records_visited_elements), + expand_rep_levels_(expand_rep_levels), + new_range_(true) {} + + ::arrow::internal::BitRunReader MakeReader(const ElementRange& range) { + return ::arrow::internal::BitRunReader(null_bitmap_, entry_offset_ + range.start, + range.Size()); + } + + IterationResult Run(ElementRange* range, ElementRange* child_range, + PathWriteContext* context) { + if (range->Empty()) { + new_range_ = true; + return kDone; + } + if (new_range_) { + context->TrimVectorLogicalVisit(*range); + } + if (expand_rep_levels_ && new_range_) { + // Expand the trailing repetition levels (one per logical vector in this + // delegation, appended by the enclosing list machinery) once, before + // any run processing. Every vector contributes + // vector_length * child_slot_multiplier leaf slots regardless of + // validity, so the multiplier is constant and per-run expansion (which + // would expand positionally wrong entries for interleaved null and + // present runs) is not needed. Only the outermost vector node of a + // nested chain expands. + RETURN_IF_ERROR(ExpandLastRepLevels(range->Size(), + vector_length_ * child_slot_multiplier_, + filler_rep_level_, context)); + } + if (null_bitmap_ == nullptr) { + ElementRange logical_range = *range; + child_range->start = range->start * vector_length_; + child_range->end = child_range->start + range->Size() * vector_length_; + if (!child_records_visited_elements_) { + context->RecordVectorPostListVisit(logical_range, *child_range); + } + if (!child_emits_present_def_levels_) { + RETURN_IF_ERROR( + context->AppendDefLevels(child_range->Size(), def_level_if_present_)); + } + range->start = range->end; + new_range_ = false; + return kNext; + } + if (new_range_) { + valid_bits_reader_ = MakeReader(*range); + } + ::arrow::internal::BitRun run = valid_bits_reader_.NextRun(); + new_range_ = false; + while (!range->Empty() && !run.set) { + range->start += run.length; + RETURN_IF_ERROR(context->AppendDefLevels( + run.length * vector_length_ * child_slot_multiplier_, def_level_if_null_)); + run = valid_bits_reader_.NextRun(); + } + if (range->Empty()) { + new_range_ = true; + return kDone; + } + + ElementRange logical_range{range->start, range->start + run.length}; + child_range->start = range->start * vector_length_; + child_range->end = child_range->start + run.length * vector_length_; + if (!child_records_visited_elements_) { + context->RecordVectorPostListVisit(logical_range, *child_range); + } + if (!child_emits_present_def_levels_) { + RETURN_IF_ERROR( + context->AppendDefLevels(run.length * vector_length_, def_level_if_present_)); + } + range->start += run.length; + return kNext; + } + + private: + const uint8_t* null_bitmap_; + int64_t entry_offset_; + int32_t vector_length_; + int32_t child_slot_multiplier_; + ::arrow::internal::BitRunReader valid_bits_reader_; + int16_t def_level_if_present_; + int16_t def_level_if_null_; + int16_t filler_rep_level_; + bool child_emits_present_def_levels_; + bool child_records_visited_elements_; + bool expand_rep_levels_; + bool new_range_ = true; +}; + // An intermediate node that handles null values. class NullableNode { public: @@ -512,6 +650,17 @@ using ListNode = ListPathNode>; using LargeListNode = ListPathNode>; using FixedSizeListNode = ListPathNode; +// Number of physical leaf slots contributed per element of the given type +// (the product of nested fixed-size-list sizes; 1 for non-nested elements). +int32_t VectorSlotMultiplier(const ::arrow::DataType& type) { + if (type.id() != ::arrow::Type::FIXED_SIZE_LIST) { + return 1; + } + const auto& list_type = + ::arrow::internal::checked_cast(type); + return list_type.list_size() * VectorSlotMultiplier(*list_type.value_type()); +} + // Contains static information derived from traversing the schema. struct PathInfo { // The vectors are expected to the same length info. @@ -519,7 +668,8 @@ struct PathInfo { // Note index order matters here. using Node = std::variant; + NullableNode, VectorNullableNode, AllPresentTerminalNode, + AllNullsTerminalNode, NoLevelTerminalNode>; std::vector path; std::shared_ptr primitive_array; @@ -527,6 +677,18 @@ struct PathInfo { int16_t max_rep_level = 0; bool has_dictionary = false; bool leaf_is_nullable = false; + bool leaf_is_vector = false; + int32_t leaf_vector_length = 1; + // Definition level after the last repeated ancestor (0 when there is none). + // For VECTOR leaves, definition levels below it mark ancestors that made + // the vector value absent (empty or null lists); they occupy one level + // entry with no slots, while levels at or above it occupy fixed strides of + // leaf_vector_length slots. + int16_t repeated_ancestor_def_level = 0; + // Definition level at or above which a leaf slot carries a value from the + // leaf array (all vector ancestors present); slots below it (but at or + // above the stride level) belong to null vector values. + int16_t vector_values_def_level = 0; }; /// Contains logic for writing a single leaf node to parquet. @@ -536,6 +698,48 @@ struct PathInfo { /// values have been calculated for root_range with the calculated /// values. It is intended to abstract the complexity of writing /// the levels and values to parquet. +// Counts the physical leaf slots of a VECTOR leaf: definition levels at or +// above the stride definition level. Entries below it are absent-ancestor +// markers carrying no slot. +int64_t CountVectorSlots(const int16_t* def_levels, int64_t def_count, + int16_t repeated_ancestor_def_level) { + if (repeated_ancestor_def_level <= 0) { + return def_count; + } + int64_t slots = 0; + for (int64_t i = 0; i < def_count; ++i) { + slots += def_levels[i] >= repeated_ancestor_def_level; + } + return slots; +} + +IterationResult ExpandLastRepLevels(int64_t logical_count, int32_t multiplier, + int16_t filler_rep_level, PathWriteContext* context) { + if (multiplier <= 1 || logical_count == 0 || context->rep_levels.length() == 0) { + return kDone; + } + const int64_t old_length = context->rep_levels.length(); + if (ARROW_PREDICT_FALSE(old_length < logical_count)) { + context->last_status = + Status::Invalid("VECTOR repetition level expansion needed ", logical_count, + " source levels but only ", old_length, " were available"); + return kError; + } + const int64_t prefix_length = old_length - logical_count; + std::vector expanded; + expanded.reserve(static_cast(prefix_length + logical_count * multiplier)); + expanded.insert(expanded.end(), context->rep_levels.data(), + context->rep_levels.data() + prefix_length); + const int16_t* source = context->rep_levels.data() + prefix_length; + for (int64_t i = 0; i < logical_count; ++i) { + expanded.push_back(source[i]); + expanded.insert(expanded.end(), multiplier - 1, filler_rep_level); + } + context->rep_levels.Reset(); + context->last_status = context->rep_levels.Append(expanded.data(), expanded.size()); + return context->last_status.ok() ? kDone : kError; +} + Status WritePath(ElementRange root_range, PathInfo* path_info, ArrowWriteContext* arrow_context, MultipathLevelBuilder::CallbackFunction writer) { @@ -543,12 +747,18 @@ Status WritePath(ElementRange root_range, PathInfo* path_info, MultipathLevelBuilderResult builder_result; builder_result.leaf_array = path_info->primitive_array; builder_result.leaf_is_nullable = path_info->leaf_is_nullable; + builder_result.leaf_is_vector = path_info->leaf_is_vector; + builder_result.leaf_vector_length = path_info->leaf_vector_length; + builder_result.vector_repeated_ancestor_def_level = + path_info->repeated_ancestor_def_level; + builder_result.vector_values_def_level = path_info->vector_values_def_level; if (path_info->max_def_level == 0) { // This case only occurs when there are no nullable or repeated // columns in the path from the root to leaf. int64_t leaf_length = builder_result.leaf_array->length(); builder_result.def_rep_level_count = leaf_length; + builder_result.leaf_slot_count = leaf_length; builder_result.post_list_visited_elements.push_back({0, leaf_length}); return writer(builder_result); } @@ -597,6 +807,12 @@ Status WritePath(ElementRange root_range, PathInfo* path_info, IterationResult operator()(LargeListNode& node) { return node.Run(stack_position, stack_position + 1, context); } + IterationResult operator()(VectorNullableNode& node) { + return node.Run(stack_position, stack_position + 1, context); + } + IterationResult operator()(NoLevelTerminalNode& node) { + return node.Run(*stack_position, context); + } ElementRange* stack_position; PathWriteContext* context; } visitor = {stack_position, &context}; @@ -616,6 +832,11 @@ Status WritePath(ElementRange root_range, PathInfo* path_info, // This case only occurs when there was a repeated element that needs to be // processed. builder_result.rep_levels = context.rep_levels.data(); + if (path_info->leaf_is_vector) { + builder_result.leaf_slot_count = + CountVectorSlots(context.def_levels.data(), context.def_levels.length(), + path_info->repeated_ancestor_def_level); + } std::swap(builder_result.post_list_visited_elements, context.visited_elements); // If it is possible when processing lists that all lists where empty. In this // case no elements would have been added to post_list_visited_elements. By @@ -624,9 +845,18 @@ Status WritePath(ElementRange root_range, PathInfo* path_info, builder_result.post_list_visited_elements.push_back({0, 0}); } } else { - builder_result.post_list_visited_elements.push_back( - {0, builder_result.leaf_array->length()}); + if (!context.visited_elements.empty()) { + std::swap(builder_result.post_list_visited_elements, context.visited_elements); + } else { + builder_result.post_list_visited_elements.push_back( + {0, builder_result.leaf_array->length()}); + } builder_result.rep_levels = nullptr; + if (path_info->leaf_is_vector) { + builder_result.leaf_slot_count = + CountVectorSlots(context.def_levels.data(), context.def_levels.length(), + path_info->repeated_ancestor_def_level); + } } builder_result.def_levels = context.def_levels.data(); @@ -661,6 +891,7 @@ struct FixupVisitor { } void operator()(NullableNode& arg) { HandleIntermediateNode(arg); } + void operator()(VectorNullableNode&) {} void operator()(AllNullsTerminalNode& arg) { // Even though no processing happens past this point we @@ -671,6 +902,7 @@ struct FixupVisitor { void operator()(NullableTerminalNode&) {} void operator()(AllPresentTerminalNode&) {} + void operator()(NoLevelTerminalNode&) {} }; PathInfo Fixup(PathInfo info) { @@ -692,7 +924,9 @@ PathInfo Fixup(PathInfo info) { class PathBuilder { public: - explicit PathBuilder(bool start_nullable) : nullable_in_parent_(start_nullable) {} + PathBuilder(bool start_nullable, bool write_fixed_size_list_as_vector) + : nullable_in_parent_(start_nullable), + write_fixed_size_list_as_vector_(write_fixed_size_list_as_vector) {} template void AddTerminalInfo(const T& array) { info_.leaf_is_nullable = nullable_in_parent_; @@ -731,13 +965,18 @@ class PathBuilder { // Increment necessary due to empty lists. info_.max_def_level++; info_.max_rep_level++; + info_.repeated_ancestor_def_level = info_.max_def_level; // raw_value_offsets() accounts for any slice offset. ListPathNode> node( VarRangeSelector{array.raw_value_offsets()}, info_.max_rep_level, info_.max_def_level - 1); info_.path.emplace_back(std::move(node)); nullable_in_parent_ = array.list_type()->value_field()->nullable(); - return VisitInline(*array.values()); + const bool saved_nullable_group = nullable_group_since_repeated_; + nullable_group_since_repeated_ = false; + Status status = VisitInline(*array.values()); + nullable_group_since_repeated_ = saved_nullable_group; + return status; } Status Visit(const ::arrow::DictionaryArray& array) { @@ -790,31 +1029,100 @@ class PathBuilder { } Status Visit(const ::arrow::StructArray& array) { + const bool struct_nullable = nullable_in_parent_; MaybeAddNullable(array); PathInfo info_backup = info_; + const bool saved_nullable_group = nullable_group_since_repeated_; + nullable_group_since_repeated_ = nullable_group_since_repeated_ || struct_nullable; for (int x = 0; x < array.num_fields(); x++) { nullable_in_parent_ = array.type()->field(x)->nullable(); RETURN_NOT_OK(VisitInline(*array.field(x))); info_ = info_backup; } + nullable_group_since_repeated_ = saved_nullable_group; return Status::OK(); } Status Visit(const ::arrow::FixedSizeListArray& array) { - MaybeAddNullable(array); int32_t list_size = array.list_type()->list_size(); + if (write_fixed_size_list_as_vector_ && list_size > 0 && + !nullable_group_since_repeated_ && + IsSupportedVectorElementType(*array.value_type())) { + const bool element_nullable = array.list_type()->value_field()->nullable(); + const bool nested_vector_path = info_.leaf_is_vector; + info_.leaf_is_vector = true; + info_.leaf_vector_length *= list_size; + const bool parent_nullable = nullable_in_parent_; + if (parent_nullable) { + info_.max_def_level++; + } + // Overwritten by nested vectors so that the final value is the present + // level of the innermost vector. + info_.vector_values_def_level = info_.max_def_level; + const bool value_type_is_struct = array.value_type()->id() == ::arrow::Type::STRUCT; + const bool value_type_is_nested_vector = + array.value_type()->id() == ::arrow::Type::FIXED_SIZE_LIST; + const bool child_emits_present_def_levels = + element_nullable || value_type_is_struct || value_type_is_nested_vector || + nested_vector_path; + const bool child_records_visited_elements = value_type_is_nested_vector; + const bool has_vector_node = parent_nullable || element_nullable || + value_type_is_struct || value_type_is_nested_vector || + nested_vector_path || info_.max_rep_level > 0; + const bool vector_node_emits_def_levels = + has_vector_node && !child_emits_present_def_levels; + if (has_vector_node) { + info_.path.emplace_back(VectorNullableNode( + parent_nullable ? array.null_bitmap_data() : nullptr, array.offset(), + list_size, VectorSlotMultiplier(*array.value_type()), info_.max_def_level, + info_.max_rep_level, child_emits_present_def_levels, + child_records_visited_elements, + /*expand_rep_levels=*/!nested_vector_path)); + } + auto values = + array.values()->Slice(array.value_offset(0), array.length() * list_size); + if (!nested_vector_path && !element_nullable && + !::arrow::is_nested(*array.value_type())) { + info_.leaf_is_nullable = false; + info_.primitive_array = values; + if (!vector_node_emits_def_levels && info_.max_rep_level > 0 && + info_.max_def_level > 0) { + info_.path.emplace_back(AllPresentTerminalNode{info_.max_def_level}); + } else { + info_.path.emplace_back(NoLevelTerminalNode{}); + } + paths_.push_back(Fixup(info_)); + return Status::OK(); + } + nullable_in_parent_ = element_nullable; + const bool saved_nullable_group = nullable_group_since_repeated_; + nullable_group_since_repeated_ = false; + Status status = VisitInline(*values); + nullable_group_since_repeated_ = saved_nullable_group; + return status; + } + + MaybeAddNullable(array); // Technically we could encode fixed size lists with two level encodings // but since we always use 3 level encoding we increment def levels as // well. info_.max_def_level++; info_.max_rep_level++; + info_.repeated_ancestor_def_level = info_.max_def_level; + // def_level_if_empty is max_def_level - 1 ("list present but empty"), not + // max_def_level ("element present"); it is only reachable for zero-length + // fixed-size lists, where every present value is an empty list. info_.path.emplace_back(FixedSizeListNode(FixedSizedRangeSelector{list_size}, - info_.max_rep_level, info_.max_def_level)); + info_.max_rep_level, + info_.max_def_level - 1)); nullable_in_parent_ = array.list_type()->value_field()->nullable(); - if (array.offset() > 0) { - return VisitInline(*array.values()->Slice(array.value_offset(0))); - } - return VisitInline(*array.values()); + const bool saved_nullable_group = nullable_group_since_repeated_; + nullable_group_since_repeated_ = false; + Status status = array.offset() > 0 + ? VisitInline(*array.values()->Slice(array.value_offset(0))) + : VisitInline(*array.values()); + nullable_group_since_repeated_ = saved_nullable_group; + return status; } Status Visit(const ::arrow::ExtensionArray& array) { @@ -840,6 +1148,12 @@ class PathBuilder { PathInfo info_; std::vector paths_; bool nullable_in_parent_; + bool write_fixed_size_list_as_vector_; + // True when a nullable group (struct) sits between the nearest repeated + // ancestor (or the root) and the current position. A null group row maps + // one definition level to vector_length leaf slots, which the VECTOR level + // machinery does not support, so such fields fall back to LIST encoding. + bool nullable_group_since_repeated_ = false; }; Status PathBuilder::VisitInline(const Array& array) { @@ -883,8 +1197,10 @@ class MultipathLevelBuilderImpl : public MultipathLevelBuilder { // static ::arrow::Result> MultipathLevelBuilder::Make( - const ::arrow::Array& array, bool array_field_nullable) { - auto constructor = std::make_unique(array_field_nullable); + const ::arrow::Array& array, bool array_field_nullable, + bool write_fixed_size_list_as_vector) { + auto constructor = std::make_unique(array_field_nullable, + write_fixed_size_list_as_vector); RETURN_NOT_OK(VisitArrayInline(array, constructor.get())); return std::make_unique(array.data(), std::move(constructor)); @@ -894,8 +1210,12 @@ ::arrow::Result> MultipathLevelBuilder::M Status MultipathLevelBuilder::Write(const Array& array, bool array_field_nullable, ArrowWriteContext* context, MultipathLevelBuilder::CallbackFunction callback) { + const bool write_fixed_size_list_as_vector = + context->properties != nullptr && + context->properties->write_fixed_size_list_as_vector(); ARROW_ASSIGN_OR_RAISE(std::unique_ptr builder, - MultipathLevelBuilder::Make(array, array_field_nullable)); + MultipathLevelBuilder::Make(array, array_field_nullable, + write_fixed_size_list_as_vector)); for (int leaf_idx = 0; leaf_idx < builder->GetLeafCount(); leaf_idx++) { RETURN_NOT_OK(builder->Write(leaf_idx, context, callback)); } diff --git a/cpp/src/parquet/arrow/path_internal.h b/cpp/src/parquet/arrow/path_internal.h index 50d2bf24291a..e2158b4ad6af 100644 --- a/cpp/src/parquet/arrow/path_internal.h +++ b/cpp/src/parquet/arrow/path_internal.h @@ -94,6 +94,32 @@ struct MultipathLevelBuilderResult { /// Whether the leaf array is nullable. bool leaf_is_nullable; + + /// Whether this leaf is produced from an Arrow FixedSizeList being written as + /// Parquet VECTOR. For nullable VECTOR elements the physical Parquet leaf is + /// below the VECTOR node, so checking only the primitive schema node is not + /// sufficient. + bool leaf_is_vector = false; + + /// Fixed number of physical leaf slots per logical VECTOR value for this leaf. + int32_t leaf_vector_length = 1; + + /// Number of physical leaf slots for a VECTOR leaf: the definition levels at + /// or above the vector's stride definition level. Equals def_rep_level_count + /// unless ancestors made some vector values absent (empty or null lists, null + /// structs), which contribute one definition level and no slot. + int64_t leaf_slot_count = 0; + + /// Definition level after the last repeated ancestor of a VECTOR leaf (0 + /// when there is none). Definition levels below it mark absent ancestors + /// (empty or null lists) with no slot. + int16_t vector_repeated_ancestor_def_level = 0; + + /// Definition level at or above which a VECTOR leaf slot carries a value + /// from the leaf array; slots between the repeated-ancestor level and this + /// level belong to null vector values at some nesting depth and carry no + /// value. + int16_t vector_values_def_level = 0; }; /// \brief Logic for being able to write out nesting (rep/def level) data that is @@ -132,7 +158,8 @@ class PARQUET_EXPORT MultipathLevelBuilder { /// the array column as nullable (as determined by its type's parent /// field). static ::arrow::Result> Make( - const ::arrow::Array& array, bool array_field_nullable); + const ::arrow::Array& array, bool array_field_nullable, + bool write_fixed_size_list_as_vector = false); virtual ~MultipathLevelBuilder() = default; diff --git a/cpp/src/parquet/arrow/path_internal_test.cc b/cpp/src/parquet/arrow/path_internal_test.cc index 0145e889ddaf..09c6acd3732b 100644 --- a/cpp/src/parquet/arrow/path_internal_test.cc +++ b/cpp/src/parquet/arrow/path_internal_test.cc @@ -551,6 +551,88 @@ TEST_F(MultipathLevelBuilderTest, TestFixedSizeList) { EXPECT_THAT(results_[0].post_list_elements[0].end, Eq(6)); } +TEST_F(MultipathLevelBuilderTest, TestFixedSizeListExperimentalVector) { + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + arrow_properties_ = builder.build(); + context_ = ArrowWriteContext(default_memory_pool(), arrow_properties_.get()); + + auto entries = field("Entries", ::arrow::int64(), /*nullable=*/false); + auto list_type = fixed_size_list(entries, 2); + auto array = ArrayFromJSON(list_type, "[[1, 2], [3, 4], [5, 6]]"); + + ASSERT_OK( + MultipathLevelBuilder::Write(*array, /*nullable=*/false, &context_, callback_)); + ASSERT_THAT(results_, SizeIs(1)); + EXPECT_TRUE(results_[0].null_rep_levels); + EXPECT_TRUE(results_[0].null_def_levels); + ASSERT_THAT(results_[0].post_list_elements, SizeIs(1)); + EXPECT_EQ(results_[0].post_list_elements[0].start, 0); + EXPECT_EQ(results_[0].post_list_elements[0].end, 6); +} + +TEST_F(MultipathLevelBuilderTest, TestFixedSizeListExperimentalVectorNullableElements) { + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + arrow_properties_ = builder.build(); + context_ = ArrowWriteContext(default_memory_pool(), arrow_properties_.get()); + + auto entries = field("Entries", ::arrow::int64(), /*nullable=*/true); + auto list_type = fixed_size_list(entries, 2); + auto array = ArrayFromJSON(list_type, "[[1, null], [3, 4], [null, 6]]"); + + ASSERT_OK( + MultipathLevelBuilder::Write(*array, /*nullable=*/false, &context_, callback_)); + ASSERT_THAT(results_, SizeIs(1)); + results_[0].CheckLevelsWithNullRepLevels(std::vector{1, 0, 1, 1, 0, 1}); + ASSERT_THAT(results_[0].post_list_elements, SizeIs(1)); + EXPECT_EQ(results_[0].post_list_elements[0].start, 0); + EXPECT_EQ(results_[0].post_list_elements[0].end, 6); +} + +TEST_F(MultipathLevelBuilderTest, + TestFixedSizeListExperimentalVectorNullableRowsAndElements) { + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + arrow_properties_ = builder.build(); + context_ = ArrowWriteContext(default_memory_pool(), arrow_properties_.get()); + + auto entries = field("Entries", ::arrow::int64(), /*nullable=*/true); + auto list_type = fixed_size_list(entries, 2); + auto array = ArrayFromJSON(list_type, "[[1, null], null, [null, 6]]"); + + ASSERT_OK( + MultipathLevelBuilder::Write(*array, /*nullable=*/true, &context_, callback_)); + ASSERT_THAT(results_, SizeIs(1)); + results_[0].CheckLevelsWithNullRepLevels(std::vector{2, 1, 0, 0, 1, 2}); + ASSERT_THAT(results_[0].post_list_elements, SizeIs(2)); + EXPECT_EQ(results_[0].post_list_elements[0].start, 0); + EXPECT_EQ(results_[0].post_list_elements[0].end, 2); + EXPECT_EQ(results_[0].post_list_elements[1].start, 4); + EXPECT_EQ(results_[0].post_list_elements[1].end, 6); +} + +TEST_F(MultipathLevelBuilderTest, TestFixedSizeListExperimentalVectorNullable) { + ArrowWriterProperties::Builder builder; + builder.enable_experimental_vector_encoding(); + arrow_properties_ = builder.build(); + context_ = ArrowWriteContext(default_memory_pool(), arrow_properties_.get()); + + auto entries = field("Entries", ::arrow::int64(), /*nullable=*/false); + auto list_type = fixed_size_list(entries, 2); + auto array = ArrayFromJSON(list_type, "[[1, 2], null, [5, 6]]"); + + ASSERT_OK( + MultipathLevelBuilder::Write(*array, /*nullable=*/true, &context_, callback_)); + ASSERT_THAT(results_, SizeIs(1)); + results_[0].CheckLevelsWithNullRepLevels(std::vector{1, 1, 0, 0, 1, 1}); + ASSERT_THAT(results_[0].post_list_elements, SizeIs(2)); + EXPECT_EQ(results_[0].post_list_elements[0].start, 0); + EXPECT_EQ(results_[0].post_list_elements[0].end, 2); + EXPECT_EQ(results_[0].post_list_elements[1].start, 4); + EXPECT_EQ(results_[0].post_list_elements[1].end, 6); +} + TEST_F(MultipathLevelBuilderTest, TestFixedSizeListMissingMiddleHasTwoVisitedRanges) { auto entries = field("Entries", ::arrow::int64(), /*nullable=*/false); auto list_type = fixed_size_list(entries, 2); diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index a60af69aec9f..a97762d9a1b2 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -26,6 +26,8 @@ #include #include "arrow/array.h" +#include "arrow/array/concatenate.h" +#include "arrow/array/util.h" #include "arrow/buffer.h" #include "arrow/extension_type.h" #include "arrow/io/memory.h" @@ -273,10 +275,21 @@ class FileReaderImpl : public FileReader { // TODO(wesm): This calculation doesn't make much sense when we have repeated // schema nodes int64_t records_to_read = 0; + DCHECK(dynamic_cast(reader) != nullptr); + const bool has_repeated_child = + static_cast(reader)->IsOrHasRepeatedChild(); for (auto row_group : row_groups) { - // Can throw exception - records_to_read += - reader_->metadata()->RowGroup(row_group)->ColumnChunk(i)->num_values(); + // Can throw exception. ColumnReader::NextBatch takes logical parent + // records. For readers without repeated children (flat and VECTOR + // columns) the row group row count is an exact bound; for VECTOR columns + // in particular, ColumnChunk num_values would over-count records by the + // vector length. + if (has_repeated_child) { + records_to_read += + reader_->metadata()->RowGroup(row_group)->ColumnChunk(i)->num_values(); + } else { + records_to_read += reader_->metadata()->RowGroup(row_group)->num_rows(); + } } #ifdef ARROW_WITH_OPENTELEMETRY std::string column_name = reader_->metadata()->schema()->Column(i)->name(); @@ -706,6 +719,196 @@ class PARQUET_NO_EXPORT FixedSizeListReader : public ListReader { } }; +// Reads Parquet VECTOR columns into Arrow FixedSizeList arrays. +// +// VECTOR stores one definition level per element even though the public Arrow result is a +// single FixedSizeList slot per row. For nullable VECTOR rows, the child reader therefore +// materializes spaced child slots and this reader collapses each vector's per-element def +// levels back into a parent validity bitmap. +class PARQUET_NO_EXPORT VectorFixedSizeListReader : public ColumnReaderImpl { + public: + VectorFixedSizeListReader(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)), + level_info_(level_info), + item_reader_(std::move(child_reader)), + list_size_(checked_cast(*field_->type()) + .list_size()) {} + + Status GetDefLevels(const int16_t** data, int64_t* length) override { + if (collapsed_def_levels_.empty()) { + *data = nullptr; + *length = vector_rows_; + } else { + *data = collapsed_def_levels_.data(); + *length = static_cast(collapsed_def_levels_.size()); + } + return Status::OK(); + } + + Status GetRepLevels(const int16_t** data, int64_t* length) override { + if (collapsed_rep_levels_.empty()) { + *data = nullptr; + *length = vector_rows_; + } else { + *data = collapsed_rep_levels_.data(); + *length = static_cast(collapsed_rep_levels_.size()); + } + return Status::OK(); + } + + bool IsOrHasRepeatedChild() const final { return false; } + + Status LoadBatch(int64_t number_of_records) final { + vector_rows_ = 0; + collapsed_def_levels_.clear(); + collapsed_rep_levels_.clear(); + vector_def_levels_.clear(); + const int64_t child_records_to_read = + level_info_.rep_level > 0 ? number_of_records : number_of_records * list_size_; + RETURN_NOT_OK(item_reader_->LoadBatch(child_records_to_read)); + + const int16_t* def_levels = nullptr; + int64_t num_levels = 0; + RETURN_NOT_OK(item_reader_->GetDefLevels(&def_levels, &num_levels)); + const int16_t* rep_levels = nullptr; + int64_t num_rep_levels = 0; + RETURN_NOT_OK(item_reader_->GetRepLevels(&rep_levels, &num_rep_levels)); + const bool collapse_rep_levels = level_info_.rep_level > 0 && rep_levels != nullptr; + if (collapse_rep_levels && num_rep_levels != num_levels) { + return Status::Invalid("VECTOR child produced ", num_rep_levels, + " repetition levels for ", num_levels, " definition levels"); + } + if (def_levels == nullptr || num_levels == 0) { + // No level information; the row count comes from the child array length. + return Status::OK(); + } + + // Existing vector values (present, or null when the vector field is + // nullable) occupy exactly list_size_ definition levels. Levels below + // the repeated-ancestor level mark ancestors that made the vector value + // absent (empty or null lists); they occupy a single level entry with no + // slots and pass through unchanged for the parent reader. + const int16_t stride_def_level = level_info_.repeated_ancestor_def_level; + const size_t reserve_hint = static_cast(num_levels / list_size_) + 1; + collapsed_def_levels_.reserve(reserve_hint); + vector_def_levels_.reserve(reserve_hint); + if (collapse_rep_levels) { + collapsed_rep_levels_.reserve(reserve_hint); + } + int64_t i = 0; + while (i < num_levels) { + if (def_levels[i] < stride_def_level) { + collapsed_def_levels_.push_back(def_levels[i]); + if (collapse_rep_levels) { + collapsed_rep_levels_.push_back(rep_levels[i]); + } + ++i; + continue; + } + if (i + list_size_ > num_levels) { + return Status::Invalid("VECTOR column ended mid-vector: ", num_levels - i, + " definition levels remaining for list_size=", list_size_); + } + const bool is_present = def_levels[i] >= level_info_.def_level; + for (int32_t k = 1; k < list_size_; ++k) { + const bool slot_is_present = def_levels[i + k] >= level_info_.def_level; + if (slot_is_present != is_present) { + return Status::Invalid( + "VECTOR parent validity changed within one fixed-size vector at row ", + vector_rows_, + "; null VECTOR rows must still emit exactly list_size null child " + "slots and all slots for one parent must agree on parent " + "validity"); + } + } + const int16_t collapsed = is_present ? level_info_.def_level : def_levels[i]; + collapsed_def_levels_.push_back(collapsed); + vector_def_levels_.push_back(collapsed); + if (collapse_rep_levels) { + collapsed_rep_levels_.push_back(rep_levels[i]); + } + ++vector_rows_; + i += list_size_; + } + return Status::OK(); + } + + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override { + std::shared_ptr child_out; + RETURN_NOT_OK(item_reader_->BuildArray(length_upper_bound * list_size_, &child_out)); + ARROW_ASSIGN_OR_RAISE(std::shared_ptr child_data, + ChunksToSingle(*child_out)); + std::shared_ptr child_array = ::arrow::MakeArray(child_data); + if (vector_rows_ == 0 && vector_def_levels_.empty()) { + // No level information was loaded; derive the row count from the child. + if (child_array->length() % list_size_ != 0) { + return Status::Invalid("VECTOR FixedSizeList child length ", + child_array->length(), + " was not divisible by list_size=", list_size_); + } + vector_rows_ = child_array->length() / list_size_; + } + + if (child_array->length() != vector_rows_ * list_size_) { + return Status::Invalid("VECTOR FixedSizeList child length ", child_array->length(), + " did not match expected ", vector_rows_ * list_size_); + } + + std::shared_ptr validity_buffer; + int64_t null_count = 0; + if (field_->nullable()) { + ARROW_ASSIGN_OR_RAISE( + validity_buffer, + AllocateResizableBuffer(bit_util::BytesForBits(vector_rows_), ctx_->pool)); + memset(validity_buffer->mutable_data(), 0, + static_cast(bit_util::BytesForBits(vector_rows_))); + if (vector_def_levels_.empty()) { + bit_util::SetBitsTo(validity_buffer->mutable_data(), 0, vector_rows_, true); + } else { + for (int64_t row = 0; row < vector_rows_; ++row) { + if (vector_def_levels_[row] == level_info_.def_level) { + bit_util::SetBit(validity_buffer->mutable_data(), row); + } else { + ++null_count; + } + } + } + validity_buffer->ZeroPadding(); + } + + auto data = std::make_shared( + field_->type(), vector_rows_, + std::vector>{null_count > 0 ? validity_buffer : nullptr}, + null_count); + data->child_data.push_back(child_data); + *out = std::make_shared(::arrow::MakeArray(std::move(data))); + return Status::OK(); + } + + const std::shared_ptr field() override { return field_; } + + private: + std::shared_ptr ctx_; + std::shared_ptr field_; + ::parquet::internal::LevelInfo level_info_; + std::unique_ptr item_reader_; + int32_t list_size_; + // Number of existing vector values (present or null); absent ancestors do + // not count. + int64_t vector_rows_ = 0; + // One entry per logical entry (existing vector or absent ancestor), for the + // parent reader. + std::vector collapsed_def_levels_; + std::vector collapsed_rep_levels_; + // One entry per existing vector value, for this reader's validity bitmap. + std::vector vector_def_levels_; +}; + class PARQUET_NO_EXPORT StructReader : public ColumnReaderImpl { public: explicit StructReader(std::shared_ptr ctx, @@ -884,8 +1087,12 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& arrow_f } std::unique_ptr input( ctx->iterator_factory(field.column_index, ctx->reader)); - *out = std::make_unique(ctx, arrow_field, std::move(input), - field.level_info); + auto leaf_field = arrow_field; + if (field.is_vector && field.level_info.def_level > 0) { + leaf_field = leaf_field->WithNullable(true); + } + *out = + std::make_unique(ctx, leaf_field, std::move(input), field.level_info); } else if (type_id == ::arrow::Type::LIST || type_id == ::arrow::Type::MAP || type_id == ::arrow::Type::FIXED_SIZE_LIST || type_id == ::arrow::Type::LARGE_LIST) { @@ -950,8 +1157,16 @@ Status GetReader(const SchemaField& field, const std::shared_ptr& arrow_f list_field->WithType(::arrow::fixed_size_list(reader_child_type, list_size)); } - *out = std::make_unique(ctx, list_field, field.level_info, - std::move(child_reader)); + if (field.is_vector) { + if (child->is_leaf() && child->column_index >= 0 && list_field->nullable()) { + DCHECK(child_reader->field()->nullable()); + } + *out = std::make_unique( + ctx, list_field, field.level_info, std::move(child_reader)); + } else { + *out = std::make_unique(ctx, list_field, field.level_info, + std::move(child_reader)); + } } else { return Status::UnknownError("Unknown list type: ", field.field->ToString()); } diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index 9c4c462c6b8c..1742bde0232c 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -30,6 +30,7 @@ #include "arrow/result.h" #include "arrow/status.h" #include "arrow/type.h" +#include "arrow/type_traits.h" #include "arrow/util/base64.h" #include "arrow/util/checked_cast.h" #include "arrow/util/key_value_metadata.h" @@ -91,9 +92,15 @@ Result> MakeArrowList( } } +// nullable_group_since_repeated is true when a nullable group (struct) sits +// between the nearest repeated ancestor (or the root) and the field being +// converted. A null group row maps one definition level to vector_length +// leaf slots, which the VECTOR level machinery does not support, so such +// FixedSizeList fields fall back to the standard LIST encoding. Status FieldToNode(const std::string& name, const std::shared_ptr& field, const WriterProperties& properties, - const ArrowWriterProperties& arrow_properties, NodePtr* out); + const ArrowWriterProperties& arrow_properties, NodePtr* out, + bool nullable_group_since_repeated = false); Status ListToNode(const std::shared_ptr<::arrow::BaseListType>& type, const std::string& name, bool nullable, int field_id, @@ -111,6 +118,83 @@ Status ListToNode(const std::shared_ptr<::arrow::BaseListType>& type, return Status::OK(); } +bool IsSupportedVectorStructNode(const Node& node) { + if (node.is_primitive()) { + return true; + } + if (!node.is_group() || node.is_repeated()) { + return false; + } + const auto& group = checked_cast(node); + if (group.logical_type() != nullptr && !group.logical_type()->is_none()) { + return false; + } + for (int i = 0; i < group.field_count(); ++i) { + if (!IsSupportedVectorStructNode(*group.field(i))) { + return false; + } + } + return true; +} + +Status ValidateSupportedVectorStructNode(const Node& node) { + if (IsSupportedVectorStructNode(node)) { + return Status::OK(); + } + return Status::NotImplemented( + "VECTOR elements only support primitive, nested VECTOR, or nested struct fields; " + "repeated/list/map descendants are not supported"); +} + +// Option B mapping for Arrow FixedSizeList -> Parquet VECTOR. +// +// The schema always uses a 3-level structure mirroring LIST, regardless of +// vector and element nullability: +// +// group (VECTOR) { +// vector group list [N] { +// element; +// } +// } +// +// VECTOR is a structural repetition type: vector nullability is carried by the +// outer group, the VECTOR-repeated middle group carries vector_length and +// contributes no repetition/definition level, and element nullability is +// carried by the element node. The outer group is annotated with the VECTOR +// logical type to make it explicit that it represents a vector value rather +// than, e.g., a single-field struct containing a vector, mirroring how LIST +// annotates list-shaped groups. +Status FixedSizeListToNode(const std::shared_ptr<::arrow::FixedSizeListType>& type, + const std::string& name, bool nullable, int field_id, + const WriterProperties& properties, + const ArrowWriterProperties& arrow_properties, NodePtr* out) { + if (type->list_size() <= 0) { + return Status::NotImplemented( + "VECTOR repetition does not support zero-length FixedSizeList values"); + } + if (!IsSupportedVectorElementType(*type->value_type())) { + return Status::NotImplemented( + "VECTOR repetition only supports fixed-width primitive, nested FixedSizeList, " + "or struct FixedSizeList elements"); + } + const auto& value_field = type->value_field(); + const std::string value_name = + arrow_properties.compliant_nested_types() ? "element" : value_field->name(); + + NodePtr element; + RETURN_NOT_OK( + FieldToNode(value_name, value_field, properties, arrow_properties, &element)); + if (type->value_type()->id() == ::arrow::Type::STRUCT) { + RETURN_NOT_OK(ValidateSupportedVectorStructNode(*element)); + } + NodePtr vector = + GroupNode::Make("list", Repetition::VECTOR, {element}, + /*logical_type=*/nullptr, /*field_id=*/-1, type->list_size()); + *out = GroupNode::Make(name, RepetitionFromNullable(nullable), {vector}, + LogicalType::Vector(), field_id); + return Status::OK(); +} + Status MapToNode(const std::shared_ptr<::arrow::MapType>& type, const std::string& name, bool nullable, int field_id, const WriterProperties& properties, const ArrowWriterProperties& arrow_properties, NodePtr* out) { @@ -153,12 +237,14 @@ Status VariantToNode( Status StructToNode(const std::shared_ptr<::arrow::StructType>& type, const std::string& name, bool nullable, int field_id, const WriterProperties& properties, - const ArrowWriterProperties& arrow_properties, NodePtr* out) { + const ArrowWriterProperties& arrow_properties, NodePtr* out, + bool nullable_group_since_repeated) { std::vector children(type->num_fields()); if (type->num_fields() != 0) { for (int i = 0; i < type->num_fields(); i++) { RETURN_NOT_OK(FieldToNode(type->field(i)->name(), type->field(i), properties, - arrow_properties, &children[i])); + arrow_properties, &children[i], + nullable_group_since_repeated || nullable)); } } else { // XXX (ARROW-10928) We could add a dummy primitive node but that would @@ -315,7 +401,8 @@ int FieldIdFromMetadata( Status FieldToNode(const std::string& name, const std::shared_ptr& field, const WriterProperties& properties, - const ArrowWriterProperties& arrow_properties, NodePtr* out) { + const ArrowWriterProperties& arrow_properties, NodePtr* out, + bool nullable_group_since_repeated) { std::shared_ptr logical_type = LogicalType::None(); ParquetType::type type; Repetition::type repetition = RepetitionFromNullable(field->nullable()); @@ -451,9 +538,34 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, case ArrowTypeId::STRUCT: { auto struct_type = std::static_pointer_cast<::arrow::StructType>(field->type()); return StructToNode(struct_type, name, field->nullable(), field_id, properties, - arrow_properties, out); + arrow_properties, out, nullable_group_since_repeated); + } + case ArrowTypeId::FIXED_SIZE_LIST: { + auto list_type = + std::static_pointer_cast<::arrow::FixedSizeListType>(field->type()); + // Experimental VECTOR encoding is opportunistic: use it for supported + // FixedSizeList fields, and preserve writability by falling back to the + // standard LIST encoding for unsupported fields. This lets schemas mix + // VECTOR-friendly dense numeric/struct vectors with FixedSizeList values + // that still need LIST machinery (for example strings). + if (arrow_properties.write_fixed_size_list_as_vector() && + !nullable_group_since_repeated && list_type->list_size() > 0 && + IsSupportedVectorElementType(*list_type->value_type())) { + Status vector_status = + FixedSizeListToNode(list_type, name, field->nullable(), field_id, properties, + arrow_properties, out); + if (vector_status.ok()) { + return vector_status; + } + if (!vector_status.IsNotImplemented()) { + return vector_status; + } + } + auto base_list_type = + std::static_pointer_cast<::arrow::BaseListType>(field->type()); + return ListToNode(base_list_type, name, field->nullable(), field_id, properties, + arrow_properties, out); } - case ArrowTypeId::FIXED_SIZE_LIST: case ArrowTypeId::LARGE_LIST: case ArrowTypeId::LIST: { auto list_type = std::static_pointer_cast<::arrow::BaseListType>(field->type()); @@ -467,7 +579,8 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, static_cast(*field->type()); std::shared_ptr<::arrow::Field> unpacked_field = ::arrow::field( name, dict_type.value_type(), field->nullable(), field->metadata()); - return FieldToNode(name, unpacked_field, properties, arrow_properties, out); + return FieldToNode(name, unpacked_field, properties, arrow_properties, out, + nullable_group_since_repeated); } case ArrowTypeId::EXTENSION: { auto ext_type = std::static_pointer_cast<::arrow::ExtensionType>(field->type()); @@ -498,7 +611,8 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, std::shared_ptr<::arrow::Field> storage_field = ::arrow::field( name, ext_type->storage_type(), field->nullable(), field->metadata()); - return FieldToNode(name, storage_field, properties, arrow_properties, out); + return FieldToNode(name, storage_field, properties, arrow_properties, out, + nullable_group_since_repeated); } case ArrowTypeId::MAP: { auto map_type = std::static_pointer_cast<::arrow::MapType>(field->type()); @@ -574,6 +688,13 @@ Status PopulateLeaf(int column_index, const std::shared_ptr& field, return Status::OK(); } +void MarkVectorSubtree(SchemaField* field) { + field->is_vector = true; + for (auto& child : field->children) { + MarkVectorSubtree(&child); + } +} + // Special case mentioned in the format spec: // If the name is array or uses the parent's name with `_tuple` appended, // this should be: @@ -599,7 +720,7 @@ Status GroupToStruct(const GroupNode& node, LevelInfo current_levels, arrow_fields.push_back(out->children[i].field); } auto struct_type = ::arrow::struct_(arrow_fields); - if (ctx->properties.get_arrow_extensions_enabled() && + if (ctx->properties.get_arrow_extensions_enabled() && node.logical_type() != nullptr && node.logical_type()->is_variant()) { auto extension_type = ::arrow::GetExtensionType("arrow.parquet.variant"); if (extension_type) { @@ -618,6 +739,99 @@ Status ListToSchemaField(const GroupNode& group, LevelInfo current_levels, SchemaTreeContext* ctx, const SchemaField* parent, SchemaField* out); +bool IsVectorGroup(const GroupNode& group) { + // The VECTOR logical type annotation is required to identify a vector group; + // this is what distinguishes a vector field from a struct-like group that + // happens to contain a VECTOR-repeated child. An unannotated group with a + // VECTOR-repeated child is rejected during conversion (see + // GroupToSchemaField) rather than silently interpreted as a vector. + return group.logical_type() != nullptr && group.logical_type()->is_vector(); +} + +// Reconstructs an Arrow FixedSizeList from the canonical 3-level vector form: +// +// group (VECTOR) { +// vector group list [N] { +// element; +// } +// } +// +// The element may be a primitive, a struct group, or a nested vector group +// (itself in canonical form). +Status VectorToSchemaField(const GroupNode& group, LevelInfo current_levels, + SchemaTreeContext* ctx, const SchemaField* parent, + SchemaField* out) { + if (group.is_repeated() || group.is_vector()) { + return Status::NotImplemented( + "Repeated VECTOR groups are not supported in the current prototype"); + } + if (group.field_count() != 1) { + return Status::Invalid("VECTOR groups must have a single VECTOR-repeated child"); + } + const Node& child = *group.field(0); + if (!child.is_vector()) { + return Status::Invalid("VECTOR groups must contain a VECTOR-repeated child"); + } + if (!child.is_group()) { + return Status::Invalid( + "VECTOR-repeated nodes must be groups containing a single element field"); + } + const auto& vector_group = static_cast(child); + if (vector_group.field_count() != 1) { + return Status::Invalid("VECTOR-repeated groups must have a single element child"); + } + const Node& element = *vector_group.field(0); + if (element.is_repeated() || element.is_vector()) { + return Status::Invalid("VECTOR elements must be required or optional"); + } + + if (group.is_optional()) { + current_levels.IncrementOptional(); + } + const LevelInfo vector_level = current_levels; + + out->children.resize(1); + SchemaField* child_field = &out->children[0]; + ctx->LinkParent(out, parent); + ctx->LinkParent(child_field, out); + + if (element.is_primitive()) { + const bool element_nullable = element.is_optional(); + if (element_nullable) { + current_levels.IncrementOptional(); + } + const auto& primitive_node = static_cast(element); + int column_index = ctx->schema->GetColumnIndex(primitive_node); + ARROW_ASSIGN_OR_RAISE(std::shared_ptr type, + GetTypeForNode(column_index, primitive_node, ctx)); + auto item_field = ::arrow::field(element.name(), type, element_nullable, + FieldIdMetadata(element.field_id())); + RETURN_NOT_OK( + PopulateLeaf(column_index, item_field, current_levels, ctx, out, child_field)); + } else { + const auto& element_group = static_cast(element); + if (IsVectorGroup(element_group)) { + // Nested vector element; the recursive call handles element nullability. + RETURN_NOT_OK( + VectorToSchemaField(element_group, current_levels, ctx, out, child_field)); + } else { + RETURN_NOT_OK(ValidateSupportedVectorStructNode(element_group)); + if (element_group.is_optional()) { + current_levels.IncrementOptional(); + } + RETURN_NOT_OK(GroupToStruct(element_group, current_levels, ctx, out, child_field)); + } + } + + MarkVectorSubtree(child_field); + out->field = ::arrow::field( + group.name(), ::arrow::fixed_size_list(child_field->field, child.vector_length()), + group.is_optional(), FieldIdMetadata(group.field_id())); + out->level_info = vector_level; + out->is_vector = true; + return Status::OK(); +} + Status MapToSchemaField(const GroupNode& group, LevelInfo current_levels, SchemaTreeContext* ctx, const SchemaField* parent, SchemaField* out) { @@ -842,6 +1056,14 @@ Status ListToSchemaField(const GroupNode& group, LevelInfo current_levels, Status GroupToSchemaField(const GroupNode& node, LevelInfo current_levels, SchemaTreeContext* ctx, const SchemaField* parent, SchemaField* out) { + if (IsVectorGroup(node)) { + return VectorToSchemaField(node, current_levels, ctx, parent, out); + } + if (node.is_vector()) { + return Status::Invalid( + "VECTOR-repeated groups must be the single child of a group annotated with " + "the VECTOR logical type"); + } if (node.logical_type()->is_list()) { return ListToSchemaField(node, current_levels, ctx, parent, out); } else if (node.logical_type()->is_map()) { @@ -903,6 +1125,13 @@ Status NodeToSchemaField(const Node& node, LevelInfo current_levels, int column_index = ctx->schema->GetColumnIndex(primitive_node); ARROW_ASSIGN_OR_RAISE(std::shared_ptr type, GetTypeForNode(column_index, primitive_node, ctx)); + if (node.is_vector()) { + // VECTOR repetition directly on a primitive is not part of the canonical + // 3-level vector structure; require the VECTOR group form instead. + return Status::Invalid( + "VECTOR repetition on primitive nodes is not supported; vectors must use " + "the 3-level VECTOR group structure"); + } if (node.is_repeated()) { // One-level list encoding, e.g. // a: repeated int32; diff --git a/cpp/src/parquet/arrow/schema.h b/cpp/src/parquet/arrow/schema.h index dd60fde43422..7ab0cc31dd05 100644 --- a/cpp/src/parquet/arrow/schema.h +++ b/cpp/src/parquet/arrow/schema.h @@ -96,6 +96,11 @@ struct PARQUET_EXPORT SchemaField { parquet::internal::LevelInfo level_info; + // True when this Arrow field is backed by a Parquet VECTOR node. VECTOR may + // be represented either directly as a primitive leaf (non-nullable elements) + // or as an intermediate VECTOR group containing a nullable element leaf. + bool is_vector = false; + bool is_leaf() const { return column_index != -1; } }; diff --git a/cpp/src/parquet/arrow/schema_internal.cc b/cpp/src/parquet/arrow/schema_internal.cc index 2e8cf764b27f..b3a6a8b3f768 100644 --- a/cpp/src/parquet/arrow/schema_internal.cc +++ b/cpp/src/parquet/arrow/schema_internal.cc @@ -20,6 +20,7 @@ #include "arrow/extension/json.h" #include "arrow/extension/uuid.h" #include "arrow/type.h" +#include "arrow/type_traits.h" #include "arrow/util/key_value_metadata.h" #include "arrow/util/logging.h" #include "arrow/util/string.h" @@ -39,6 +40,62 @@ using ::arrow::internal::checked_cast; namespace { +// Whether a type may appear inside a struct VECTOR element. This must mirror +// the writer-side schema validation (ValidateSupportedVectorStructNode): +// list/map descendants require repetition levels and force the whole column +// back to the standard LIST encoding. Both the schema conversion and the +// write-path level generation gate on this predicate so they cannot disagree +// about the VECTOR-vs-LIST fallback. +bool IsSupportedVectorStructFieldType(const ::arrow::DataType& type) { + switch (type.id()) { + case ::arrow::Type::LIST: + case ::arrow::Type::LARGE_LIST: + case ::arrow::Type::LIST_VIEW: + case ::arrow::Type::LARGE_LIST_VIEW: + case ::arrow::Type::MAP: + return false; + case ::arrow::Type::FIXED_SIZE_LIST: + // A fixed-size list nested inside a struct element would give the + // struct's leaves different per-row multiplicities, which the write-path + // level generation does not support; fall back to LIST encoding. + return false; + case ::arrow::Type::STRUCT: { + const auto& struct_type = checked_cast(type); + for (const auto& field : struct_type.fields()) { + if (!IsSupportedVectorStructFieldType(*field->type())) { + return false; + } + } + return true; + } + case ::arrow::Type::DICTIONARY: + return IsSupportedVectorStructFieldType( + *checked_cast(type).value_type()); + case ::arrow::Type::EXTENSION: + return IsSupportedVectorStructFieldType( + *checked_cast(type).storage_type()); + default: + return true; + } +} + +} // namespace + +bool IsSupportedVectorElementType(const ::arrow::DataType& type) { + if (type.id() == ::arrow::Type::FIXED_SIZE_LIST) { + const auto& list_type = checked_cast(type); + return list_type.list_size() > 0 && + IsSupportedVectorElementType(*list_type.value_type()); + } + if (type.id() == ::arrow::Type::STRUCT) { + return IsSupportedVectorStructFieldType(type); + } + return !::arrow::is_nested(type) && ::arrow::is_fixed_width(type) && + type.id() != ::arrow::Type::DICTIONARY && type.id() != ::arrow::Type::EXTENSION; +} + +namespace { + Result> MakeArrowDecimal(const LogicalType& logical_type, bool smallest_decimal_enabled) { const auto& decimal = checked_cast(logical_type); diff --git a/cpp/src/parquet/arrow/schema_internal.h b/cpp/src/parquet/arrow/schema_internal.h index 09ad891aad3b..48b20338af50 100644 --- a/cpp/src/parquet/arrow/schema_internal.h +++ b/cpp/src/parquet/arrow/schema_internal.h @@ -40,4 +40,6 @@ Result> GetArrowType( const ArrowReaderProperties& reader_properties, const std::shared_ptr& metadata = nullptr); +bool IsSupportedVectorElementType(const ::arrow::DataType& type); + } // namespace parquet::arrow diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index 4b2b06e5e097..df32d4d43ea5 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -20,12 +20,14 @@ #include #include #include +#include #include #include #include #include #include "arrow/array.h" +#include "arrow/array/concatenate.h" #include "arrow/extension_type.h" #include "arrow/ipc/writer.h" #include "arrow/record_batch.h" @@ -107,6 +109,101 @@ bool HasNullableRoot(const SchemaManifest& schema_manifest, return nullable; } +int64_t CountVisitedValues(const std::vector& visited_elements) { + return std::accumulate( + visited_elements.begin(), visited_elements.end(), int64_t{0}, + [](int64_t total, const ElementRange& range) { return total + range.Size(); }); +} + +Result> MaterializeVectorLeafArray( + const MultipathLevelBuilderResult& result, ArrowWriteContext* ctx, + bool* leaf_is_nullable) { + const auto& visited_elements = result.post_list_visited_elements; + DCHECK_GT(visited_elements.size(), 0); + + // Nullable VECTOR rows still need one child slot per vector element so that the + // generic spaced leaf writer can align leaf slots with the VECTOR def levels. Arrow + // FixedSizeList already stores child slots for null parent rows, so preserve a + // zero-copy slice over the complete child range and let WriteArrow's def-level-derived + // validity bitmap suppress values belonging to null vector rows. + // The materialized array must hold one entry per physical leaf slot; absent + // ancestors (empty or null lists, null structs) contribute definition levels + // but no slots, so the target is leaf_slot_count, not def_rep_level_count. + if (CountVisitedValues(visited_elements) == result.leaf_slot_count && + visited_elements.size() == 1) { + const ElementRange& range = visited_elements[0]; + auto values = result.leaf_array->Slice(range.start, range.Size()); + if (values->null_count() != 0) { + *leaf_is_nullable = true; + } + return values; + } + + // General path, driven by the definition levels (the only stream that + // distinguishes the three kinds of entries): + // def < ancestor level: absent ancestor, no slot, no value; + // def < values level: slot of a null vector value, no value; + // def >= values level: slot carrying the next value from the visited + // ranges (in order; ranges skip null vector slots). + const int16_t ancestor_def_level = result.vector_repeated_ancestor_def_level; + const int16_t values_def_level = result.vector_values_def_level; + const int16_t* def_levels = result.def_levels; + ::arrow::ArrayVector parts; + size_t range_idx = 0; + int64_t range_used = 0; + int64_t pending_nulls = 0; + auto flush_nulls = [&]() -> Status { + if (pending_nulls > 0) { + ARROW_ASSIGN_OR_RAISE(auto null_values, + ::arrow::MakeArrayOfNull(result.leaf_array->type(), + pending_nulls, ctx->memory_pool)); + parts.push_back(std::move(null_values)); + pending_nulls = 0; + } + return Status::OK(); + }; + int64_t i = 0; + while (i < result.def_rep_level_count) { + const int16_t def_level = def_levels[i]; + if (def_level < ancestor_def_level) { + ++i; + continue; + } + if (def_level < values_def_level) { + ++pending_nulls; + ++i; + continue; + } + // Batch a run of value slots from the current visited range. + int64_t run = 1; + while (i + run < result.def_rep_level_count && + def_levels[i + run] >= values_def_level && + range_idx < visited_elements.size() && + range_used + run < visited_elements[range_idx].Size()) { + ++run; + } + if (range_idx >= visited_elements.size() || + range_used + run > visited_elements[range_idx].Size()) { + return Status::Invalid("VECTOR leaf values out of sync with definition levels"); + } + RETURN_NOT_OK(flush_nulls()); + const ElementRange& range = visited_elements[range_idx]; + parts.push_back(result.leaf_array->Slice(range.start + range_used, run)); + range_used += run; + if (range_used == static_cast(range.Size())) { + ++range_idx; + range_used = 0; + } + i += run; + } + RETURN_NOT_OK(flush_nulls()); + *leaf_is_nullable = true; + if (parts.empty()) { + return result.leaf_array->Slice(0, 0); + } + return ::arrow::Concatenate(parts, ctx->memory_pool); +} + Status GetSchemaMetadata(const ::arrow::Schema& schema, ::arrow::MemoryPool* pool, const ArrowWriterProperties& properties, std::shared_ptr* out) { @@ -169,17 +266,22 @@ class ArrowColumnWriterV2 { leaf_idx, ctx, [&](const MultipathLevelBuilderResult& result) { size_t visited_component_size = result.post_list_visited_elements.size(); DCHECK_GT(visited_component_size, 0); - if (visited_component_size != 1) { + std::shared_ptr values_array; + bool leaf_is_nullable = result.leaf_is_nullable; + if (result.leaf_is_vector) { + ARROW_ASSIGN_OR_RAISE(values_array, MaterializeVectorLeafArray( + result, ctx, &leaf_is_nullable)); + } else if (visited_component_size == 1) { + const ElementRange& range = result.post_list_visited_elements[0]; + values_array = result.leaf_array->Slice(range.start, range.Size()); + } else { return Status::NotImplemented( "Lists with non-zero length null components are not supported"); } - const ElementRange& range = result.post_list_visited_elements[0]; - std::shared_ptr 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, result.leaf_is_nullable); + ctx, leaf_is_nullable); })); } @@ -201,7 +303,7 @@ class ArrowColumnWriterV2 { static ::arrow::Result> Make( const ChunkedArray& data, int64_t offset, const int64_t size, const SchemaManifest& schema_manifest, RowGroupWriter* row_group_writer, - int start_leaf_column_index = -1) { + bool write_fixed_size_list_as_vector, int start_leaf_column_index = -1) { int64_t absolute_position = 0; int chunk_index = 0; int64_t chunk_offset = 0; @@ -271,8 +373,10 @@ class ArrowColumnWriterV2 { std::shared_ptr array_to_write = chunk.Slice(chunk_offset, chunk_write_size); if (array_to_write->length() > 0) { - ARROW_ASSIGN_OR_RAISE(std::unique_ptr builder, - MultipathLevelBuilder::Make(*array_to_write, is_nullable)); + ARROW_ASSIGN_OR_RAISE( + std::unique_ptr builder, + MultipathLevelBuilder::Make(*array_to_write, is_nullable, + write_fixed_size_list_as_vector)); if (leaf_count != builder->GetLeafCount()) { return Status::UnknownError("data type leaf_count != builder_leaf_count", leaf_count, " ", builder->GetLeafCount()); @@ -375,10 +479,10 @@ class FileWriterImpl : public FileWriter { if (row_group_writer_->buffered()) { return Status::Invalid("Cannot write column chunk into the buffered row group."); } - ARROW_ASSIGN_OR_RAISE( - std::unique_ptr writer, - ArrowColumnWriterV2::Make(*data, offset, size, schema_manifest_, - row_group_writer_)); + ARROW_ASSIGN_OR_RAISE(std::unique_ptr writer, + ArrowColumnWriterV2::Make( + *data, offset, size, schema_manifest_, row_group_writer_, + arrow_properties_->write_fixed_size_list_as_vector())); return writer->Write(&column_write_context_); } @@ -459,8 +563,10 @@ class FileWriterImpl : public FileWriter { ChunkedArray chunked_array{batch.column(i)}; ARROW_ASSIGN_OR_RAISE( std::unique_ptr writer, - ArrowColumnWriterV2::Make(chunked_array, offset, size, schema_manifest_, - row_group_writer_, column_index_start)); + ArrowColumnWriterV2::Make( + chunked_array, offset, size, schema_manifest_, row_group_writer_, + arrow_properties_->write_fixed_size_list_as_vector(), + column_index_start)); column_index_start += writer->leaf_count(); if (arrow_properties_->use_threads()) { writers.emplace_back(std::move(writer)); diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index b3ed46ee2d28..df6fb8cba700 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -83,6 +83,27 @@ namespace parquet { namespace { +std::vector<::parquet::internal::Chunk> AlignCdcChunksToVectorBoundaries( + std::vector<::parquet::internal::Chunk> chunks, int32_t vector_length) { + if (vector_length <= 1 || chunks.size() <= 1) { + return chunks; + } + + std::vector<::parquet::internal::Chunk> aligned; + aligned.reserve(chunks.size()); + auto pending = chunks.front(); + for (size_t i = 1; i < chunks.size(); ++i) { + if (pending.levels_to_write % vector_length == 0) { + aligned.push_back(pending); + pending = chunks[i]; + } else { + pending.levels_to_write += chunks[i].levels_to_write; + } + } + aligned.push_back(pending); + return aligned; +} + // Visitor that extracts the value buffer from a FlatArray at a given offset. struct ValueBufferSlicer { template @@ -1185,6 +1206,38 @@ inline void DoInBatchesNonRepeated(int64_t num_levels, int64_t batch_size, } } +// DoInBatches for VECTOR columns without repeated ancestors. num_levels +// counts physical leaf slots and must be a whole number of vectors; batching +// happens on vector boundaries so a vector value is never split across pages. +template +inline void DoInBatchesVectorNonRepeated(int64_t num_levels, int64_t batch_size, + int64_t max_rows_per_page, int32_t vector_length, + Action&& action, + GetBufferedRows&& curr_page_buffered_rows) { + ARROW_DCHECK_GT(vector_length, 0); + if (num_levels % vector_length != 0) { + throw ParquetException("VECTOR columns must be written in whole-vector batches"); + } + const int64_t total_rows = num_levels / vector_length; + int64_t row_offset = 0; + while (row_offset < total_rows) { + int64_t page_buffered_rows = curr_page_buffered_rows(); + ARROW_DCHECK_LE(page_buffered_rows, max_rows_per_page); + + int64_t max_batch_rows = std::max(1, batch_size / vector_length); + max_batch_rows = std::min(max_batch_rows, total_rows - row_offset); + max_batch_rows = std::min(max_batch_rows, max_rows_per_page - page_buffered_rows); + if (max_batch_rows == 0) { + max_batch_rows = 1; + } + int64_t level_offset = row_offset * vector_length; + int64_t level_count = max_batch_rows * vector_length; + + action(level_offset, level_count, /*check_page_limit=*/true); + row_offset += max_batch_rows; + } +} + // DoInBatches for repeated columns template inline void DoInBatchesRepeated(const int16_t* def_levels, const int16_t* rep_levels, @@ -1240,12 +1293,20 @@ inline void DoInBatchesRepeated(const int16_t* def_levels, const int16_t* rep_le template inline void DoInBatches(const int16_t* def_levels, const int16_t* rep_levels, int64_t num_levels, int64_t batch_size, int64_t max_rows_per_page, - bool pages_change_on_record_boundaries, Action&& action, + bool pages_change_on_record_boundaries, bool is_vector, + int32_t vector_length, Action&& action, GetBufferedRows&& curr_page_buffered_rows) { if (!rep_levels) { - DoInBatchesNonRepeated(num_levels, batch_size, max_rows_per_page, - std::forward(action), - std::forward(curr_page_buffered_rows)); + if (is_vector) { + DoInBatchesVectorNonRepeated( + num_levels, batch_size, max_rows_per_page, vector_length, + std::forward(action), + std::forward(curr_page_buffered_rows)); + } else { + DoInBatchesNonRepeated(num_levels, batch_size, max_rows_per_page, + std::forward(action), + std::forward(curr_page_buffered_rows)); + } } else { DoInBatchesRepeated(def_levels, rep_levels, num_levels, batch_size, max_rows_per_page, pages_change_on_record_boundaries, std::forward(action), @@ -1323,7 +1384,11 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, } pages_change_on_record_boundaries_ = properties->data_page_version() == ParquetDataPageVersion::V2 || - properties->page_index_enabled(descr_->path()); + properties->page_index_enabled(descr_->path()) || + // A vector value must not be split across pages; for VECTOR columns + // below repeated ancestors a record contains whole vectors, so + // breaking pages on record boundaries preserves the invariant. + descr_->in_vector_column(); } int64_t Close() override { return ColumnWriterImpl::Close(); } @@ -1366,7 +1431,8 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, }; DoInBatches(def_levels, rep_levels, num_values, properties_->write_batch_size(), properties_->max_rows_per_page(), pages_change_on_record_boundaries(), - WriteChunk, [this]() { return num_buffered_rows_; }); + descr_->in_vector_column(), descr_->effective_vector_length(), WriteChunk, + [this]() { return num_buffered_rows_; }); return value_offset; } @@ -1417,16 +1483,20 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, }; DoInBatches(def_levels, rep_levels, num_values, properties_->write_batch_size(), properties_->max_rows_per_page(), pages_change_on_record_boundaries(), - WriteChunk, [this]() { return num_buffered_rows_; }); + descr_->in_vector_column(), descr_->effective_vector_length(), WriteChunk, + [this]() { return num_buffered_rows_; }); } Status WriteArrow(const int16_t* def_levels, const int16_t* rep_levels, int64_t num_levels, const ::arrow::Array& leaf_array, ArrowWriteContext* ctx, bool leaf_field_nullable) override { BEGIN_PARQUET_CATCH_EXCEPTIONS + const bool is_vector = descr_->in_vector_column(); // Leaf nulls are canonical when there is only a single null element after a list - // and it is at the leaf. + // and it is at the leaf. VECTOR parent nulls may also be materialized as spaced + // null slots in the leaf array, so do not treat those as single nullable elements. bool single_nullable_element = + !is_vector && (level_info_.def_level == level_info_.repeated_ancestor_def_level + 1) && leaf_field_nullable; if (!leaf_field_nullable && leaf_array.null_count() != 0) { @@ -1446,6 +1516,10 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, DCHECK(content_defined_chunker_.has_value()); auto chunks = content_defined_chunker_->GetChunks(def_levels, rep_levels, num_levels, leaf_array); + if (is_vector && descr_->max_repetition_level() == 0) { + chunks = AlignCdcChunksToVectorBoundaries(std::move(chunks), + descr_->effective_vector_length()); + } for (size_t i = 0; i < chunks.size(); i++) { auto chunk = chunks[i]; auto chunk_array = leaf_array.Slice(chunk.value_offset); @@ -1700,6 +1774,13 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, } WriteRepetitionLevels(num_levels, rep_levels); + } else if (descr_->in_vector_column()) { + const int32_t vector_length = descr_->effective_vector_length(); + if (vector_length <= 0 || num_levels % vector_length != 0) { + throw ParquetException("VECTOR columns must be written in whole-vector batches"); + } + rows_written_ += num_levels / vector_length; + num_buffered_rows_ += num_levels / vector_length; } else { // Each value is exactly one row rows_written_ += num_levels; @@ -1794,6 +1875,13 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, } } WriteRepetitionLevels(num_levels, rep_levels); + } else if (descr_->in_vector_column()) { + const int32_t vector_length = descr_->effective_vector_length(); + if (vector_length <= 0 || num_levels % vector_length != 0) { + throw ParquetException("VECTOR columns must be written in whole-vector batches"); + } + rows_written_ += num_levels / vector_length; + num_buffered_rows_ += num_levels / vector_length; } else { // Each value is exactly one row rows_written_ += num_levels; @@ -2079,6 +2167,7 @@ Status TypedColumnWriterImpl::WriteArrowDictionary( PARQUET_CATCH_NOT_OK( DoInBatches(def_levels, rep_levels, num_levels, properties_->write_batch_size(), properties_->max_rows_per_page(), pages_change_on_record_boundaries(), + descr_->in_vector_column(), descr_->effective_vector_length(), WriteIndicesChunk, [this]() { return num_buffered_rows_; })); return Status::OK(); } @@ -2541,6 +2630,7 @@ Status TypedColumnWriterImpl::WriteArrowDense( PARQUET_CATCH_NOT_OK( DoInBatches(def_levels, rep_levels, num_levels, properties_->write_batch_size(), properties_->max_rows_per_page(), pages_change_on_record_boundaries(), + descr_->in_vector_column(), descr_->effective_vector_length(), WriteChunk, [this]() { return num_buffered_rows_; })); return Status::OK(); } diff --git a/cpp/src/parquet/parquet.thrift b/cpp/src/parquet/parquet.thrift index e3cc5adb9648..bd4112c1a504 100644 --- a/cpp/src/parquet/parquet.thrift +++ b/cpp/src/parquet/parquet.thrift @@ -190,6 +190,61 @@ enum FieldRepetitionType { /** The field is repeated and can contain 0 or more values */ REPEATED = 2; + + /** EXPERIMENTAL: The field repeats a fixed number of times per parent value. + * + * A VECTOR field repeats exactly vector_length times for every parent value + * and, unlike REPEATED, does not increase the maximum definition or + * repetition level of its descendants: readers reconstruct vector values + * from the fixed multiplicity declared in the schema instead of decoding + * repetition levels. + * + * Vector fields MUST use the following 3-level structure, mirroring LIST: + * + * group (VECTOR) { + * vector group list [vector_length] { + * element; + * } + * } + * + * - The outer group MUST be REQUIRED or OPTIONAL and carries the + * nullability of the vector value itself. It MUST be annotated with the + * VECTOR logical type and MUST have exactly one child: the + * VECTOR-repeated group. + * - The VECTOR-repeated middle group carries vector_length and MUST have + * exactly one child: the element field. + * - The element field MUST be REQUIRED or OPTIONAL and MAY be a primitive + * type, a group of fields (a struct), or another vector group in the same + * form. For nested vectors a leaf stores the product of the + * vector_length values of its VECTOR ancestors physical values per parent + * record. + * + * Data layout rules: + * - Writers MUST emit exactly vector_length child slots for every parent + * record in which the vector value is present or null; null vector values + * emit vector_length slots carrying the definition level of the null + * vector. Ancestors that make the vector value absent altogether (for + * example an empty or null list) contribute a single level entry with + * their usual definition level and no slots. Column chunk num_values and + * null counts therefore include the padding slots of null vector values. + * - Every slot of an existing vector occupies one position in both level + * streams. The repetition level of the first slot encodes record + * structure as usual; the repetition levels of the remaining + * vector_length - 1 slots MUST be written as the column's maximum + * repetition level and MUST be ignored by readers. Stream consumption is + * definition-driven: after reading a position whose definition level + * indicates an existing vector, a reader consumes that position and the + * following vector_length - 1 positions as one vector value. + * - Writers MUST NOT split the slots of one vector value across data pages; + * pages begin and end on whole-vector boundaries. + * - vector_length MUST be positive: a zero-length vector value would + * contribute no slots at all, leaving row counts and row nullability + * unrepresentable in this layout. Writers MUST represent zero-length + * fixed-size lists with the LIST encoding instead. + * + * Readers that do not understand VECTOR are expected to reject the file. + */ + VECTOR = 3; } /** @@ -462,6 +517,23 @@ struct GeographyType { 2: optional EdgeInterpolationAlgorithm algorithm; } +/** + * Embedded Vector logical type annotation + * + * EXPERIMENTAL: annotates the outer group of a fixed-size vector field; see + * FieldRepetitionType.VECTOR for the full structure and data layout rules. + * The annotated group MUST be REQUIRED or OPTIONAL and MUST have exactly one + * child: a VECTOR-repeated group whose SchemaElement.vector_length carries + * the fixed multiplicity. The annotation is required so that readers can + * distinguish a vector field from a struct-like group that happens to + * contain a VECTOR-repeated child. + * + * Allowed for group nodes only. Defined as a typedef of the (empty) + * NullType struct to keep the generated code small; a format-level proposal + * would define a distinct empty struct. + */ +typedef NullType VectorType + /** * LogicalType annotations to replace ConvertedType. * @@ -495,6 +567,7 @@ union LogicalType { 16: VariantType VARIANT // no compatible ConvertedType 17: GeometryType GEOMETRY // no compatible ConvertedType 18: GeographyType GEOGRAPHY // no compatible ConvertedType + 19: VectorType VECTOR // no compatible ConvertedType } /** @@ -557,6 +630,18 @@ struct SchemaElement { * for some logical types to ensure forward-compatibility in format v1. */ 10: optional LogicalType logicalType + + /** + * EXPERIMENTAL: The fixed number of times the field repeats per parent + * value. + * + * MUST be set, and positive, when repetition_type is VECTOR; MUST NOT be + * set otherwise. Zero-length vectors are not representable as VECTOR and + * use the LIST encoding (see FieldRepetitionType.VECTOR). For nested + * VECTOR fields the number of physical leaf values per parent record is + * the product of vector_length over the leaf's VECTOR ancestors. + */ + 12: optional i32 vector_length; } /** @@ -1090,7 +1175,7 @@ union ColumnOrder { * - If the min is +0, the row group may contain -0 values as well. * - If the max is -0, the row group may contain +0 values as well. * - When looking for NaN values, min and max should be ignored. - * + * * When writing statistics the following rules should be followed: * - NaNs should not be written to min or max statistics fields. * - If the computed max value is zero (whether negative or positive), diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 6634bac4f684..2cf375ccefe4 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1298,6 +1298,7 @@ class PARQUET_EXPORT ArrowWriterProperties { truncated_timestamps_allowed_(false), store_schema_(false), compliant_nested_types_(true), + write_fixed_size_list_as_vector_(false), engine_version_(V2), use_threads_(kArrowDefaultUseThreads), executor_(NULLPTR), @@ -1367,6 +1368,14 @@ class PARQUET_EXPORT ArrowWriterProperties { return this; } + /// \brief EXPERIMENTAL: encode supported Arrow FixedSizeList values as Parquet + /// VECTOR. Unsupported FixedSizeList values, including zero-length lists, continue + /// to use the standard LIST encoding. + Builder* enable_experimental_vector_encoding() { + write_fixed_size_list_as_vector_ = true; + return this; + } + /// Set the version of the Parquet writer engine. Builder* set_engine_version(EngineVersion version) { engine_version_ = version; @@ -1409,7 +1418,8 @@ class PARQUET_EXPORT ArrowWriterProperties { return std::shared_ptr(new ArrowWriterProperties( write_timestamps_as_int96_, coerce_timestamps_enabled_, coerce_timestamps_unit_, truncated_timestamps_allowed_, store_schema_, compliant_nested_types_, - engine_version_, use_threads_, executor_, write_time_adjusted_to_utc_)); + write_fixed_size_list_as_vector_, engine_version_, use_threads_, executor_, + write_time_adjusted_to_utc_)); } private: @@ -1421,6 +1431,7 @@ class PARQUET_EXPORT ArrowWriterProperties { bool store_schema_; bool compliant_nested_types_; + bool write_fixed_size_list_as_vector_; EngineVersion engine_version_; bool use_threads_; @@ -1447,6 +1458,10 @@ class PARQUET_EXPORT ArrowWriterProperties { /// "element". bool compliant_nested_types() const { return compliant_nested_types_; } + bool write_fixed_size_list_as_vector() const { + return write_fixed_size_list_as_vector_; + } + /// \brief The underlying engine version to use when writing Arrow data. /// /// V2 is currently the latest V1 is considered deprecated but left in @@ -1471,6 +1486,7 @@ class PARQUET_EXPORT ArrowWriterProperties { ::arrow::TimeUnit::type coerce_timestamps_unit, bool truncated_timestamps_allowed, bool store_schema, bool compliant_nested_types, + bool write_fixed_size_list_as_vector, EngineVersion engine_version, bool use_threads, ::arrow::internal::Executor* executor, bool write_time_adjusted_to_utc) @@ -1480,6 +1496,7 @@ class PARQUET_EXPORT ArrowWriterProperties { truncated_timestamps_allowed_(truncated_timestamps_allowed), store_schema_(store_schema), compliant_nested_types_(compliant_nested_types), + write_fixed_size_list_as_vector_(write_fixed_size_list_as_vector), engine_version_(engine_version), use_threads_(use_threads), executor_(executor), @@ -1491,6 +1508,7 @@ class PARQUET_EXPORT ArrowWriterProperties { const bool truncated_timestamps_allowed_; const bool store_schema_; const bool compliant_nested_types_; + const bool write_fixed_size_list_as_vector_; const EngineVersion engine_version_; const bool use_threads_; ::arrow::internal::Executor* executor_; diff --git a/cpp/src/parquet/schema.cc b/cpp/src/parquet/schema.cc index 0cfa49c21c16..04f829e6e50e 100644 --- a/cpp/src/parquet/schema.cc +++ b/cpp/src/parquet/schema.cc @@ -50,6 +50,16 @@ void CheckColumnBounds(int column_index, size_t max_columns) { } } +void ValidateVectorProperties(Repetition::type repetition, int32_t vector_length) { + if (repetition == Repetition::VECTOR) { + if (vector_length <= 0) { + throw ParquetException("VECTOR nodes must specify a positive vector_length"); + } + } else if (vector_length != -1) { + throw ParquetException("Only VECTOR nodes may specify vector_length"); + } +} + } // namespace namespace schema { @@ -117,7 +127,7 @@ const std::shared_ptr Node::path() const { bool Node::EqualsInternal(const Node* other) const { return type_ == other->type_ && name_ == other->name_ && repetition_ == other->repetition_ && converted_type_ == other->converted_type_ && - field_id_ == other->field_id() && + field_id_ == other->field_id() && vector_length_ == other->vector_length() && logical_type_->Equals(*(other->logical_type())); } @@ -128,10 +138,12 @@ void Node::SetParent(const Node* parent) { parent_ = parent; } PrimitiveNode::PrimitiveNode(const std::string& name, Repetition::type repetition, Type::type type, ConvertedType::type converted_type, - int length, int precision, int scale, int id) - : Node(Node::PRIMITIVE, name, repetition, converted_type, id), + int length, int precision, int scale, int id, + int32_t vector_length) + : Node(Node::PRIMITIVE, name, repetition, converted_type, id, vector_length), physical_type_(type), type_length_(length) { + ValidateVectorProperties(repetition, vector_length); std::stringstream ss; // PARQUET-842: In an earlier revision, decimal_metadata_.isset was being @@ -241,10 +253,12 @@ PrimitiveNode::PrimitiveNode(const std::string& name, Repetition::type repetitio PrimitiveNode::PrimitiveNode(const std::string& name, Repetition::type repetition, std::shared_ptr logical_type, - Type::type physical_type, int physical_length, int id) - : Node(Node::PRIMITIVE, name, repetition, std::move(logical_type), id), + Type::type physical_type, int physical_length, int id, + int32_t vector_length) + : Node(Node::PRIMITIVE, name, repetition, std::move(logical_type), id, vector_length), physical_type_(physical_type), type_length_(physical_length) { + ValidateVectorProperties(repetition, vector_length); std::stringstream error; if (logical_type_) { // Check for logical type <=> node type consistency @@ -315,8 +329,11 @@ void PrimitiveNode::VisitConst(Node::ConstVisitor* visitor) const { // Group node GroupNode::GroupNode(const std::string& name, Repetition::type repetition, - const NodeVector& fields, ConvertedType::type converted_type, int id) - : Node(Node::GROUP, name, repetition, converted_type, id), fields_(fields) { + const NodeVector& fields, ConvertedType::type converted_type, int id, + int32_t vector_length) + : Node(Node::GROUP, name, repetition, converted_type, id, vector_length), + fields_(fields) { + ValidateVectorProperties(repetition, vector_length); // For forward compatibility, create an equivalent logical type logical_type_ = LogicalType::FromConvertedType(converted_type_); if (!(logical_type_ && (logical_type_->is_nested() || logical_type_->is_none()) && @@ -334,8 +351,11 @@ GroupNode::GroupNode(const std::string& name, Repetition::type repetition, GroupNode::GroupNode(const std::string& name, Repetition::type repetition, const NodeVector& fields, - std::shared_ptr logical_type, int id) - : Node(Node::GROUP, name, repetition, std::move(logical_type), id), fields_(fields) { + std::shared_ptr logical_type, int id, + int32_t vector_length) + : Node(Node::GROUP, name, repetition, std::move(logical_type), id, vector_length), + fields_(fields) { + ValidateVectorProperties(repetition, vector_length); if (logical_type_) { // Check for logical type <=> node type consistency if (logical_type_->is_nested()) { @@ -423,19 +443,23 @@ std::unique_ptr GroupNode::FromParquet(const void* opaque_element, if (element->__isset.field_id) { field_id = element->field_id; } + int32_t vector_length = -1; + if (element->__isset.vector_length) { + vector_length = element->vector_length; + } std::unique_ptr group_node; if (element->__isset.logicalType) { // updated writer with logical type present - group_node = std::unique_ptr( - new GroupNode(element->name, LoadEnumSafe(&element->repetition_type), fields, - LogicalType::FromThrift(element->logicalType), field_id)); + group_node = std::unique_ptr(new GroupNode( + element->name, LoadEnumSafe(&element->repetition_type), fields, + LogicalType::FromThrift(element->logicalType), field_id, vector_length)); } else { group_node = std::unique_ptr(new GroupNode( element->name, LoadEnumSafe(&element->repetition_type), fields, (element->__isset.converted_type ? LoadEnumSafe(&element->converted_type) : ConvertedType::NONE), - field_id)); + field_id, vector_length)); } return std::unique_ptr(group_node.release()); @@ -449,25 +473,30 @@ std::unique_ptr PrimitiveNode::FromParquet(const void* opaque_element) { if (element->__isset.field_id) { field_id = element->field_id; } + int32_t vector_length = -1; + if (element->__isset.vector_length) { + vector_length = element->vector_length; + } std::unique_ptr primitive_node; if (element->__isset.logicalType) { // updated writer with logical type present - primitive_node = std::unique_ptr( - new PrimitiveNode(element->name, LoadEnumSafe(&element->repetition_type), - LogicalType::FromThrift(element->logicalType), - LoadEnumSafe(&element->type), element->type_length, field_id)); - } else if (element->__isset.converted_type) { - // legacy writer with converted type present primitive_node = std::unique_ptr(new PrimitiveNode( element->name, LoadEnumSafe(&element->repetition_type), - LoadEnumSafe(&element->type), LoadEnumSafe(&element->converted_type), - element->type_length, element->precision, element->scale, field_id)); + LogicalType::FromThrift(element->logicalType), LoadEnumSafe(&element->type), + element->type_length, field_id, vector_length)); + } else if (element->__isset.converted_type) { + // legacy writer with converted type present + primitive_node = std::unique_ptr( + new PrimitiveNode(element->name, LoadEnumSafe(&element->repetition_type), + LoadEnumSafe(&element->type), + LoadEnumSafe(&element->converted_type), element->type_length, + element->precision, element->scale, field_id, vector_length)); } else { // logical type not present primitive_node = std::unique_ptr(new PrimitiveNode( element->name, LoadEnumSafe(&element->repetition_type), NoLogicalType::Make(), - LoadEnumSafe(&element->type), element->type_length, field_id)); + LoadEnumSafe(&element->type), element->type_length, field_id, vector_length)); } // Return as unique_ptr to the base type @@ -499,6 +528,9 @@ void GroupNode::ToParquet(void* opaque_element) const { if (field_id_ >= 0) { element->__set_field_id(field_id_); } + if (is_vector()) { + element->__set_vector_length(vector_length_); + } if (logical_type_ && logical_type_->is_serialized()) { element->__set_logicalType(logical_type_->ToThrift()); } @@ -524,6 +556,9 @@ void PrimitiveNode::ToParquet(void* opaque_element) const { if (field_id_ >= 0) { element->__set_field_id(field_id_); } + if (is_vector()) { + element->__set_vector_length(vector_length_); + } if (logical_type_ && logical_type_->is_serialized() && // TODO(tpboudreau): remove the following conjunct to enable serialization // of IntervalTypes after parquet.thrift recognizes them @@ -634,6 +669,9 @@ static void PrintRepLevel(Repetition::type repetition, std::ostream& stream) { case Repetition::REPEATED: stream << "repeated"; break; + case Repetition::VECTOR: + stream << "vector"; + break; default: break; } @@ -710,14 +748,19 @@ struct SchemaPrinter : public Node::ConstVisitor { stream_ << " "; PrintType(node, stream_); stream_ << " field_id=" << node->field_id() << " " << node->name(); + if (node->is_vector()) { + stream_ << " [" << node->vector_length() << "]"; + } PrintConvertedType(node, stream_); stream_ << ";" << std::endl; } void Visit(const GroupNode* node) { PrintRepLevel(node->repetition(), stream_); - stream_ << " group " - << "field_id=" << node->field_id() << " " << node->name(); + stream_ << " group " << "field_id=" << node->field_id() << " " << node->name(); + if (node->is_vector()) { + stream_ << " [" << node->vector_length() << "]"; + } auto lt = node->converted_type(); const auto& la = node->logical_type(); if (la && la->is_valid() && !la->is_none()) { @@ -836,6 +879,9 @@ void SchemaDescriptor::BuildTree(const NodePtr& node, int16_t max_def_level, // between an empty list and a list with an item in it. ++max_rep_level; ++max_def_level; + } else if (node->is_vector()) { + // VECTOR fields repeat a fixed number of times per parent value without + // increasing the maximum definition or repetition level. } // Now, walk the schema and create a ColumnDescriptor for each leaf node @@ -848,8 +894,23 @@ void SchemaDescriptor::BuildTree(const NodePtr& node, int16_t max_def_level, node_to_leaf_index_[static_cast(node.get())] = static_cast(leaves_.size()); + // Determine the product of all VECTOR-repeated ancestors (including this leaf) + // so nested VECTOR shapes such as vector<3, vector<4, int32>> report the total + // number of physical leaf values contributed per parent record. + int32_t effective_vector_length = -1; + for (const Node* cursor = node.get(); cursor != nullptr; cursor = cursor->parent()) { + if (cursor->is_vector()) { + if (effective_vector_length < 0) { + effective_vector_length = cursor->vector_length(); + } else { + effective_vector_length *= cursor->vector_length(); + } + } + } + // Primitive node, append to leaves - leaves_.push_back(ColumnDescriptor(node, max_def_level, max_rep_level, this)); + leaves_.push_back(ColumnDescriptor(node, max_def_level, max_rep_level, this, + effective_vector_length)); leaf_to_base_.emplace(static_cast(leaves_.size()) - 1, base); leaf_to_idx_.emplace(node->path()->ToDotString(), static_cast(leaves_.size()) - 1); @@ -866,10 +927,12 @@ int SchemaDescriptor::GetColumnIndex(const PrimitiveNode& node) const { ColumnDescriptor::ColumnDescriptor(schema::NodePtr node, int16_t max_definition_level, int16_t max_repetition_level, - const SchemaDescriptor* schema_descr) + const SchemaDescriptor* schema_descr, + int32_t effective_vector_length) : node_(std::move(node)), max_definition_level_(max_definition_level), - max_repetition_level_(max_repetition_level) { + max_repetition_level_(max_repetition_level), + effective_vector_length_(effective_vector_length) { if (!node_->is_primitive()) { throw ParquetException("Must be a primitive type"); } @@ -879,7 +942,8 @@ ColumnDescriptor::ColumnDescriptor(schema::NodePtr node, int16_t max_definition_ bool ColumnDescriptor::Equals(const ColumnDescriptor& other) const { return primitive_node_->Equals(other.primitive_node_) && max_repetition_level() == other.max_repetition_level() && - max_definition_level() == other.max_definition_level(); + max_definition_level() == other.max_definition_level() && + effective_vector_length() == other.effective_vector_length(); } const ColumnDescriptor* SchemaDescriptor::Column(int i) const { @@ -930,8 +994,13 @@ std::string ColumnDescriptor::ToString() const { << " physical_type: " << TypeToString(physical_type()) << "," << std::endl << " converted_type: " << ConvertedTypeToString(converted_type()) << "," << std::endl - << " logical_type: " << logical_type()->ToString() << "," << std::endl - << " max_definition_level: " << max_definition_level() << "," << std::endl + << " logical_type: " << logical_type()->ToString() << "," << std::endl; + + if (schema_node()->is_vector()) { + ss << " vector_length: " << schema_node()->vector_length() << "," << std::endl; + } + + ss << " max_definition_level: " << max_definition_level() << "," << std::endl << " max_repetition_level: " << max_repetition_level() << "," << std::endl; if (physical_type() == ::parquet::Type::FIXED_LEN_BYTE_ARRAY) { diff --git a/cpp/src/parquet/schema.h b/cpp/src/parquet/schema.h index 1addc73bd367..3b21cb6a40dd 100644 --- a/cpp/src/parquet/schema.h +++ b/cpp/src/parquet/schema.h @@ -109,6 +109,11 @@ class PARQUET_EXPORT Node { bool is_repeated() const { return repetition_ == Repetition::REPEATED; } + /// \brief True when this node specifically has VECTOR repetition; use + /// ColumnDescriptor::in_vector_column() to ask whether a column is part of + /// a VECTOR subtree. + bool is_vector() const { return repetition_ == Repetition::VECTOR; } + bool is_required() const { return repetition_ == Repetition::REQUIRED; } virtual bool Equals(const Node* other) const = 0; @@ -128,6 +133,10 @@ class PARQUET_EXPORT Node { /// Thrift. int field_id() const { return field_id_; } + /// \brief The fixed number of values per parent when repetition == VECTOR. + /// Returns -1 when this node is not VECTOR-repeated. + int32_t vector_length() const { return vector_length_; } + const Node* parent() const { return parent_; } const std::shared_ptr path() const; @@ -155,21 +164,25 @@ class PARQUET_EXPORT Node { friend class GroupNode; Node(Node::type type, const std::string& name, Repetition::type repetition, - ConvertedType::type converted_type = ConvertedType::NONE, int field_id = -1) + ConvertedType::type converted_type = ConvertedType::NONE, int field_id = -1, + int32_t vector_length = -1) : type_(type), name_(name), repetition_(repetition), converted_type_(converted_type), field_id_(field_id), + vector_length_(vector_length), parent_(NULLPTR) {} Node(Node::type type, const std::string& name, Repetition::type repetition, - std::shared_ptr logical_type, int field_id = -1) + std::shared_ptr logical_type, int field_id = -1, + int32_t vector_length = -1) : type_(type), name_(name), repetition_(repetition), logical_type_(std::move(logical_type)), field_id_(field_id), + vector_length_(vector_length), parent_(NULLPTR) {} Node::type type_; @@ -178,6 +191,7 @@ class PARQUET_EXPORT Node { ConvertedType::type converted_type_{ConvertedType::NONE}; std::shared_ptr logical_type_; int field_id_; + int32_t vector_length_; // Nodes should not be shared, they have a single parent. const Node* parent_; @@ -205,9 +219,9 @@ class PARQUET_EXPORT PrimitiveNode : public Node { Type::type type, ConvertedType::type converted_type = ConvertedType::NONE, int length = -1, int precision = -1, int scale = -1, - int field_id = -1) { + int field_id = -1, int32_t vector_length = -1) { return NodePtr(new PrimitiveNode(name, repetition, type, converted_type, length, - precision, scale, field_id)); + precision, scale, field_id, vector_length)); } // If no logical type, pass LogicalType::None() or nullptr @@ -215,9 +229,10 @@ class PARQUET_EXPORT PrimitiveNode : public Node { static inline NodePtr Make(const std::string& name, Repetition::type repetition, std::shared_ptr logical_type, Type::type primitive_type, int primitive_length = -1, - int field_id = -1) { + int field_id = -1, int32_t vector_length = -1) { return NodePtr(new PrimitiveNode(name, repetition, std::move(logical_type), - primitive_type, primitive_length, field_id)); + primitive_type, primitive_length, field_id, + vector_length)); } bool Equals(const Node* other) const override; @@ -239,11 +254,13 @@ class PARQUET_EXPORT PrimitiveNode : public Node { private: PrimitiveNode(const std::string& name, Repetition::type repetition, Type::type type, ConvertedType::type converted_type = ConvertedType::NONE, int length = -1, - int precision = -1, int scale = -1, int field_id = -1); + int precision = -1, int scale = -1, int field_id = -1, + int32_t vector_length = -1); PrimitiveNode(const std::string& name, Repetition::type repetition, std::shared_ptr logical_type, - Type::type primitive_type, int primitive_length = -1, int field_id = -1); + Type::type primitive_type, int primitive_length = -1, int field_id = -1, + int32_t vector_length = -1); Type::type physical_type_; int32_t type_length_; @@ -270,8 +287,9 @@ class PARQUET_EXPORT GroupNode : public Node { static inline NodePtr Make(const std::string& name, Repetition::type repetition, const NodeVector& fields, ConvertedType::type converted_type = ConvertedType::NONE, - int field_id = -1) { - return NodePtr(new GroupNode(name, repetition, fields, converted_type, field_id)); + int field_id = -1, int32_t vector_length = -1) { + return NodePtr( + new GroupNode(name, repetition, fields, converted_type, field_id, vector_length)); } // If no logical type, pass nullptr @@ -279,9 +297,9 @@ class PARQUET_EXPORT GroupNode : public Node { static inline NodePtr Make(const std::string& name, Repetition::type repetition, const NodeVector& fields, std::shared_ptr logical_type, - int field_id = -1) { - return NodePtr( - new GroupNode(name, repetition, fields, std::move(logical_type), field_id)); + int field_id = -1, int32_t vector_length = -1) { + return NodePtr(new GroupNode(name, repetition, fields, std::move(logical_type), + field_id, vector_length)); } bool Equals(const Node* other) const override; @@ -307,11 +325,12 @@ class PARQUET_EXPORT GroupNode : public Node { private: GroupNode(const std::string& name, Repetition::type repetition, const NodeVector& fields, - ConvertedType::type converted_type = ConvertedType::NONE, int field_id = -1); + ConvertedType::type converted_type = ConvertedType::NONE, int field_id = -1, + int32_t vector_length = -1); GroupNode(const std::string& name, Repetition::type repetition, const NodeVector& fields, std::shared_ptr logical_type, - int field_id = -1); + int field_id = -1, int32_t vector_length = -1); NodeVector fields_; bool EqualsInternal(const GroupNode* other) const; @@ -358,7 +377,8 @@ class PARQUET_EXPORT ColumnDescriptor { public: ColumnDescriptor(schema::NodePtr node, int16_t max_definition_level, int16_t max_repetition_level, - const SchemaDescriptor* schema_descr = NULLPTR); + const SchemaDescriptor* schema_descr = NULLPTR, + int32_t effective_vector_length = -1); bool Equals(const ColumnDescriptor& other) const; @@ -366,6 +386,16 @@ class PARQUET_EXPORT ColumnDescriptor { int16_t max_repetition_level() const { return max_repetition_level_; } + /// \brief The fixed number of leaf values per parent record contributed by + /// this column, computed as the product of all VECTOR-repeated ancestors + /// (including the leaf itself if it is VECTOR-repeated). Returns -1 when the + /// column is not part of a VECTOR-repeated subtree. + int32_t effective_vector_length() const { return effective_vector_length_; } + + /// \brief True when this column belongs to a VECTOR-repeated subtree + /// (the leaf or some ancestor has Repetition::VECTOR). + bool in_vector_column() const { return effective_vector_length_ > 0; } + Type::type physical_type() const { return primitive_node_->physical_type(); } ConvertedType::type converted_type() const { return primitive_node_->converted_type(); } @@ -402,6 +432,7 @@ class PARQUET_EXPORT ColumnDescriptor { int16_t max_definition_level_; int16_t max_repetition_level_; + int32_t effective_vector_length_; }; // Container for the converted Parquet schema with a computed information from diff --git a/cpp/src/parquet/schema_test.cc b/cpp/src/parquet/schema_test.cc index 2950a7df70f8..35c1f8bd9332 100644 --- a/cpp/src/parquet/schema_test.cc +++ b/cpp/src/parquet/schema_test.cc @@ -43,7 +43,8 @@ namespace schema { static inline SchemaElement NewPrimitive(const std::string& name, FieldRepetitionType::type repetition, - Type::type type, int field_id = -1) { + Type::type type, int field_id = -1, + int32_t vector_length = -1) { SchemaElement result; result.__set_name(name); result.__set_repetition_type(repetition); @@ -51,12 +52,16 @@ static inline SchemaElement NewPrimitive(const std::string& name, if (field_id >= 0) { result.__set_field_id(field_id); } + if (vector_length >= 0) { + result.__set_vector_length(vector_length); + } return result; } static inline SchemaElement NewGroup(const std::string& name, FieldRepetitionType::type repetition, - int num_children, int field_id = -1) { + int num_children, int field_id = -1, + int32_t vector_length = -1) { SchemaElement result; result.__set_name(name); result.__set_repetition_type(repetition); @@ -65,6 +70,9 @@ static inline SchemaElement NewGroup(const std::string& name, if (field_id >= 0) { result.__set_field_id(field_id); } + if (vector_length >= 0) { + result.__set_vector_length(vector_length); + } return result; } @@ -217,6 +225,30 @@ TEST_F(TestPrimitiveNode, FromParquet) { ASSERT_EQ(12, prim_node_->decimal_metadata().precision); } +TEST_F(TestPrimitiveNode, VectorFromParquet) { + SchemaElement elt = + NewPrimitive(name_, FieldRepetitionType::VECTOR, Type::FLOAT, field_id_, 8); + + ASSERT_NO_FATAL_FAILURE(Convert(&elt)); + ASSERT_EQ(name_, prim_node_->name()); + ASSERT_EQ(field_id_, prim_node_->field_id()); + ASSERT_EQ(Repetition::VECTOR, prim_node_->repetition()); + ASSERT_TRUE(prim_node_->is_vector()); + ASSERT_EQ(8, prim_node_->vector_length()); + ASSERT_EQ(Type::FLOAT, prim_node_->physical_type()); +} + +TEST_F(TestPrimitiveNode, VectorValidation) { + ASSERT_THROW(PrimitiveNode::Make("vec", Repetition::VECTOR, Type::FLOAT), + ParquetException); + ASSERT_THROW(PrimitiveNode::Make("scalar", Repetition::REQUIRED, Type::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, 4), + ParquetException); + ASSERT_THROW(PrimitiveNode::Make("empty", Repetition::VECTOR, Type::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, 0), + ParquetException); +} + TEST_F(TestPrimitiveNode, Equals) { PrimitiveNode node1("foo", Repetition::REQUIRED, Type::INT32); PrimitiveNode node2("foo", Repetition::REQUIRED, Type::INT64); @@ -365,6 +397,47 @@ TEST_F(TestGroupNode, Attrs) { ASSERT_EQ(ConvertedType::LIST, node2.converted_type()); } +TEST_F(TestGroupNode, VectorAttrs) { + auto node = GroupNode::Make("vec", Repetition::VECTOR, Fields1(), ConvertedType::NONE, + /*field_id=*/-1, + /*vector_length=*/4); + + ASSERT_TRUE(node->is_vector()); + ASSERT_EQ(Repetition::VECTOR, node->repetition()); + ASSERT_EQ(4, node->vector_length()); +} + +TEST_F(TestGroupNode, VectorLogicalTypeRoundTrip) { + // A canonical 3-level vector group annotated with the Vector logical type + // round-trips through the Thrift intermediary. + auto element = PrimitiveNode::Make("element", Repetition::REQUIRED, Type::FLOAT); + auto list = GroupNode::Make("list", Repetition::VECTOR, {element}, + /*logical_type=*/nullptr, /*field_id=*/-1, + /*vector_length=*/3); + auto group = + GroupNode::Make("embedding", Repetition::OPTIONAL, {list}, LogicalType::Vector()); + + ASSERT_TRUE(group->logical_type()->is_vector()); + ASSERT_TRUE(group->logical_type()->is_nested()); + ASSERT_EQ(ConvertedType::NONE, group->converted_type()); + + format::SchemaElement group_element; + group->ToParquet(&group_element); + ASSERT_TRUE(group_element.__isset.logicalType); + ASSERT_TRUE(group_element.logicalType.__isset.VECTOR); + + format::SchemaElement list_element; + list->ToParquet(&list_element); + format::SchemaElement element_element; + element->ToParquet(&element_element); + std::unique_ptr roundtripped = GroupNode::FromParquet( + &group_element, + {GroupNode::FromParquet(&list_element, + {PrimitiveNode::FromParquet(&element_element)})}); + ASSERT_TRUE(roundtripped->logical_type()->is_vector()); + ASSERT_TRUE(group->Equals(roundtripped.get())); +} + TEST_F(TestGroupNode, Equals) { NodeVector f1 = Fields1(); NodeVector f2 = Fields1(); @@ -799,6 +872,39 @@ TEST_F(TestSchemaDescriptor, BuildTree) { ASSERT_EQ(nleaves, descr_.num_columns()); } +TEST_F(TestSchemaDescriptor, BuildTreeVector) { + NodePtr element = PrimitiveNode::Make("element", Repetition::VECTOR, Type::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, 3); + NodePtr embedding = GroupNode::Make("embedding", Repetition::OPTIONAL, {element}); + + descr_.Init(GroupNode::Make("schema", Repetition::REPEATED, {embedding})); + + ASSERT_EQ(1, descr_.num_columns()); + const ColumnDescriptor* col = descr_.Column(0); + EXPECT_EQ(1, col->max_definition_level()); + EXPECT_EQ(0, col->max_repetition_level()); + EXPECT_TRUE(col->schema_node()->is_vector()); + EXPECT_EQ(3, col->schema_node()->vector_length()); + EXPECT_EQ(3, col->effective_vector_length()); +} + +TEST_F(TestSchemaDescriptor, BuildTreeNestedVector) { + NodePtr element = PrimitiveNode::Make("element", Repetition::VECTOR, Type::FLOAT, + ConvertedType::NONE, -1, -1, -1, -1, 4); + NodePtr inner = GroupNode::Make("inner", Repetition::VECTOR, {element}, + /*converted_type=*/ConvertedType::NONE, + /*field_id=*/-1, /*vector_length=*/3); + NodePtr embedding = GroupNode::Make("embedding", Repetition::OPTIONAL, {inner}); + + descr_.Init(GroupNode::Make("schema", Repetition::REPEATED, {embedding})); + + ASSERT_EQ(1, descr_.num_columns()); + const ColumnDescriptor* col = descr_.Column(0); + EXPECT_EQ(1, col->max_definition_level()); + EXPECT_EQ(0, col->max_repetition_level()); + EXPECT_EQ(12, col->effective_vector_length()); +} + TEST_F(TestSchemaDescriptor, HasRepeatedFields) { NodeVector fields; NodePtr schema; @@ -886,6 +992,28 @@ TEST(TestSchemaPrinter, Examples) { ASSERT_EQ(expected, result); } +TEST(TestSchemaPrinter, Vector) { + NodePtr schema = GroupNode::Make( + "schema", Repetition::REPEATED, + {GroupNode::Make( + "embedding", Repetition::OPTIONAL, + {GroupNode::Make( + "vector", Repetition::VECTOR, + {PrimitiveNode::Make("element", Repetition::REQUIRED, Type::FLOAT)}, + /*logical_type=*/nullptr, -1, 3)}, + LogicalType::Vector())}); + + std::string expected = R"(repeated group field_id=-1 schema { + optional group field_id=-1 embedding (Vector) { + vector group field_id=-1 vector [3] { + required float field_id=-1 element; + } + } +} +)"; + ASSERT_EQ(expected, Print(schema)); +} + static void ConfirmFactoryEquivalence( ConvertedType::type converted_type, const std::shared_ptr& from_make, diff --git a/cpp/src/parquet/types.cc b/cpp/src/parquet/types.cc index fb4eb92a7544..1c3277172d22 100644 --- a/cpp/src/parquet/types.cc +++ b/cpp/src/parquet/types.cc @@ -602,6 +602,8 @@ std::shared_ptr LogicalType::FromThrift( } return VariantLogicalType::Make(spec_version); + } else if (type.__isset.VECTOR) { + return VectorLogicalType::Make(); } else { // Sentinel type for one we do not recognize return UndefinedLogicalType::Make(); @@ -673,6 +675,10 @@ std::shared_ptr LogicalType::Variant(int8_t spec_version) { return VariantLogicalType::Make(spec_version); } +std::shared_ptr LogicalType::Vector() { + return VectorLogicalType::Make(); +} + std::shared_ptr LogicalType::None() { return NoLogicalType::Make(); } /* @@ -758,6 +764,7 @@ class LogicalType::Impl { class Geometry; class Geography; class Variant; + class Vector; class No; class Undefined; @@ -839,6 +846,7 @@ bool LogicalType::is_geography() const { bool LogicalType::is_variant() const { return impl_->type() == LogicalType::Type::VARIANT; } +bool LogicalType::is_vector() const { return impl_->type() == LogicalType::Type::VECTOR; } bool LogicalType::is_none() const { return impl_->type() == LogicalType::Type::NONE; } bool LogicalType::is_valid() const { return impl_->type() != LogicalType::Type::UNDEFINED; @@ -847,7 +855,8 @@ bool LogicalType::is_invalid() const { return !is_valid(); } bool LogicalType::is_nested() const { return impl_->type() == LogicalType::Type::LIST || impl_->type() == LogicalType::Type::MAP || - impl_->type() == LogicalType::Type::VARIANT; + impl_->type() == LogicalType::Type::VARIANT || + impl_->type() == LogicalType::Type::VECTOR; } bool LogicalType::is_nonnested() const { return !is_nested(); } bool LogicalType::is_serialized() const { return impl_->is_serialized(); } @@ -2016,6 +2025,20 @@ std::shared_ptr VariantLogicalType::Make(const int8_t spec_ve return logical_type; } +class LogicalType::Impl::Vector final : public LogicalType::Impl::Incompatible, + public LogicalType::Impl::Inapplicable { + public: + friend class VectorLogicalType; + + OVERRIDE_TOSTRING(Vector) + OVERRIDE_TOTHRIFT(VectorType, VECTOR) + + private: + Vector() : LogicalType::Impl(LogicalType::Type::VECTOR, SortOrder::UNKNOWN) {} +}; + +GENERATE_MAKE(Vector) + class LogicalType::Impl::No final : public LogicalType::Impl::SimpleCompatible, public LogicalType::Impl::UniversalApplicable { public: diff --git a/cpp/src/parquet/types.h b/cpp/src/parquet/types.h index ad4df5119e75..35643a27f172 100644 --- a/cpp/src/parquet/types.h +++ b/cpp/src/parquet/types.h @@ -112,7 +112,13 @@ class LogicalType; // Mirrors parquet::FieldRepetitionType struct Repetition { - enum type { REQUIRED = 0, OPTIONAL = 1, REPEATED = 2, /*Always last*/ UNDEFINED = 3 }; + enum type { + REQUIRED = 0, + OPTIONAL = 1, + REPEATED = 2, + VECTOR = 3, + /*Always last*/ UNDEFINED = 4 + }; }; // Reference: @@ -162,6 +168,7 @@ class PARQUET_EXPORT LogicalType { GEOMETRY, GEOGRAPHY, VARIANT, + VECTOR, NONE // Not a real logical type; should always be last element }; }; @@ -231,6 +238,10 @@ class PARQUET_EXPORT LogicalType { static std::shared_ptr Variant( int8_t specVersion = kVariantSpecVersion); + /// \brief Create an experimental Vector logical type, annotating a group + /// node whose single child has VECTOR repetition. + static std::shared_ptr Vector(); + static std::shared_ptr Geometry(std::string crs = ""); static std::shared_ptr Geography( @@ -293,6 +304,7 @@ class PARQUET_EXPORT LogicalType { bool is_geometry() const; bool is_geography() const; bool is_variant() const; + bool is_vector() const; bool is_none() const; /// \brief Return true if this logical type is of a known type. bool is_valid() const; @@ -342,6 +354,16 @@ class PARQUET_EXPORT ListLogicalType : public LogicalType { ListLogicalType() = default; }; +/// \brief Allowed for group nodes only. EXPERIMENTAL: annotates a group whose +/// single child has VECTOR repetition. +class PARQUET_EXPORT VectorLogicalType : public LogicalType { + public: + static std::shared_ptr Make(); + + private: + VectorLogicalType() = default; +}; + /// \brief Allowed for physical type BYTE_ARRAY, must be encoded as UTF-8. class PARQUET_EXPORT EnumLogicalType : public LogicalType { public: diff --git a/python/pyarrow/_dataset_parquet.pyx b/python/pyarrow/_dataset_parquet.pyx index 534f7790923a..6567b95a9ab1 100644 --- a/python/pyarrow/_dataset_parquet.pyx +++ b/python/pyarrow/_dataset_parquet.pyx @@ -620,6 +620,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): "coerce_timestamps", "allow_truncated_timestamps", "use_compliant_nested_type", + "experimental_vector_encoding", } setters = set() @@ -676,6 +677,11 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): writer_engine_version="V2", use_compliant_nested_type=( self._properties["use_compliant_nested_type"] + ), + store_schema=True, + write_time_adjusted_to_utc=False, + experimental_vector_encoding=( + self._properties["experimental_vector_encoding"] ) ) @@ -705,6 +711,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): coerce_timestamps=None, allow_truncated_timestamps=False, use_compliant_nested_type=True, + experimental_vector_encoding=False, encryption_properties=None, write_batch_size=None, dictionary_pagesize_limit=None, diff --git a/python/pyarrow/_parquet.pxd b/python/pyarrow/_parquet.pxd index 36fc2ccf2f33..fa0a06e53c7e 100644 --- a/python/pyarrow/_parquet.pxd +++ b/python/pyarrow/_parquet.pxd @@ -69,6 +69,7 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties( use_compliant_nested_type=*, store_schema=*, write_time_adjusted_to_utc=*, + experimental_vector_encoding=*, ) except * diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 2358a961ebd9..39cb5dac42f0 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -2281,7 +2281,8 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties( writer_engine_version=None, use_compliant_nested_type=True, store_schema=True, - write_time_adjusted_to_utc=False) except *: + write_time_adjusted_to_utc=False, + experimental_vector_encoding=False) except *: """Arrow writer properties""" cdef: shared_ptr[ArrowWriterProperties] arrow_properties @@ -2322,6 +2323,11 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties( else: arrow_props.disable_compliant_nested_types() + # experimental_vector_encoding + + if experimental_vector_encoding: + arrow_props.enable_experimental_vector_encoding() + # writer_engine_version if writer_engine_version == "V1": @@ -2396,7 +2402,8 @@ cdef class ParquetWriter(_Weakrefable): store_decimal_as_integer=False, use_content_defined_chunking=False, write_time_adjusted_to_utc=False, - bloom_filter_options=None): + bloom_filter_options=None, + experimental_vector_encoding=False): cdef: shared_ptr[WriterProperties] properties shared_ptr[ArrowWriterProperties] arrow_properties @@ -2443,6 +2450,7 @@ cdef class ParquetWriter(_Weakrefable): use_compliant_nested_type=use_compliant_nested_type, store_schema=store_schema, write_time_adjusted_to_utc=write_time_adjusted_to_utc, + experimental_vector_encoding=experimental_vector_encoding, ) pool = maybe_unbox_memory_pool(memory_pool) diff --git a/python/pyarrow/includes/libparquet.pxd b/python/pyarrow/includes/libparquet.pxd index a834bd5dfa0c..b1f7e0240607 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -528,6 +528,7 @@ cdef extern from "parquet/api/writer.h" namespace "parquet" nogil: Builder* store_schema() Builder* enable_compliant_nested_types() Builder* disable_compliant_nested_types() + Builder* enable_experimental_vector_encoding() Builder* set_engine_version(ArrowWriterEngineVersion version) Builder* set_time_adjusted_to_utc(c_bool adjusted) shared_ptr[ArrowWriterProperties] build() diff --git a/python/pyarrow/parquet/core.py b/python/pyarrow/parquet/core.py index 080bfa55c234..6d281400ce91 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -876,6 +876,10 @@ def _sanitize_table(table, new_schema, flavor): dictionary_pagesize_limit : int, default None Specify the dictionary page size limit per row group. If None, use the default 1MB. +experimental_vector_encoding : bool, default False + EXPERIMENTAL: Encode supported fixed-size list values, including the storage + of fixed-shape tensor extension arrays, using Parquet VECTOR repetition. + Unsupported fixed-size list values continue using standard LIST encoding. store_schema : bool, default True By default, the Arrow schema is serialized and stored in the Parquet file metadata (in the "ARROW:schema" key). When reading the file, @@ -1083,6 +1087,7 @@ def __init__(self, where, schema, filesystem=None, store_decimal_as_integer=False, write_time_adjusted_to_utc=False, max_rows_per_page=None, + experimental_vector_encoding=False, **options): if use_deprecated_int96_timestamps is None: # Use int96 timestamps for Spark @@ -1138,6 +1143,7 @@ def __init__(self, where, schema, filesystem=None, store_decimal_as_integer=store_decimal_as_integer, write_time_adjusted_to_utc=write_time_adjusted_to_utc, max_rows_per_page=max_rows_per_page, + experimental_vector_encoding=experimental_vector_encoding, **options) self.is_open = True @@ -2017,6 +2023,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_time_adjusted_to_utc=False, max_rows_per_page=None, bloom_filter_options=None, + experimental_vector_encoding=False, **kwargs): # Implementor's note: when adding keywords here / updating defaults, also # update it in write_to_dataset and _dataset_parquet.pyx ParquetFileWriteOptions @@ -2051,6 +2058,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_time_adjusted_to_utc=write_time_adjusted_to_utc, max_rows_per_page=max_rows_per_page, bloom_filter_options=bloom_filter_options, + experimental_vector_encoding=experimental_vector_encoding, **kwargs) as writer: writer.write_table(table, row_group_size=row_group_size) except Exception: diff --git a/python/pyarrow/tests/parquet/test_data_types.py b/python/pyarrow/tests/parquet/test_data_types.py index c546bc1532ac..3c965b2325e4 100644 --- a/python/pyarrow/tests/parquet/test_data_types.py +++ b/python/pyarrow/tests/parquet/test_data_types.py @@ -581,6 +581,28 @@ def test_json_extension_type(storage_type): store_schema=False) +def test_fixed_shape_tensor_vector_encoding(): + tensor_type = pa.fixed_shape_tensor(pa.float32(), [2, 3]) + storage = pa.array([ + [1, 2, 3, 4, 5, 6], + None, + [7, 8, 9, 10, 11, 12], + ], pa.list_(pa.float32(), 6)) + arr = pa.ExtensionArray.from_storage(tensor_type, storage) + table = pa.table({"tensor": arr}) + + sink = pa.BufferOutputStream() + pq.write_table(table, sink, experimental_vector_encoding=True) + buf = sink.getvalue() + + metadata = pq.read_metadata(pa.BufferReader(buf)) + assert "vector" in str(metadata.schema).lower() + + result = pq.read_table(pa.BufferReader(buf)) + assert result.schema == table.schema + assert result.equals(table) + + def test_uuid_extension_type(): data = [ b'\xe4`\xf9p\x83QGN\xac\x7f\xa4g>K\xa8\xcb',