diff --git a/cpp/src/generated/parquet_types.cpp b/cpp/src/generated/parquet_types.cpp index 0ee973f2a2d6..d3978bd151db 100644 --- a/cpp/src/generated/parquet_types.cpp +++ b/cpp/src/generated/parquet_types.cpp @@ -2255,6 +2255,213 @@ void GeographyType::printTo(std::ostream& out) const { } + +VectorElementLogicalType::~VectorElementLogicalType() noexcept { +} + +VectorElementLogicalType::VectorElementLogicalType() noexcept { +} + +void VectorElementLogicalType::__set_DECIMAL(const DecimalType& val) { + this->DECIMAL = val; +__isset.DECIMAL = true; +} + +void VectorElementLogicalType::__set_INTEGER(const IntType& val) { + this->INTEGER = val; +__isset.INTEGER = true; +} + +void VectorElementLogicalType::__set_UUID(const UUIDType& val) { + this->UUID = val; +__isset.UUID = true; +} + +void VectorElementLogicalType::__set_FLOAT16(const Float16Type& val) { + this->FLOAT16 = val; +__isset.FLOAT16 = true; +} + +std::ostream& operator<<(std::ostream& out, const VectorElementLogicalType& obj) +{ + obj.printTo(out); + return out; +} + + +void swap(VectorElementLogicalType &a, VectorElementLogicalType &b) { + using ::std::swap; + swap(a.DECIMAL, b.DECIMAL); + swap(a.INTEGER, b.INTEGER); + swap(a.UUID, b.UUID); + swap(a.FLOAT16, b.FLOAT16); + swap(a.__isset, b.__isset); +} + +bool VectorElementLogicalType::operator==(const VectorElementLogicalType & rhs) const +{ + if (__isset.DECIMAL != rhs.__isset.DECIMAL) + return false; + else if (__isset.DECIMAL && !(DECIMAL == rhs.DECIMAL)) + return false; + if (__isset.INTEGER != rhs.__isset.INTEGER) + return false; + else if (__isset.INTEGER && !(INTEGER == rhs.INTEGER)) + return false; + if (__isset.UUID != rhs.__isset.UUID) + return false; + else if (__isset.UUID && !(UUID == rhs.UUID)) + return false; + if (__isset.FLOAT16 != rhs.__isset.FLOAT16) + return false; + else if (__isset.FLOAT16 && !(FLOAT16 == rhs.FLOAT16)) + return false; + return true; +} + +VectorElementLogicalType::VectorElementLogicalType(const VectorElementLogicalType& other119) noexcept { + DECIMAL = other119.DECIMAL; + INTEGER = other119.INTEGER; + UUID = other119.UUID; + FLOAT16 = other119.FLOAT16; + __isset = other119.__isset; +} +VectorElementLogicalType::VectorElementLogicalType(VectorElementLogicalType&& other120) noexcept { + DECIMAL = std::move(other120.DECIMAL); + INTEGER = std::move(other120.INTEGER); + UUID = std::move(other120.UUID); + FLOAT16 = std::move(other120.FLOAT16); + __isset = other120.__isset; +} +VectorElementLogicalType& VectorElementLogicalType::operator=(const VectorElementLogicalType& other121) noexcept { + DECIMAL = other121.DECIMAL; + INTEGER = other121.INTEGER; + UUID = other121.UUID; + FLOAT16 = other121.FLOAT16; + __isset = other121.__isset; + return *this; +} +VectorElementLogicalType& VectorElementLogicalType::operator=(VectorElementLogicalType&& other122) noexcept { + DECIMAL = std::move(other122.DECIMAL); + INTEGER = std::move(other122.INTEGER); + UUID = std::move(other122.UUID); + FLOAT16 = std::move(other122.FLOAT16); + __isset = other122.__isset; + return *this; +} +void VectorElementLogicalType::printTo(std::ostream& out) const { + using ::apache::thrift::to_string; + out << "VectorElementLogicalType("; + out << "DECIMAL="; (__isset.DECIMAL ? (out << to_string(DECIMAL)) : (out << "")); + out << ", " << "INTEGER="; (__isset.INTEGER ? (out << to_string(INTEGER)) : (out << "")); + out << ", " << "UUID="; (__isset.UUID ? (out << to_string(UUID)) : (out << "")); + out << ", " << "FLOAT16="; (__isset.FLOAT16 ? (out << to_string(FLOAT16)) : (out << "")); + out << ")"; +} + + +VectorType::~VectorType() noexcept { +} + +VectorType::VectorType() noexcept + : length(0), + element_type(static_cast(0)), + element_type_length(0) { +} + +void VectorType::__set_length(const int32_t val) { + this->length = val; +__isset.length = true; +} + +void VectorType::__set_element_type(const Type::type val) { + this->element_type = val; +__isset.element_type = true; +} + +void VectorType::__set_element_type_length(const int32_t val) { + this->element_type_length = val; +__isset.element_type_length = true; +} + +void VectorType::__set_element_logical_type(const VectorElementLogicalType& val) { + this->element_logical_type = val; +__isset.element_logical_type = true; +} +std::ostream& operator<<(std::ostream& out, const VectorType& obj) +{ + obj.printTo(out); + return out; +} + + +void swap(VectorType &a, VectorType &b) { + using ::std::swap; + swap(a.length, b.length); + swap(a.element_type, b.element_type); + swap(a.element_type_length, b.element_type_length); + swap(a.element_logical_type, b.element_logical_type); + swap(a.__isset, b.__isset); +} + +bool VectorType::operator==(const VectorType & rhs) const +{ + if (!(length == rhs.length)) + return false; + if (!(element_type == rhs.element_type)) + return false; + if (__isset.element_type_length != rhs.__isset.element_type_length) + return false; + else if (__isset.element_type_length && !(element_type_length == rhs.element_type_length)) + return false; + if (__isset.element_logical_type != rhs.__isset.element_logical_type) + return false; + else if (__isset.element_logical_type && !(element_logical_type == rhs.element_logical_type)) + return false; + return true; +} + +VectorType::VectorType(const VectorType& other124) noexcept { + length = other124.length; + element_type = other124.element_type; + element_type_length = other124.element_type_length; + element_logical_type = other124.element_logical_type; + __isset = other124.__isset; +} +VectorType::VectorType(VectorType&& other125) noexcept { + length = other125.length; + element_type = other125.element_type; + element_type_length = other125.element_type_length; + element_logical_type = std::move(other125.element_logical_type); + __isset = other125.__isset; +} +VectorType& VectorType::operator=(const VectorType& other126) noexcept { + length = other126.length; + element_type = other126.element_type; + element_type_length = other126.element_type_length; + element_logical_type = other126.element_logical_type; + __isset = other126.__isset; + return *this; +} +VectorType& VectorType::operator=(VectorType&& other127) noexcept { + length = other127.length; + element_type = other127.element_type; + element_type_length = other127.element_type_length; + element_logical_type = std::move(other127.element_logical_type); + __isset = other127.__isset; + return *this; +} +void VectorType::printTo(std::ostream& out) const { + using ::apache::thrift::to_string; + out << "VectorType("; + out << "length=" << to_string(length); + out << ", " << "element_type=" << to_string(element_type); + out << ", " << "element_type_length="; (__isset.element_type_length ? (out << to_string(element_type_length)) : (out << "")); + out << ", " << "element_logical_type="; (__isset.element_logical_type ? (out << to_string(element_logical_type)) : (out << "")); + out << ")"; +} + + LogicalType::~LogicalType() noexcept { } @@ -2345,6 +2552,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 +2583,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 +2657,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 +2682,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 +2703,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 +2724,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 +2746,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 +2770,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 << ")"; } diff --git a/cpp/src/generated/parquet_types.h b/cpp/src/generated/parquet_types.h index 1f1e254f5cf2..5d4a73d193e0 100644 --- a/cpp/src/generated/parquet_types.h +++ b/cpp/src/generated/parquet_types.h @@ -1620,8 +1620,127 @@ void swap(GeographyType &a, GeographyType &b); std::ostream& operator<<(std::ostream& out, const GeographyType& obj); +typedef struct _VectorElementLogicalType__isset { + _VectorElementLogicalType__isset() : DECIMAL(false), INTEGER(false), UUID(false), FLOAT16(false) {} + bool DECIMAL :1; + bool INTEGER :1; + bool UUID :1; + bool FLOAT16 :1; +} _VectorElementLogicalType__isset; + +/** + * Logical type annotation for fixed-length dense vectors. + * + * VectorType describes the logical value independently from a particular + * physical representation. This version defines its use on FIXED_LEN_BYTE_ARRAY + * primitive nodes, where each non-null value stores length fixed-width elements + * packed contiguously. For length > 0, the SchemaElement type_length must equal + * length multiplied by the byte width of the element physical type. For + * length == 0, the SchemaElement type_length must be 1 and the single physical + * byte is padding with no logical meaning. + */ +class VectorElementLogicalType { + public: + + VectorElementLogicalType(const VectorElementLogicalType&) noexcept; + VectorElementLogicalType(VectorElementLogicalType&&) noexcept; + VectorElementLogicalType& operator=(const VectorElementLogicalType&) noexcept; + VectorElementLogicalType& operator=(VectorElementLogicalType&&) noexcept; + VectorElementLogicalType() noexcept; + + virtual ~VectorElementLogicalType() noexcept; + DecimalType DECIMAL; + IntType INTEGER; + UUIDType UUID; + Float16Type FLOAT16; + + _VectorElementLogicalType__isset __isset; + + void __set_DECIMAL(const DecimalType& val); + + void __set_INTEGER(const IntType& val); + + void __set_UUID(const UUIDType& val); + + void __set_FLOAT16(const Float16Type& val); + + bool operator == (const VectorElementLogicalType & rhs) const; + bool operator != (const VectorElementLogicalType &rhs) const { + return !(*this == rhs); + } + + bool operator < (const VectorElementLogicalType & ) const; + + template + uint32_t read(Protocol_* iprot); + template + uint32_t write(Protocol_* oprot) const; + + virtual void printTo(std::ostream& out) const; +}; + +void swap(VectorElementLogicalType &a, VectorElementLogicalType &b); + +std::ostream& operator<<(std::ostream& out, const VectorElementLogicalType& obj); + +typedef struct _VectorType__isset { + _VectorType__isset() : length(false), element_type(false), element_type_length(false), element_logical_type(false) {} + bool length :1; + bool element_type :1; + bool element_type_length :1; + bool element_logical_type :1; +} _VectorType__isset; + +class VectorType { + public: + + VectorType(const VectorType&) noexcept; + VectorType(VectorType&&) noexcept; + VectorType& operator=(const VectorType&) noexcept; + VectorType& operator=(VectorType&&) noexcept; + VectorType() noexcept; + + virtual ~VectorType() noexcept; + int32_t length; + /** + * + * @see Type + */ + Type::type element_type; + int32_t element_type_length; + VectorElementLogicalType element_logical_type; + + _VectorType__isset __isset; + + void __set_length(const int32_t val); + + void __set_element_type(const Type::type val); + + void __set_element_type_length(const int32_t val); + + void __set_element_logical_type(const VectorElementLogicalType& val); + + bool operator == (const VectorType & rhs) const; + bool operator != (const VectorType &rhs) const { + return !(*this == rhs); + } + + bool operator < (const VectorType & ) const; + + template + uint32_t read(Protocol_* iprot); + template + uint32_t write(Protocol_* oprot) const; + + virtual void printTo(std::ostream& out) const; +}; + +void swap(VectorType &a, VectorType &b); + +std::ostream& operator<<(std::ostream& out, const VectorType& 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 +1758,7 @@ typedef struct _LogicalType__isset { bool VARIANT :1; bool GEOMETRY :1; bool GEOGRAPHY :1; + bool VECTOR :1; } _LogicalType__isset; /** @@ -1675,6 +1795,7 @@ class LogicalType { VariantType VARIANT; GeometryType GEOMETRY; GeographyType GEOGRAPHY; + VectorType VECTOR; _LogicalType__isset __isset; @@ -1712,6 +1833,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); diff --git a/cpp/src/generated/parquet_types.tcc b/cpp/src/generated/parquet_types.tcc index 78e3e2549394..dea5348d054d 100644 --- a/cpp/src/generated/parquet_types.tcc +++ b/cpp/src/generated/parquet_types.tcc @@ -1625,6 +1625,209 @@ uint32_t GeographyType::write(Protocol_* oprot) const { return xfer; } + +template +uint32_t VectorElementLogicalType::read(Protocol_* iprot) { + + ::apache::thrift::protocol::TInputRecursionTracker tracker(*iprot); + uint32_t xfer = 0; + std::string fname; + ::apache::thrift::protocol::TType ftype; + int16_t fid; + + xfer += iprot->readStructBegin(fname); + + using ::apache::thrift::protocol::TProtocolException; + + + while (true) + { + xfer += iprot->readFieldBegin(fname, ftype, fid); + if (ftype == ::apache::thrift::protocol::T_STOP) { + break; + } + switch (fid) + { + case 1: + if (ftype == ::apache::thrift::protocol::T_STRUCT) { + xfer += this->DECIMAL.read(iprot); + this->__isset.DECIMAL = true; + } else { + xfer += iprot->skip(ftype); + } + break; + case 2: + if (ftype == ::apache::thrift::protocol::T_STRUCT) { + xfer += this->INTEGER.read(iprot); + this->__isset.INTEGER = true; + } else { + xfer += iprot->skip(ftype); + } + break; + case 3: + if (ftype == ::apache::thrift::protocol::T_STRUCT) { + xfer += this->UUID.read(iprot); + this->__isset.UUID = true; + } else { + xfer += iprot->skip(ftype); + } + break; + case 4: + if (ftype == ::apache::thrift::protocol::T_STRUCT) { + xfer += this->FLOAT16.read(iprot); + this->__isset.FLOAT16 = true; + } else { + xfer += iprot->skip(ftype); + } + break; + default: + xfer += iprot->skip(ftype); + break; + } + xfer += iprot->readFieldEnd(); + } + + xfer += iprot->readStructEnd(); + + return xfer; +} + +template +uint32_t VectorElementLogicalType::write(Protocol_* oprot) const { + uint32_t xfer = 0; + ::apache::thrift::protocol::TOutputRecursionTracker tracker(*oprot); + xfer += oprot->writeStructBegin("VectorElementLogicalType"); + + if (this->__isset.DECIMAL) { + xfer += oprot->writeFieldBegin("DECIMAL", ::apache::thrift::protocol::T_STRUCT, 1); + xfer += this->DECIMAL.write(oprot); + xfer += oprot->writeFieldEnd(); + } + if (this->__isset.INTEGER) { + xfer += oprot->writeFieldBegin("INTEGER", ::apache::thrift::protocol::T_STRUCT, 2); + xfer += this->INTEGER.write(oprot); + xfer += oprot->writeFieldEnd(); + } + if (this->__isset.UUID) { + xfer += oprot->writeFieldBegin("UUID", ::apache::thrift::protocol::T_STRUCT, 3); + xfer += this->UUID.write(oprot); + xfer += oprot->writeFieldEnd(); + } + if (this->__isset.FLOAT16) { + xfer += oprot->writeFieldBegin("FLOAT16", ::apache::thrift::protocol::T_STRUCT, 4); + xfer += this->FLOAT16.write(oprot); + xfer += oprot->writeFieldEnd(); + } + xfer += oprot->writeFieldStop(); + xfer += oprot->writeStructEnd(); + return xfer; +} + +template +uint32_t VectorType::read(Protocol_* iprot) { + + ::apache::thrift::protocol::TInputRecursionTracker tracker(*iprot); + uint32_t xfer = 0; + std::string fname; + ::apache::thrift::protocol::TType ftype; + int16_t fid; + + xfer += iprot->readStructBegin(fname); + + using ::apache::thrift::protocol::TProtocolException; + + bool isset_length = false; + bool isset_element_type = false; + + while (true) + { + xfer += iprot->readFieldBegin(fname, ftype, fid); + if (ftype == ::apache::thrift::protocol::T_STOP) { + break; + } + switch (fid) + { + case 1: + if (ftype == ::apache::thrift::protocol::T_I32) { + xfer += iprot->readI32(this->length); + this->__isset.length = true; + isset_length = true; + } else { + xfer += iprot->skip(ftype); + } + break; + case 2: + if (ftype == ::apache::thrift::protocol::T_I32) { + int32_t ecast123; + xfer += iprot->readI32(ecast123); + this->element_type = static_cast(ecast123); + this->__isset.element_type = true; + isset_element_type = true; + } else { + xfer += iprot->skip(ftype); + } + break; + case 3: + if (ftype == ::apache::thrift::protocol::T_I32) { + xfer += iprot->readI32(this->element_type_length); + this->__isset.element_type_length = true; + } else { + xfer += iprot->skip(ftype); + } + break; + case 4: + if (ftype == ::apache::thrift::protocol::T_STRUCT) { + xfer += this->element_logical_type.read(iprot); + this->__isset.element_logical_type = true; + } else { + xfer += iprot->skip(ftype); + } + break; + default: + xfer += iprot->skip(ftype); + break; + } + xfer += iprot->readFieldEnd(); + } + + xfer += iprot->readStructEnd(); + + if (!isset_length) + throw TProtocolException(TProtocolException::INVALID_DATA); + if (!isset_element_type) + throw TProtocolException(TProtocolException::INVALID_DATA); + return xfer; +} + +template +uint32_t VectorType::write(Protocol_* oprot) const { + uint32_t xfer = 0; + ::apache::thrift::protocol::TOutputRecursionTracker tracker(*oprot); + xfer += oprot->writeStructBegin("VectorType"); + + xfer += oprot->writeFieldBegin("length", ::apache::thrift::protocol::T_I32, 1); + xfer += oprot->writeI32(this->length); + xfer += oprot->writeFieldEnd(); + + xfer += oprot->writeFieldBegin("element_type", ::apache::thrift::protocol::T_I32, 2); + xfer += oprot->writeI32(static_cast(this->element_type)); + xfer += oprot->writeFieldEnd(); + + if (this->__isset.element_type_length) { + xfer += oprot->writeFieldBegin("element_type_length", ::apache::thrift::protocol::T_I32, 3); + xfer += oprot->writeI32(this->element_type_length); + xfer += oprot->writeFieldEnd(); + } + if (this->__isset.element_logical_type) { + xfer += oprot->writeFieldBegin("element_logical_type", ::apache::thrift::protocol::T_STRUCT, 4); + xfer += this->element_logical_type.write(oprot); + xfer += oprot->writeFieldEnd(); + } + xfer += oprot->writeFieldStop(); + xfer += oprot->writeStructEnd(); + return xfer; +} + template uint32_t LogicalType::read(Protocol_* iprot) { @@ -1783,6 +1986,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 +2097,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; diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index d29458bf226b..6433aa4b8966 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -3341,6 +3341,50 @@ TEST(ArrowReadWrite, FixedSizeList) { CheckSimpleRoundtrip(table, 2, props_store_schema); } +TEST(ArrowReadWrite, FixedSizeListAsVector) { + using ::arrow::field; + using ::arrow::fixed_size_list; + using ::arrow::internal::checked_cast; + + auto type = fixed_size_list(field("element", ::arrow::float32(), false), /*size=*/3); + auto array = ::arrow::ArrayFromJSON(type, R"([[1, 2, 3], null, [4, 5, 6]])"); + auto table = ::arrow::Table::Make(::arrow::schema({field("root", type)}), {array}); + auto writer_props = + ArrowWriterProperties::Builder().enable_fixed_size_list_as_vector()->build(); + + std::shared_ptr buffer; + ASSERT_NO_FATAL_FAILURE(WriteTableToBuffer(table, 2, writer_props, &buffer)); + auto parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + const auto* descr = parquet_reader->metadata()->schema()->Column(0); + ASSERT_EQ(Type::FIXED_LEN_BYTE_ARRAY, descr->physical_type()); + ASSERT_EQ(12, descr->type_length()); + ASSERT_TRUE(descr->logical_type()->is_vector()); + const auto& vector_type = + checked_cast(*descr->logical_type()); + ASSERT_EQ(Type::FLOAT, vector_type.element_type()); + ASSERT_EQ(3, vector_type.length()); + + CheckSimpleRoundtrip(table, 2, writer_props); + + auto zero_length_type = + fixed_size_list(field("element", ::arrow::float32(), false), /*size=*/0); + auto zero_length_array = ::arrow::ArrayFromJSON(zero_length_type, R"([[], null, []])"); + auto zero_length_table = ::arrow::Table::Make( + ::arrow::schema({field("root", zero_length_type)}), {zero_length_array}); + ASSERT_NO_FATAL_FAILURE( + WriteTableToBuffer(zero_length_table, 2, writer_props, &buffer)); + parquet_reader = ParquetFileReader::Open(std::make_shared(buffer)); + descr = parquet_reader->metadata()->schema()->Column(0); + ASSERT_EQ(Type::FIXED_LEN_BYTE_ARRAY, descr->physical_type()); + ASSERT_EQ(1, descr->type_length()); + ASSERT_TRUE(descr->logical_type()->is_vector()); + const auto& zero_length_vector_type = + checked_cast(*descr->logical_type()); + ASSERT_EQ(0, zero_length_vector_type.length()); + + CheckSimpleRoundtrip(zero_length_table, 2, writer_props); +} + TEST(ArrowReadWrite, ListOfStructOfList2) { using ::arrow::field; using ::arrow::list; diff --git a/cpp/src/parquet/arrow/reader.cc b/cpp/src/parquet/arrow/reader.cc index a60af69aec9f..a810cf8f72e8 100644 --- a/cpp/src/parquet/arrow/reader.cc +++ b/cpp/src/parquet/arrow/reader.cc @@ -544,7 +544,54 @@ class LeafReader : public ColumnReaderImpl { std::shared_ptr record_reader_; }; -// Column reader for extension arrays +// Column reader for packed VECTOR arrays +class VectorFLBAReader : public ColumnReaderImpl { + public: + VectorFLBAReader(std::shared_ptr field, + std::unique_ptr storage_reader) + : field_(std::move(field)), storage_reader_(std::move(storage_reader)) {} + + Status GetDefLevels(const int16_t** data, int64_t* length) override { + return storage_reader_->GetDefLevels(data, length); + } + + Status GetRepLevels(const int16_t** data, int64_t* length) override { + return storage_reader_->GetRepLevels(data, length); + } + + Status LoadBatch(int64_t number_of_records) final { + return storage_reader_->LoadBatch(number_of_records); + } + + Status BuildArray(int64_t length_upper_bound, + std::shared_ptr* out) override { + std::shared_ptr storage; + RETURN_NOT_OK(storage_reader_->BuildArray(length_upper_bound, &storage)); + ARROW_ASSIGN_OR_RAISE(std::shared_ptr storage_data, + ChunksToSingle(*storage)); + + const auto& list_type = + checked_cast(*field_->type()); + const int64_t child_length = storage_data->length * list_type.list_size(); + auto child_data = ::arrow::ArrayData::Make(list_type.value_type(), child_length, + {nullptr, storage_data.buffers[1]}, + /*null_count=*/0); + auto data = ::arrow::ArrayData::Make( + field_->type(), storage_data->length, {storage_data->buffers[0]}, + {std::move(child_data)}, storage_data->null_count); + *out = std::make_shared(::arrow::MakeArray(std::move(data))); + return Status::OK(); + } + + bool IsOrHasRepeatedChild() const final { return false; } + + const std::shared_ptr field() override { return field_; } + + private: + std::shared_ptr field_; + std::unique_ptr storage_reader_; +}; + class ExtensionReader : public ColumnReaderImpl { public: ExtensionReader(std::shared_ptr field, @@ -884,8 +931,17 @@ 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); + const auto* descr = ctx->reader->metadata()->schema()->Column(field.column_index); + if (descr->logical_type()->is_vector()) { + auto storage_field = + arrow_field->WithType(::arrow::fixed_size_binary(descr->type_length())); + *out = std::make_unique( + arrow_field, std::make_unique(ctx, storage_field, std::move(input), + field.level_info)); + } else { + *out = std::make_unique(ctx, arrow_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) { diff --git a/cpp/src/parquet/arrow/schema.cc b/cpp/src/parquet/arrow/schema.cc index 9c4c462c6b8c..9f1123d16ba9 100644 --- a/cpp/src/parquet/arrow/schema.cc +++ b/cpp/src/parquet/arrow/schema.cc @@ -78,6 +78,46 @@ Repetition::type RepetitionFromNullable(bool is_nullable) { return is_nullable ? Repetition::OPTIONAL : Repetition::REQUIRED; } +struct VectorElementStorage { + ParquetType::type type; + int32_t byte_width; + int32_t type_length = -1; + std::shared_ptr logical_type; +}; + +Result GetVectorElementStorage(const ::arrow::DataType& type) { + switch (type.id()) { + case ArrowTypeId::INT32: + return VectorElementStorage{ParquetType::INT32, 4, -1, nullptr}; + case ArrowTypeId::UINT32: + return VectorElementStorage{ParquetType::INT32, 4, -1, + LogicalType::Int(32, /*is_signed=*/false)}; + case ArrowTypeId::INT64: + return VectorElementStorage{ParquetType::INT64, 8, -1, nullptr}; + case ArrowTypeId::UINT64: + return VectorElementStorage{ParquetType::INT64, 8, -1, + LogicalType::Int(64, /*is_signed=*/false)}; + case ArrowTypeId::FLOAT: + return VectorElementStorage{ParquetType::FLOAT, 4, -1, nullptr}; + case ArrowTypeId::DOUBLE: + return VectorElementStorage{ParquetType::DOUBLE, 8, -1, nullptr}; + case ArrowTypeId::HALF_FLOAT: + return VectorElementStorage{ParquetType::FIXED_LEN_BYTE_ARRAY, 2, 2, + LogicalType::Float16()}; + case ArrowTypeId::FIXED_SIZE_BINARY: { + const auto& fixed_size_binary_type = + ::arrow::internal::checked_cast(type); + return VectorElementStorage{ParquetType::FIXED_LEN_BYTE_ARRAY, + fixed_size_binary_type.byte_width(), + fixed_size_binary_type.byte_width(), nullptr}; + } + default: + return Status::NotImplemented( + "VECTOR as FIXED_LEN_BYTE_ARRAY is only implemented for int32, uint32, " + "int64, uint64, float, double, float16, and fixed_size_binary elements"); + } +} + Result> MakeArrowList( std::shared_ptr field, const ArrowReaderProperties& props) { switch (props.list_type()) { @@ -453,7 +493,26 @@ Status FieldToNode(const std::string& name, const std::shared_ptr& field, return StructToNode(struct_type, name, field->nullable(), field_id, properties, arrow_properties, out); } - case ArrowTypeId::FIXED_SIZE_LIST: + case ArrowTypeId::FIXED_SIZE_LIST: { + auto fixed_size_list_type = + std::static_pointer_cast<::arrow::FixedSizeListType>(field->type()); + if (arrow_properties.write_fixed_size_list_as_vector() && + !fixed_size_list_type->value_field()->nullable()) { + ARROW_ASSIGN_OR_RAISE( + auto element_storage, + GetVectorElementStorage(*fixed_size_list_type->value_type())); + type = ParquetType::FIXED_LEN_BYTE_ARRAY; + length = fixed_size_list_type->list_size() == 0 + ? 1 + : element_storage.byte_width * fixed_size_list_type->list_size(); + logical_type = LogicalType::Vector( + element_storage.type, fixed_size_list_type->list_size(), + element_storage.type_length, element_storage.logical_type); + break; + } + return ListToNode(fixed_size_list_type, name, field->nullable(), field_id, + properties, arrow_properties, out); + } case ArrowTypeId::LARGE_LIST: case ArrowTypeId::LIST: { auto list_type = std::static_pointer_cast<::arrow::BaseListType>(field->type()); diff --git a/cpp/src/parquet/arrow/schema_internal.cc b/cpp/src/parquet/arrow/schema_internal.cc index 2e8cf764b27f..c56c713cc863 100644 --- a/cpp/src/parquet/arrow/schema_internal.cc +++ b/cpp/src/parquet/arrow/schema_internal.cc @@ -191,10 +191,90 @@ Result> FromByteArray( } } +Result> FromVectorElement( + const VectorLogicalType& vector_type, + const ArrowReaderProperties& reader_properties) { + const auto element_logical_type = vector_type.element_logical_type(); + const bool has_logical_type = element_logical_type != nullptr; + + switch (vector_type.element_type()) { + case ParquetType::INT32: + if (!has_logical_type) { + return ::arrow::int32(); + } + if (element_logical_type->is_int()) { + const auto& integer = checked_cast(*element_logical_type); + if (integer.bit_width() != 32) { + return Status::NotImplemented("VECTOR IntType elements with bit width ", + integer.bit_width(), + " require value conversion from physical INT32"); + } + return integer.is_signed() ? ::arrow::int32() : ::arrow::uint32(); + } + break; + case ParquetType::INT64: + if (!has_logical_type) { + return ::arrow::int64(); + } + if (element_logical_type->is_int()) { + const auto& integer = checked_cast(*element_logical_type); + if (integer.bit_width() != 64) { + return Status::NotImplemented("VECTOR IntType elements with bit width ", + integer.bit_width(), + " require value conversion from physical INT64"); + } + return integer.is_signed() ? ::arrow::int64() : ::arrow::uint64(); + } + break; + case ParquetType::FLOAT: + if (!has_logical_type) { + return ::arrow::float32(); + } + break; + case ParquetType::DOUBLE: + if (!has_logical_type) { + return ::arrow::float64(); + } + break; + case ParquetType::FIXED_LEN_BYTE_ARRAY: + if (!has_logical_type) { + return ::arrow::fixed_size_binary(vector_type.element_type_length()); + } + if (element_logical_type->is_float16()) { + return ::arrow::float16(); + } + if (element_logical_type->is_UUID()) { + if (vector_type.element_type_length() == 16 && + reader_properties.get_arrow_extensions_enabled()) { + return ::arrow::extension::uuid(); + } + return ::arrow::fixed_size_binary(vector_type.element_type_length()); + } + break; + default: + break; + } + + return Status::NotImplemented("Unsupported VECTOR element logical type ", + has_logical_type ? element_logical_type->ToString() + : LogicalType::None()->ToString(), + " for physical element type ", + TypeToString(vector_type.element_type())); +} + Result> FromFLBA( const LogicalType& logical_type, int32_t physical_length, const ArrowReaderProperties& reader_properties) { switch (logical_type.type()) { + case LogicalType::Type::VECTOR: { + const auto& vector_type = + ::arrow::internal::checked_cast(logical_type); + ARROW_ASSIGN_OR_RAISE(auto value_type, + FromVectorElement(vector_type, reader_properties)); + return ::arrow::fixed_size_list( + ::arrow::field("element", std::move(value_type), /*nullable=*/false), + vector_type.length()); + } case LogicalType::Type::DECIMAL: return MakeArrowDecimal(logical_type, reader_properties.smallest_decimal_enabled()); case LogicalType::Type::FLOAT16: diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index 4b2b06e5e097..037ffad66a44 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -18,6 +18,7 @@ #include "parquet/arrow/writer.h" #include +#include #include #include #include @@ -56,6 +57,7 @@ using arrow::ExtensionArray; using arrow::ExtensionType; using arrow::Field; using arrow::FixedSizeBinaryArray; +using arrow::FixedSizeListArray; using arrow::ListArray; using arrow::MemoryPool; using arrow::NumericArray; @@ -107,6 +109,65 @@ bool HasNullableRoot(const SchemaManifest& schema_manifest, return nullable; } +Result> FixedSizeListToFixedSizeBinary( + const ChunkedArray& data, int64_t offset, int64_t size, int32_t storage_byte_width, + MemoryPool* pool) { + auto sliced = data.Slice(offset, size); + std::vector> chunks; + chunks.reserve(sliced->num_chunks()); + + for (const auto& chunk : sliced->chunks()) { + const auto& list_array = checked_cast(*chunk); + const auto& list_type = + checked_cast(*chunk->type()); + const auto& value_type = + checked_cast(*list_type.value_type()); + const int32_t value_byte_width = value_type.bit_width() / 8; + const int32_t expected_storage_byte_width = + list_type.list_size() == 0 ? storage_byte_width + : list_type.list_size() * value_byte_width; + if (expected_storage_byte_width != storage_byte_width) { + return Status::Invalid("VECTOR storage byte width mismatch: expected ", + expected_storage_byte_width, " but Parquet schema has ", + storage_byte_width); + } + const int64_t data_size = chunk->length() * storage_byte_width; + + auto values = list_array.values(); + if (values->null_count() != 0) { + for (int64_t i = 0; i < list_array.length(); ++i) { + if (list_array.IsNull(i)) { + continue; + } + const int64_t value_offset = list_array.value_offset(i); + for (int32_t j = 0; j < list_type.list_size(); ++j) { + if (values->IsNull(value_offset + j)) { + return Status::Invalid( + "Cannot write VECTOR as FIXED_LEN_BYTE_ARRAY when elements are null"); + } + } + } + } + ARROW_ASSIGN_OR_RAISE(auto data_buffer, AllocateBuffer(data_size, pool)); + if (list_type.list_size() == 0) { + std::memset(data_buffer->mutable_data(), 0, data_size); + } else { + const auto& values_data = values->data()->buffers[1]; + const uint8_t* source = + values_data->data() + list_array.value_offset(0) * value_byte_width; + std::memcpy(data_buffer->mutable_data(), source, data_size); + } + + auto storage_type = ::arrow::fixed_size_binary(storage_byte_width); + auto array_data = + ::arrow::ArrayData::Make(storage_type, chunk->length(), + {chunk->data()->buffers[0], std::move(data_buffer)}, + chunk->null_count(), /*offset=*/0); + chunks.push_back(::arrow::MakeArray(std::move(array_data))); + } + return std::make_shared(std::move(chunks)); +} + Status GetSchemaMetadata(const ::arrow::Schema& schema, ::arrow::MemoryPool* pool, const ArrowWriterProperties& properties, std::shared_ptr* out) { @@ -375,9 +436,20 @@ class FileWriterImpl : public FileWriter { if (row_group_writer_->buffered()) { return Status::Invalid("Cannot write column chunk into the buffered row group."); } + std::shared_ptr data_to_write = data; + const int column_index = row_group_writer_->current_column() + 1; + const auto* column_descr = writer_->schema()->Column(column_index); + const auto* parquet_node = column_descr->schema_node().get(); + if (parquet_node->logical_type()->is_vector()) { + ARROW_ASSIGN_OR_RAISE(data_to_write, + FixedSizeListToFixedSizeBinary( + *data, offset, size, column_descr->type_length(), + column_write_context_.memory_pool)); + offset = 0; + } ARROW_ASSIGN_OR_RAISE( std::unique_ptr writer, - ArrowColumnWriterV2::Make(*data, offset, size, schema_manifest_, + ArrowColumnWriterV2::Make(*data_to_write, offset, size, schema_manifest_, row_group_writer_)); return writer->Write(&column_write_context_); } @@ -456,11 +528,22 @@ class FileWriterImpl : public FileWriter { int column_index_start = 0; for (int i = 0; i < batch.num_columns(); i++) { - 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)); + std::shared_ptr chunked_array = + std::make_shared(batch.column(i)); + int64_t column_offset = offset; + const auto* column_descr = writer_->schema()->Column(column_index_start); + const auto* parquet_node = column_descr->schema_node().get(); + if (parquet_node->logical_type()->is_vector()) { + ARROW_ASSIGN_OR_RAISE(chunked_array, FixedSizeListToFixedSizeBinary( + *chunked_array, offset, size, + column_descr->type_length(), + column_write_context_.memory_pool)); + column_offset = 0; + } + ARROW_ASSIGN_OR_RAISE(std::unique_ptr writer, + ArrowColumnWriterV2::Make( + *chunked_array, column_offset, size, schema_manifest_, + row_group_writer_, 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/parquet.thrift b/cpp/src/parquet/parquet.thrift index e3cc5adb9648..38ef4df7a6f4 100644 --- a/cpp/src/parquet/parquet.thrift +++ b/cpp/src/parquet/parquet.thrift @@ -462,6 +462,31 @@ struct GeographyType { 2: optional EdgeInterpolationAlgorithm algorithm; } +/** + * Logical type annotation for fixed-length dense vectors. + * + * VectorType describes the logical value independently from a particular + * physical representation. This version defines its use on FIXED_LEN_BYTE_ARRAY + * primitive nodes, where each non-null value stores length fixed-width elements + * packed contiguously. For length > 0, the SchemaElement type_length must equal + * length multiplied by the byte width of the element physical type. For + * length == 0, the SchemaElement type_length must be 1 and the single physical + * byte is padding with no logical meaning. + */ +union VectorElementLogicalType { + 1: DecimalType DECIMAL + 2: IntType INTEGER + 3: UUIDType UUID + 4: Float16Type FLOAT16 +} + +struct VectorType { + 1: required i32 length; + 2: required Type element_type; + 3: optional i32 element_type_length; + 4: optional VectorElementLogicalType element_logical_type; +} + /** * LogicalType annotations to replace ConvertedType. * @@ -495,6 +520,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 } /** diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 6634bac4f684..8d207e4a6081 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,13 @@ class PARQUET_EXPORT ArrowWriterProperties { return this; } + /// \brief EXPERIMENTAL: Write eligible Arrow FixedSizeList columns as + /// FIXED_LEN_BYTE_ARRAY annotated with the Parquet VECTOR logical type. + Builder* enable_fixed_size_list_as_vector() { + 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 +1417,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 +1430,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 +1457,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 +1485,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 +1495,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 +1507,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/types.cc b/cpp/src/parquet/types.cc index fb4eb92a7544..d23af22f5744 100644 --- a/cpp/src/parquet/types.cc +++ b/cpp/src/parquet/types.cc @@ -519,6 +519,28 @@ std::shared_ptr LogicalType::FromConvertedType( return UndefinedLogicalType::Make(); } +namespace { + +std::shared_ptr VectorElementLogicalTypeFromThrift( + const format::VectorElementLogicalType& type) { + if (type.__isset.DECIMAL) { + return LogicalType::Decimal(type.DECIMAL.precision, type.DECIMAL.scale); + } + if (type.__isset.INTEGER) { + return LogicalType::Int(static_cast(type.INTEGER.bitWidth), + type.INTEGER.isSigned); + } + if (type.__isset.UUID) { + return LogicalType::UUID(); + } + if (type.__isset.FLOAT16) { + return LogicalType::Float16(); + } + return nullptr; +} + +} // namespace + std::shared_ptr LogicalType::FromThrift( const format::LogicalType& type) { if (type.__isset.STRING) { @@ -602,6 +624,16 @@ std::shared_ptr LogicalType::FromThrift( } return VariantLogicalType::Make(spec_version); + } else if (type.__isset.VECTOR) { + std::shared_ptr element_logical_type; + if (type.VECTOR.__isset.element_logical_type) { + element_logical_type = + VectorElementLogicalTypeFromThrift(type.VECTOR.element_logical_type); + } + return VectorLogicalType::Make( + static_cast(type.VECTOR.element_type), type.VECTOR.length, + type.VECTOR.__isset.element_type_length ? type.VECTOR.element_type_length : -1, + std::move(element_logical_type)); } else { // Sentinel type for one we do not recognize return UndefinedLogicalType::Make(); @@ -673,6 +705,13 @@ std::shared_ptr LogicalType::Variant(int8_t spec_version) { return VariantLogicalType::Make(spec_version); } +std::shared_ptr LogicalType::Vector( + parquet::Type::type element_type, int32_t length, int32_t element_type_length, + std::shared_ptr element_logical_type) { + return VectorLogicalType::Make(element_type, length, element_type_length, + std::move(element_logical_type)); +} + std::shared_ptr LogicalType::None() { return NoLogicalType::Make(); } /* @@ -758,6 +797,7 @@ class LogicalType::Impl { class Geometry; class Geography; class Variant; + class Vector; class No; class Undefined; @@ -839,6 +879,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; @@ -2016,6 +2057,202 @@ std::shared_ptr VariantLogicalType::Make(const int8_t spec_ve return logical_type; } +namespace { + +int32_t VectorElementByteWidth(parquet::Type::type element_type, + int32_t element_type_length) { + switch (element_type) { + case Type::INT32: + case Type::FLOAT: + return 4; + case Type::INT64: + case Type::DOUBLE: + return 8; + case Type::FIXED_LEN_BYTE_ARRAY: + return element_type_length > 0 ? element_type_length : 0; + default: + return 0; + } +} + +int64_t VectorPhysicalTypeLength(int32_t length, int32_t element_width) { + return length == 0 ? 1 : static_cast(length) * element_width; +} + +bool IsSupportedVectorElementLogicalType(const LogicalType& logical_type) { + return logical_type.is_decimal() || logical_type.is_int() || logical_type.is_UUID() || + logical_type.is_float16(); +} + +format::VectorElementLogicalType VectorElementLogicalTypeToThrift( + const LogicalType& logical_type) { + format::VectorElementLogicalType out; + format::LogicalType thrift = logical_type.ToThrift(); + if (thrift.__isset.DECIMAL) { + out.__set_DECIMAL(thrift.DECIMAL); + } else if (thrift.__isset.INTEGER) { + out.__set_INTEGER(thrift.INTEGER); + } else if (thrift.__isset.UUID) { + out.__set_UUID(thrift.UUID); + } else if (thrift.__isset.FLOAT16) { + out.__set_FLOAT16(thrift.FLOAT16); + } else { + throw ParquetException("Unsupported VECTOR element logical type: ", + logical_type.ToString()); + } + return out; +} + +} // namespace + +class LogicalType::Impl::Vector final : public LogicalType::Impl::Incompatible, + public LogicalType::Impl::Applicable { + public: + friend class VectorLogicalType; + + bool is_applicable(parquet::Type::type primitive_type, + int32_t primitive_length = -1) const override; + std::string ToString() const override; + std::string ToJSON() const override; + format::LogicalType ToThrift() const override; + bool Equals(const LogicalType& other) const override; + + parquet::Type::type element_type() const { return element_type_; } + int32_t length() const { return length_; } + int32_t element_type_length() const { return element_type_length_; } + std::shared_ptr element_logical_type() const { + return element_logical_type_; + } + + private: + Vector(parquet::Type::type element_type, int32_t length, int32_t element_type_length, + std::shared_ptr element_logical_type) + : LogicalType::Impl(LogicalType::Type::VECTOR, SortOrder::UNKNOWN), + element_type_(element_type), + length_(length), + element_type_length_(element_type_length), + element_logical_type_(std::move(element_logical_type)) {} + + parquet::Type::type element_type_; + int32_t length_; + int32_t element_type_length_; + std::shared_ptr element_logical_type_; +}; + +bool LogicalType::Impl::Vector::is_applicable(parquet::Type::type primitive_type, + int32_t primitive_length) const { + const int32_t element_width = + VectorElementByteWidth(element_type_, element_type_length_); + return primitive_type == parquet::Type::FIXED_LEN_BYTE_ARRAY && element_width > 0 && + primitive_length == VectorPhysicalTypeLength(length_, element_width); +} + +parquet::Type::type VectorLogicalType::element_type() const { + return (dynamic_cast(*impl_)).element_type(); +} + +int32_t VectorLogicalType::length() const { + return (dynamic_cast(*impl_)).length(); +} + +int32_t VectorLogicalType::element_type_length() const { + return (dynamic_cast(*impl_)).element_type_length(); +} + +std::shared_ptr VectorLogicalType::element_logical_type() const { + return (dynamic_cast(*impl_)).element_logical_type(); +} + +std::string LogicalType::Impl::Vector::ToString() const { + std::stringstream type; + type << "Vector(element_type=" << TypeToString(element_type_); + if (element_type_length_ > 0) { + type << "(" << element_type_length_ << ")"; + } + type << ", length=" << length_; + if (element_logical_type_ != nullptr) { + type << ", element_logical_type=" << element_logical_type_->ToString(); + } + type << ")"; + return type.str(); +} + +std::string LogicalType::Impl::Vector::ToJSON() const { + std::stringstream json; + json << R"({"Type": "Vector", "element_type": ")" << TypeToString(element_type_) + << R"(", "length": )" << length_; + if (element_type_length_ > 0) { + json << R"(, "element_type_length": )" << element_type_length_; + } + if (element_logical_type_ != nullptr) { + json << R"(, "element_logical_type": )" << element_logical_type_->ToJSON(); + } + json << "}"; + return json.str(); +} + +format::LogicalType LogicalType::Impl::Vector::ToThrift() const { + format::LogicalType type; + format::VectorType vector_type; + vector_type.__set_length(length_); + vector_type.__set_element_type(static_cast(element_type_)); + if (element_type_length_ > 0) { + vector_type.__set_element_type_length(element_type_length_); + } + if (element_logical_type_ != nullptr) { + vector_type.__set_element_logical_type( + VectorElementLogicalTypeToThrift(*element_logical_type_)); + } + type.__set_VECTOR(vector_type); + return type; +} + +bool LogicalType::Impl::Vector::Equals(const LogicalType& other) const { + if (other.type() != LogicalType::Type::VECTOR) { + return false; + } + const auto& vector = dynamic_cast(other); + if (vector.element_type() != element_type_ || vector.length() != length_ || + vector.element_type_length() != element_type_length_) { + return false; + } + const auto other_element_logical_type = vector.element_logical_type(); + if (element_logical_type_ == nullptr || other_element_logical_type == nullptr) { + return element_logical_type_ == nullptr && other_element_logical_type == nullptr; + } + return element_logical_type_->Equals(*other_element_logical_type); +} + +std::shared_ptr VectorLogicalType::Make( + parquet::Type::type element_type, int32_t length, int32_t element_type_length, + std::shared_ptr element_logical_type) { + if (length < 0) { + throw ParquetException("Vector logical type requires non-negative length"); + } + if (VectorElementByteWidth(element_type, element_type_length) == 0) { + throw ParquetException("Unsupported VECTOR element type: ", + TypeToString(element_type)); + } + if (element_logical_type != nullptr && !element_logical_type->is_none()) { + if (!IsSupportedVectorElementLogicalType(*element_logical_type)) { + throw ParquetException("Unsupported VECTOR element logical type: ", + element_logical_type->ToString()); + } + if (!element_logical_type->is_applicable(element_type, element_type_length)) { + throw ParquetException(element_logical_type->ToString(), + " cannot annotate VECTOR element type ", + TypeToString(element_type)); + } + } else { + element_logical_type = nullptr; + } + + auto logical_type = std::shared_ptr(new VectorLogicalType()); + logical_type->impl_.reset(new LogicalType::Impl::Vector( + element_type, length, element_type_length, std::move(element_logical_type))); + return logical_type; +} + 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..c3318bae11f7 100644 --- a/cpp/src/parquet/types.h +++ b/cpp/src/parquet/types.h @@ -162,6 +162,7 @@ class PARQUET_EXPORT LogicalType { GEOMETRY, GEOGRAPHY, VARIANT, + VECTOR, NONE // Not a real logical type; should always be last element }; }; @@ -230,6 +231,9 @@ class PARQUET_EXPORT LogicalType { static std::shared_ptr Float16(); static std::shared_ptr Variant( int8_t specVersion = kVariantSpecVersion); + static std::shared_ptr Vector( + parquet::Type::type element_type, int32_t length, int32_t element_type_length = -1, + std::shared_ptr element_logical_type = nullptr); static std::shared_ptr Geometry(std::string crs = ""); @@ -293,6 +297,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; @@ -509,6 +514,22 @@ class PARQUET_EXPORT VariantLogicalType : public LogicalType { VariantLogicalType() = default; }; +/// \brief Allowed for FIXED_LEN_BYTE_ARRAY primitive nodes only. +class PARQUET_EXPORT VectorLogicalType : public LogicalType { + public: + static std::shared_ptr Make( + parquet::Type::type element_type, int32_t length, int32_t element_type_length = -1, + std::shared_ptr element_logical_type = nullptr); + + parquet::Type::type element_type() const; + int32_t length() const; + int32_t element_type_length() const; + std::shared_ptr element_logical_type() const; + + private: + VectorLogicalType() = default; +}; + /// \brief Allowed for any physical type. class PARQUET_EXPORT NoLogicalType : public LogicalType { public: diff --git a/cpp/src/parquet/types_test.cc b/cpp/src/parquet/types_test.cc index 6c77662d58f7..c0965080df22 100644 --- a/cpp/src/parquet/types_test.cc +++ b/cpp/src/parquet/types_test.cc @@ -20,6 +20,7 @@ #include #include "arrow/util/endian.h" +#include "generated/parquet_types.h" #include "parquet/types.h" namespace parquet { @@ -35,6 +36,38 @@ TEST(TestTypeToString, PhysicalTypes) { ASSERT_STREQ("FIXED_LEN_BYTE_ARRAY", TypeToString(Type::FIXED_LEN_BYTE_ARRAY).c_str()); } +TEST(TestLogicalType, Vector) { + auto vector = LogicalType::Vector(Type::FLOAT, 768); + ASSERT_TRUE(vector->is_vector()); + ASSERT_TRUE(vector->is_applicable(Type::FIXED_LEN_BYTE_ARRAY, 3072)); + ASSERT_FALSE(vector->is_applicable(Type::FIXED_LEN_BYTE_ARRAY, 3071)); + + auto thrift = vector->ToThrift(); + ASSERT_TRUE(thrift.__isset.VECTOR); + ASSERT_EQ(768, thrift.VECTOR.length); + ASSERT_EQ(format::Type::FLOAT, thrift.VECTOR.element_type); + ASSERT_FALSE(thrift.VECTOR.__isset.element_logical_type); + + auto roundtrip = LogicalType::FromThrift(thrift); + ASSERT_TRUE(vector->Equals(*roundtrip)); + + auto zero_length = LogicalType::Vector(Type::FLOAT, 0); + ASSERT_TRUE(zero_length->is_applicable(Type::FIXED_LEN_BYTE_ARRAY, 1)); + ASSERT_FALSE(zero_length->is_applicable(Type::FIXED_LEN_BYTE_ARRAY, 0)); + + auto float16_vector = + LogicalType::Vector(Type::FIXED_LEN_BYTE_ARRAY, 16, + /*element_type_length=*/2, LogicalType::Float16()); + ASSERT_TRUE(float16_vector->is_applicable(Type::FIXED_LEN_BYTE_ARRAY, 32)); + thrift = float16_vector->ToThrift(); + ASSERT_TRUE(thrift.VECTOR.__isset.element_type_length); + ASSERT_EQ(2, thrift.VECTOR.element_type_length); + ASSERT_TRUE(thrift.VECTOR.__isset.element_logical_type); + ASSERT_TRUE(thrift.VECTOR.element_logical_type.__isset.FLOAT16); + roundtrip = LogicalType::FromThrift(thrift); + ASSERT_TRUE(float16_vector->Equals(*roundtrip)); +} + TEST(TestConvertedTypeToString, ConvertedTypes) { ASSERT_STREQ("NONE", ConvertedTypeToString(ConvertedType::NONE).c_str()); ASSERT_STREQ("UTF8", ConvertedTypeToString(ConvertedType::UTF8).c_str()); diff --git a/python/pyarrow/_dataset_parquet.pyx b/python/pyarrow/_dataset_parquet.pyx index 534f7790923a..b1757050db48 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", + "write_fixed_size_list_as_vector", } setters = set() @@ -676,6 +677,9 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): writer_engine_version="V2", use_compliant_nested_type=( self._properties["use_compliant_nested_type"] + ), + write_fixed_size_list_as_vector=( + self._properties["write_fixed_size_list_as_vector"] ) ) @@ -705,6 +709,7 @@ cdef class ParquetFileWriteOptions(FileWriteOptions): coerce_timestamps=None, allow_truncated_timestamps=False, use_compliant_nested_type=True, + write_fixed_size_list_as_vector=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..df8c180d48a2 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=*, + write_fixed_size_list_as_vector=*, ) except * diff --git a/python/pyarrow/_parquet.pyx b/python/pyarrow/_parquet.pyx index 2358a961ebd9..6cf02ffca3ad 100644 --- a/python/pyarrow/_parquet.pyx +++ b/python/pyarrow/_parquet.pyx @@ -1464,6 +1464,7 @@ cdef logical_type_name_from_enum(ParquetLogicalTypeId type_): ParquetLogicalType_JSON: 'JSON', ParquetLogicalType_BSON: 'BSON', ParquetLogicalType_UUID: 'UUID', + ParquetLogicalType_VECTOR: 'VECTOR', ParquetLogicalType_NONE: 'NONE', }.get(type_, 'UNKNOWN') @@ -2281,7 +2282,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, + write_fixed_size_list_as_vector=False) except *: """Arrow writer properties""" cdef: shared_ptr[ArrowWriterProperties] arrow_properties @@ -2322,6 +2324,11 @@ cdef shared_ptr[ArrowWriterProperties] _create_arrow_writer_properties( else: arrow_props.disable_compliant_nested_types() + # write_fixed_size_list_as_vector + + if write_fixed_size_list_as_vector: + arrow_props.enable_fixed_size_list_as_vector() + # writer_engine_version if writer_engine_version == "V1": @@ -2396,6 +2403,7 @@ cdef class ParquetWriter(_Weakrefable): store_decimal_as_integer=False, use_content_defined_chunking=False, write_time_adjusted_to_utc=False, + write_fixed_size_list_as_vector=False, bloom_filter_options=None): cdef: shared_ptr[WriterProperties] properties @@ -2443,6 +2451,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, + write_fixed_size_list_as_vector=write_fixed_size_list_as_vector, ) pool = maybe_unbox_memory_pool(memory_pool) diff --git a/python/pyarrow/includes/libparquet.pxd b/python/pyarrow/includes/libparquet.pxd index df353cc7805f..01b185bbd69a 100644 --- a/python/pyarrow/includes/libparquet.pxd +++ b/python/pyarrow/includes/libparquet.pxd @@ -69,6 +69,7 @@ cdef extern from "parquet/api/schema.h" namespace "parquet" nogil: ParquetLogicalType_UUID" parquet::LogicalType::Type::UUID" ParquetLogicalType_GEOMETRY" parquet::LogicalType::Type::GEOMETRY" ParquetLogicalType_GEOGRAPHY" parquet::LogicalType::Type::GEOGRAPHY" + ParquetLogicalType_VECTOR" parquet::LogicalType::Type::VECTOR" ParquetLogicalType_NONE" parquet::LogicalType::Type::NONE" enum ParquetTimeUnit" parquet::LogicalType::TimeUnit::unit": @@ -529,6 +530,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_fixed_size_list_as_vector() 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 ff880fdcf52c..2820c9bb22af 100644 --- a/python/pyarrow/parquet/core.py +++ b/python/pyarrow/parquet/core.py @@ -884,6 +884,9 @@ def _sanitize_table(table, new_schema, flavor): it will restore the timezone (Parquet only stores the UTC values without timezone), or columns with duration type will be restored from the int64 Parquet column. +write_fixed_size_list_as_vector : bool, default False + If True, write eligible Arrow FixedSizeList columns as Parquet + FIXED_LEN_BYTE_ARRAY columns annotated with VECTOR. write_page_index : bool, default False Whether to write a page index in general for all columns. Writing statistics to the page index disables the old method of writing @@ -1077,6 +1080,7 @@ def __init__(self, where, schema, filesystem=None, write_batch_size=None, dictionary_pagesize_limit=None, store_schema=True, + write_fixed_size_list_as_vector=False, write_page_index=False, write_page_checksum=False, sorting_columns=None, @@ -1132,6 +1136,7 @@ def __init__(self, where, schema, filesystem=None, write_batch_size=write_batch_size, dictionary_pagesize_limit=dictionary_pagesize_limit, store_schema=store_schema, + write_fixed_size_list_as_vector=write_fixed_size_list_as_vector, write_page_index=write_page_index, write_page_checksum=write_page_checksum, sorting_columns=sorting_columns, @@ -2010,6 +2015,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_batch_size=None, dictionary_pagesize_limit=None, store_schema=True, + write_fixed_size_list_as_vector=False, write_page_index=False, write_page_checksum=False, sorting_columns=None, @@ -2044,6 +2050,7 @@ def write_table(table, where, row_group_size=None, version='2.6', write_batch_size=write_batch_size, dictionary_pagesize_limit=dictionary_pagesize_limit, store_schema=store_schema, + write_fixed_size_list_as_vector=write_fixed_size_list_as_vector, write_page_index=write_page_index, write_page_checksum=write_page_checksum, sorting_columns=sorting_columns, diff --git a/python/pyarrow/tests/parquet/test_data_types.py b/python/pyarrow/tests/parquet/test_data_types.py index c546bc1532ac..b3e8bf0fd102 100644 --- a/python/pyarrow/tests/parquet/test_data_types.py +++ b/python/pyarrow/tests/parquet/test_data_types.py @@ -604,6 +604,23 @@ def test_uuid_extension_type(): store_schema=False) +def test_fixed_size_list_as_vector(tempdir): + typ = pa.list_(pa.field("element", pa.float32(), nullable=False), 3) + arr = pa.array([[1, 2, 3], None, [4, 5, 6]], type=typ) + table = pa.table({"embedding": arr}) + path = tempdir / "vector.parquet" + + pq.write_table(table, path, write_fixed_size_list_as_vector=True) + + metadata = pq.read_metadata(path) + col = metadata.schema[0] + assert col.physical_type == "FIXED_LEN_BYTE_ARRAY" + assert col.logical_type.type == "VECTOR" + + result = pq.read_table(path) + assert result.equals(table) + + def test_undefined_logical_type(parquet_test_datadir): test_file = f"{parquet_test_datadir}/unknown-logical-type.parquet"