Initial commit

Based on dc4ba6a

1 files changed, 102 insertions, 0 deletions

diff --git a/lib/libtmi8/src/kv6_parquet.cpp b/lib/libtmi8/src/kv6_parquet.cpp
new file mode 100644
index 0000000..ca70b7f
--- /dev/null
+++ b/lib/libtmi8/src/kv6_parquet.cpp
@@ -0,0 +1,102 @@
+// vim:set sw=2 ts=2 sts et:
+
+#include <tmi8/kv6_parquet.hpp>
+
+ParquetBuilder::ParquetBuilder() {
+  std::shared_ptr<arrow::Field> field_type, field_data_owner_code, field_line_planning_number, field_operating_day,
+                                field_journey_number, field_reinforcement_number, field_timestamp, field_source,
+                                field_punctuality, field_user_stop_code, field_passage_sequence_number,
+                                field_vehicle_number, field_block_code, field_wheelchair_accessible,
+                                field_number_of_coaches, field_rd_y, field_rd_x, field_distance_since_last_user_stop;
+  field_type                          = arrow::field("type", arrow::utf8());
+  field_data_owner_code               = arrow::field("data_owner_code", arrow::utf8());
+  field_line_planning_number          = arrow::field("line_planning_number", arrow::utf8());
+  field_operating_day                 = arrow::field("operating_day", arrow::date32());
+  field_journey_number                = arrow::field("journey_number", arrow::uint32());
+  field_reinforcement_number          = arrow::field("reinforcement_number", arrow::uint8());
+  field_timestamp                     = arrow::field("timestamp", arrow::timestamp(arrow::TimeUnit::SECOND));
+  field_source                        = arrow::field("source", arrow::utf8());
+  field_punctuality                   = arrow::field("punctuality", arrow::int16());
+  field_user_stop_code                = arrow::field("user_stop_code", arrow::utf8());
+  field_passage_sequence_number       = arrow::field("passage_sequence_number", arrow::uint16());
+  field_vehicle_number                = arrow::field("vehicle_number", arrow::uint32());
+  field_block_code                    = arrow::field("block_code", arrow::uint32());
+  field_wheelchair_accessible         = arrow::field("wheelchair_accessible", arrow::utf8());
+  field_number_of_coaches             = arrow::field("number_of_coaches", arrow::uint8());
+  field_rd_y                          = arrow::field("rd_y", arrow::int32());
+  field_rd_x                          = arrow::field("rd_x", arrow::int32());
+  field_distance_since_last_user_stop = arrow::field("distance_since_last_user_stop", arrow::uint32());
+
+  schema = arrow::schema({ field_type, field_data_owner_code, field_line_planning_number,
+                           field_operating_day, field_journey_number,
+                           field_reinforcement_number, field_timestamp, field_source,
+                           field_punctuality, field_user_stop_code,
+                           field_passage_sequence_number, field_vehicle_number,
+                           field_block_code, field_wheelchair_accessible,
+                           field_number_of_coaches, field_rd_y, field_rd_x,
+                           field_distance_since_last_user_stop });
+}
+
+arrow::Result<std::shared_ptr<arrow::Table>> ParquetBuilder::getTable() {
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> types,                          types.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> data_owner_codes,               data_owner_codes.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> line_planning_numbers,          line_planning_numbers.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> operating_days,                 operating_days.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> journey_numbers,                journey_numbers.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> reinforcement_numbers,          reinforcement_numbers.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> timestamps,                     timestamps.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> sources,                        sources.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> punctualities,                  punctualities.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> user_stop_codes,                user_stop_codes.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> passage_sequence_numbers,       passage_sequence_numbers.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> vehicle_numbers,                vehicle_numbers.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> block_codes,                    block_codes.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> wheelchair_accessibles,         wheelchair_accessibles.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> number_of_coaches,              number_of_coaches.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> rd_ys,                          rd_ys.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> rd_xs,                          rd_xs.Finish());
+  ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::Array> distance_since_last_user_stops, distance_since_last_user_stops.Finish());
+
+  std::vector<std::shared_ptr<arrow::Array>> columns = { types, data_owner_codes, line_planning_numbers, operating_days,
+                                                         journey_numbers, reinforcement_numbers, timestamps, sources,
+                                                         punctualities, user_stop_codes, passage_sequence_numbers,
+                                                         vehicle_numbers, block_codes, wheelchair_accessibles,
+                                                         number_of_coaches, rd_ys, rd_xs,
+                                                         distance_since_last_user_stops };
+  return arrow::Result(arrow::Table::Make(schema, columns));
+}
+
+arrow::Status writeArrowRecordsAsParquetFile(arrow::RecordBatchReader &rbr, std::filesystem::path filename) {
+  std::shared_ptr<parquet::WriterProperties> props = parquet::WriterProperties::Builder()
+    .compression(arrow::Compression::ZSTD)
+    ->created_by("oeuf-libtmi8")
+    ->version(parquet::ParquetVersion::PARQUET_2_6)
+    ->data_page_version(parquet::ParquetDataPageVersion::V2)
+    ->max_row_group_length(MAX_PARQUET_CHUNK)
+    ->build();
+
+  std::shared_ptr<parquet::ArrowWriterProperties> arrow_props = parquet::ArrowWriterProperties::Builder()
+    .store_schema()->build();
+
+  std::shared_ptr<arrow::io::FileOutputStream> out_file;
+  std::string filename_str = filename;
+  ARROW_ASSIGN_OR_RAISE(out_file, arrow::io::FileOutputStream::Open(filename_str + ".part"));
+
+  ARROW_ASSIGN_OR_RAISE(auto writer,
+    parquet::arrow::FileWriter::Open(*rbr.schema(), arrow::default_memory_pool(), out_file, props, arrow_props));
+  for (const auto &batchr : rbr) {
+    ARROW_ASSIGN_OR_RAISE(auto batch, batchr);
+    ARROW_RETURN_NOT_OK(writer->WriteRecordBatch(*batch));
+  }
+  ARROW_RETURN_NOT_OK(writer->Close());
+  ARROW_RETURN_NOT_OK(out_file->Close());
+
+  std::filesystem::rename(filename_str + ".part", filename);
+
+  return arrow::Status::OK();
+}
+
+arrow::Status writeArrowTableAsParquetFile(const arrow::Table &table, std::filesystem::path filename) {
+  auto tbr = arrow::TableBatchReader(table);
+  return writeArrowRecordsAsParquetFile(tbr, filename);
+}