-
Notifications
You must be signed in to change notification settings - Fork 31
feat: support fixed-size arrays in data readers #726
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
f629311
02501f3
34b889b
17a7638
ead76cf
902ce97
3e92342
486a112
db76fdc
204b39a
52ccfb7
355cb08
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,13 +16,181 @@ | |
| #include "parquet/arrow_column.h" | ||
|
|
||
| #include <arrow/array/array_binary.h> | ||
| #include <arrow/array/array_nested.h> | ||
| #include <arrow/type.h> | ||
|
|
||
| #include "neug/common/columns/columns_utils.h" | ||
| #include "neug/common/columns/value_columns.h" | ||
| #include "neug/common/types/value.h" | ||
| #include "neug/utils/exception/exception.h" | ||
|
|
||
| namespace neug { | ||
|
|
||
| static int64_t arrow_time_to_milliseconds(int64_t value, | ||
| arrow::TimeUnit::type unit) { | ||
| switch (unit) { | ||
| case arrow::TimeUnit::SECOND: | ||
| return value * Interval::MSECS_PER_SEC; | ||
| case arrow::TimeUnit::MILLI: | ||
| return value; | ||
| case arrow::TimeUnit::MICRO: | ||
| return value / Interval::MICROS_PER_MSEC; | ||
| case arrow::TimeUnit::NANO: | ||
| return value / (Interval::MICROS_PER_MSEC * 1000); | ||
| } | ||
| THROW_NOT_SUPPORTED_EXCEPTION("Unsupported Arrow time unit"); | ||
| } | ||
|
Comment on lines
+29
to
+42
|
||
|
|
||
| // Keep this mapping in sync with arrow_value_at when adding Arrow types. | ||
| static DataType arrow_type_to_neug_type(const arrow::DataType& type) { | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. [Low] Both functions (~50 lines each) manually map every Arrow type to a NeuG type or Value. Adding a new Arrow type requires updating both functions, and it's easy to forget one. Suggestion: Consider adding a comment like
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
all fixed in 486a112 |
||
| switch (type.id()) { | ||
| case arrow::Type::BOOL: | ||
| return DataType::BOOLEAN; | ||
| case arrow::Type::INT32: | ||
| return DataType::INT32; | ||
| case arrow::Type::INT64: | ||
| return DataType::INT64; | ||
| case arrow::Type::UINT32: | ||
| return DataType::UINT32; | ||
| case arrow::Type::UINT64: | ||
| return DataType::UINT64; | ||
| case arrow::Type::FLOAT: | ||
| return DataType::FLOAT; | ||
| case arrow::Type::DOUBLE: | ||
| return DataType::DOUBLE; | ||
| case arrow::Type::STRING: | ||
| case arrow::Type::LARGE_STRING: | ||
| return DataType::VARCHAR; | ||
| case arrow::Type::DATE32: | ||
| case arrow::Type::DATE64: | ||
| return DataType::DATE; | ||
| case arrow::Type::TIMESTAMP: | ||
| return DataType::TIMESTAMP_MS; | ||
| case arrow::Type::DURATION: | ||
| return DataType::INTERVAL; | ||
|
shirly121 marked this conversation as resolved.
shirly121 marked this conversation as resolved.
|
||
| case arrow::Type::FIXED_SIZE_LIST: { | ||
| const auto& array_type = static_cast<const arrow::FixedSizeListType&>(type); | ||
| return DataType::Array(arrow_type_to_neug_type(*array_type.value_type()), | ||
| array_type.list_size()); | ||
| } | ||
| case arrow::Type::LIST: | ||
| case arrow::Type::LARGE_LIST: | ||
| THROW_NOT_SUPPORTED_EXCEPTION( | ||
| "Parquet LIST is not supported as ARRAY. Specify a fixed-size Arrow " | ||
| "list / NeuG ARRAY type instead."); | ||
| default: | ||
| THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + type.ToString()); | ||
| } | ||
| } | ||
|
|
||
| static Value arrow_value_at(const arrow::Array& array, int64_t index, | ||
| const DataType& type) { | ||
| if (array.IsNull(index)) { | ||
| return Value(type); | ||
| } | ||
| switch (array.type_id()) { | ||
| case arrow::Type::BOOL: | ||
| return Value::BOOLEAN( | ||
| static_cast<const arrow::BooleanArray&>(array).Value(index)); | ||
| case arrow::Type::INT32: | ||
| return Value::INT32( | ||
| static_cast<const arrow::Int32Array&>(array).Value(index)); | ||
| case arrow::Type::INT64: | ||
| return Value::INT64( | ||
| static_cast<const arrow::Int64Array&>(array).Value(index)); | ||
| case arrow::Type::UINT32: | ||
| return Value::UINT32( | ||
| static_cast<const arrow::UInt32Array&>(array).Value(index)); | ||
| case arrow::Type::UINT64: | ||
| return Value::UINT64( | ||
| static_cast<const arrow::UInt64Array&>(array).Value(index)); | ||
| case arrow::Type::FLOAT: | ||
| return Value::FLOAT( | ||
| static_cast<const arrow::FloatArray&>(array).Value(index)); | ||
| case arrow::Type::DOUBLE: | ||
| return Value::DOUBLE( | ||
| static_cast<const arrow::DoubleArray&>(array).Value(index)); | ||
| case arrow::Type::STRING: | ||
| return Value::STRING(std::string( | ||
| static_cast<const arrow::StringArray&>(array).GetView(index))); | ||
| case arrow::Type::LARGE_STRING: | ||
| return Value::STRING(std::string( | ||
| static_cast<const arrow::LargeStringArray&>(array).GetView(index))); | ||
| case arrow::Type::DATE32: { | ||
| Date date; | ||
| date.from_num_days( | ||
| static_cast<const arrow::Date32Array&>(array).Value(index)); | ||
| return Value::DATE(date); | ||
| } | ||
| case arrow::Type::DATE64: { | ||
| constexpr int64_t MILLIS_PER_DAY = 24LL * 60 * 60 * 1000; | ||
| Date date; | ||
| date.from_num_days(static_cast<int32_t>( | ||
| static_cast<const arrow::Date64Array&>(array).Value(index) / | ||
| MILLIS_PER_DAY)); | ||
| return Value::DATE(date); | ||
| } | ||
| case arrow::Type::TIMESTAMP: { | ||
| const auto value = | ||
| static_cast<const arrow::TimestampArray&>(array).Value(index); | ||
| const auto& timestamp_type = | ||
| static_cast<const arrow::TimestampType&>(*array.type()); | ||
| return Value::TIMESTAMPMS( | ||
| DateTime(arrow_time_to_milliseconds(value, timestamp_type.unit()))); | ||
| } | ||
|
shirly121 marked this conversation as resolved.
|
||
| case arrow::Type::DURATION: { | ||
| const auto value = | ||
| static_cast<const arrow::DurationArray&>(array).Value(index); | ||
| const auto& duration_type = | ||
| static_cast<const arrow::DurationType&>(*array.type()); | ||
| Interval interval; | ||
| interval.from_mill_seconds( | ||
| arrow_time_to_milliseconds(value, duration_type.unit())); | ||
| return Value::INTERVAL(interval); | ||
| } | ||
| case arrow::Type::FIXED_SIZE_LIST: { | ||
| const auto& list = static_cast<const arrow::FixedSizeListArray&>(array); | ||
| const auto& array_type = | ||
| static_cast<const arrow::FixedSizeListType&>(*array.type()); | ||
| const auto child_type = ArrayType::GetChildType(type); | ||
| std::vector<Value> values; | ||
| values.reserve(array_type.list_size()); | ||
| const auto offset = list.value_offset(index); | ||
| for (int32_t i = 0; i < array_type.list_size(); ++i) { | ||
| values.push_back(arrow_value_at(*list.values(), offset + i, child_type)); | ||
| } | ||
| return Value::ARRAY(type, std::move(values)); | ||
| } | ||
| default: | ||
| THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + | ||
| array.type()->ToString()); | ||
| } | ||
| } | ||
|
|
||
| static std::shared_ptr<IContextColumn> convert_fixed_size_list_arrays( | ||
| const std::vector<std::shared_ptr<arrow::Array>>& arrays) { | ||
| const auto type = arrow_type_to_neug_type(*arrays.front()->type()); | ||
| auto builder = ColumnsUtils::create_builder(type); | ||
|
Comment on lines
+170
to
+173
|
||
| size_t size = 0; | ||
| for (const auto& array : arrays) { | ||
| size += array->length(); | ||
|
Comment on lines
+170
to
+176
|
||
| } | ||
| builder->reserve(size); | ||
| for (const auto& array : arrays) { | ||
| if (!array->type()->Equals(*arrays.front()->type())) { | ||
| THROW_SCHEMA_MISMATCH("Parquet ARRAY chunks have different types"); | ||
| } | ||
| for (int64_t i = 0; i < array->length(); ++i) { | ||
| if (array->IsNull(i)) { | ||
| builder->push_back_null(); | ||
| continue; | ||
| } | ||
| builder->push_back_elem(arrow_value_at(*array, i, type)); | ||
| } | ||
| } | ||
| return builder->finish(); | ||
| } | ||
|
|
||
| /// Convert numeric arrow arrays directly to ValueColumn<CppT>. | ||
| template <typename ArrowArrayT, typename CppT> | ||
| static std::shared_ptr<IContextColumn> convert_numeric_arrays( | ||
|
|
@@ -89,7 +257,11 @@ static std::shared_ptr<IContextColumn> convert_date64_arrays( | |
| if (typed->IsNull(j)) { | ||
| builder.push_back_null(); | ||
| } else { | ||
| builder.push_back_opt(Date(typed->Value(j))); | ||
| constexpr int64_t MILLIS_PER_DAY = 24LL * 60 * 60 * 1000; | ||
| Date date; | ||
| date.from_num_days( | ||
| static_cast<int32_t>(typed->Value(j) / MILLIS_PER_DAY)); | ||
| builder.push_back_opt(date); | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -106,7 +278,10 @@ static std::shared_ptr<IContextColumn> convert_timestamp_arrays( | |
| if (typed->IsNull(j)) { | ||
| builder.push_back_null(); | ||
| } else { | ||
| builder.push_back_opt(DateTime(typed->Value(j))); | ||
| const auto& timestamp_type = | ||
| static_cast<const arrow::TimestampType&>(*typed->type()); | ||
| builder.push_back_opt(DateTime(arrow_time_to_milliseconds( | ||
| typed->Value(j), timestamp_type.unit()))); | ||
| } | ||
| } | ||
| } | ||
|
|
@@ -144,6 +319,13 @@ std::shared_ptr<IContextColumn> arrow_arrays_to_value_column( | |
| return convert_date64_arrays(arrays); | ||
| case arrow::Type::TIMESTAMP: | ||
| return convert_timestamp_arrays(arrays); | ||
| case arrow::Type::FIXED_SIZE_LIST: | ||
| return convert_fixed_size_list_arrays(arrays); | ||
| case arrow::Type::LIST: | ||
| case arrow::Type::LARGE_LIST: | ||
| THROW_NOT_SUPPORTED_EXCEPTION( | ||
| "Parquet LIST columns are not supported; use a fixed-size ARRAY " | ||
| "type such as FLOAT[3]"); | ||
| default: | ||
| THROW_NOT_SUPPORTED_EXCEPTION("Unsupported arrow type: " + | ||
| arrow_type->ToString()); | ||
|
|
||
Uh oh!
There was an error while loading. Please reload this page.