Line data Source code
1 : /****************************************************************************/
2 : // Eclipse SUMO, Simulation of Urban MObility; see https://eclipse.dev/sumo
3 : // Copyright (C) 2012-2026 German Aerospace Center (DLR) and others.
4 : // This program and the accompanying materials are made available under the
5 : // terms of the Eclipse Public License 2.0 which is available at
6 : // https://www.eclipse.org/legal/epl-2.0/
7 : // This Source Code may also be made available under the following Secondary
8 : // Licenses when the conditions for such availability set forth in the Eclipse
9 : // Public License 2.0 are satisfied: GNU General Public License, version 2
10 : // or later which is available at
11 : // https://www.gnu.org/licenses/old-licenses/gpl-2.0-standalone.html
12 : // SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-or-later
13 : /****************************************************************************/
14 : /// @file ParquetFormatter.cpp
15 : /// @author Michael Behrisch
16 : /// @date 2025-06-17
17 : ///
18 : // An output formatter for Parquet files
19 : /****************************************************************************/
20 : #include <config.h>
21 :
22 : #ifdef _MSC_VER
23 : #pragma warning(push)
24 : /* Disable warning about unused parameters */
25 : #pragma warning(disable: 4100)
26 : /* Disable warning about hidden function arrow::io::Writable::Write */
27 : #pragma warning(disable: 4266)
28 : /* Disable warning about padded memory layout */
29 : #pragma warning(disable: 4324)
30 : /* Disable warning about this in initializers */
31 : #pragma warning(disable: 4355)
32 : /* Disable warning about changed memory layout due to virtual base class */
33 : #pragma warning(disable: 4435)
34 : /* Disable warning about declaration hiding class member */
35 : #pragma warning(disable: 4458)
36 : /* Disable warning about implicit conversion of int to bool */
37 : #pragma warning(disable: 4800)
38 : #endif
39 : #include <arrow/api.h>
40 : #include <arrow/io/api.h>
41 : #include <parquet/arrow/writer.h>
42 : #ifdef _MSC_VER
43 : #pragma warning(pop)
44 : #endif
45 :
46 : #include <utils/common/MsgHandler.h>
47 : #include <utils/common/ToString.h>
48 : #include "ParquetFormatter.h"
49 :
50 :
51 : // ===========================================================================
52 : // helper class definitions
53 : // ===========================================================================
54 : #ifdef _MSC_VER
55 : #pragma warning(push)
56 : /* Disable warning about hidden function arrow::io::Writable::Write */
57 : #pragma warning(disable: 4266)
58 : #endif
59 : class ArrowOStreamWrapper : public arrow::io::OutputStream {
60 : public:
61 : ArrowOStreamWrapper(std::ostream& out)
62 59 : : myOStream(out), myAmOpen(true) {}
63 :
64 0 : arrow::Status Close() override {
65 0 : myAmOpen = false;
66 0 : return arrow::Status::OK();
67 : }
68 :
69 0 : arrow::Status Flush() override {
70 0 : myOStream.flush();
71 0 : return arrow::Status::OK();
72 : }
73 :
74 4417 : arrow::Result<int64_t> Tell() const override {
75 4417 : return myOStream.tellp();
76 : }
77 :
78 0 : bool closed() const override {
79 0 : return !myAmOpen;
80 : }
81 :
82 2780 : arrow::Status Write(const void* data, int64_t nbytes) override {
83 2780 : if (!myAmOpen) {
84 : return arrow::Status::IOError("Write on closed stream");
85 : }
86 2780 : myOStream.write(reinterpret_cast<const char*>(data), nbytes);
87 2780 : if (!myOStream) {
88 : return arrow::Status::IOError("Failed to write to ostream");
89 : }
90 : return arrow::Status::OK();
91 : }
92 :
93 : private:
94 : std::ostream& myOStream;
95 : bool myAmOpen;
96 : };
97 : #ifdef _MSC_VER
98 : #pragma warning(pop)
99 : #endif
100 :
101 :
102 : // ===========================================================================
103 : // ParquetFormatter::Impl definition
104 : // ===========================================================================
105 : struct ParquetFormatter::Impl {
106 59 : Impl(const std::string& columnNames, const int batchSize)
107 177 : : myHeaderFormat(columnNames), myBatchSize(batchSize) {}
108 :
109 : /// @brief the format to use for the column names
110 : const std::string myHeaderFormat;
111 :
112 : /// @brief the column names if we write the full name
113 : std::vector<std::string> myFullHeader;
114 :
115 : /// @brief the compression to use
116 : parquet::Compression::type myCompression = parquet::Compression::UNCOMPRESSED;
117 :
118 : /// @brief the number of rows to write per batch
119 : const int myBatchSize;
120 :
121 : /// @brief the table schema
122 : std::shared_ptr<arrow::Schema> mySchema = arrow::schema({});
123 :
124 : /// @brief the output stream writer
125 : std::unique_ptr<parquet::arrow::FileWriter> myParquetWriter;
126 :
127 : /// @brief the content array builders for the table
128 : std::vector<std::shared_ptr<arrow::ArrayBuilder> > myBuilders;
129 :
130 : /// @brief The name and number of attributes in the currently open XML elements
131 : std::vector<std::pair<const std::string, int> > myXMLStack;
132 :
133 : /// @brief the current attribute / column values
134 : std::vector<std::shared_ptr<arrow::Scalar> > myValues;
135 :
136 : /// @brief the maximum depth of the XML hierarchy
137 : int myMaxDepth = 1000;
138 :
139 : /// @brief whether the schema has been constructed completely
140 : bool myWroteHeader = false;
141 :
142 : /// @brief whether the columns should be checked for completeness
143 : bool myCheckColumns = false;
144 :
145 : /// @brief whether there is still unwritten data
146 : bool myNeedsWrite = false;
147 :
148 : /// @brief whether any root attribute have been encountered
149 : bool myHaveRootAttrs = false;
150 :
151 : /// @brief the attributes which are expected for a complete row (including null values)
152 : SumoXMLAttrMask myExpectedAttrs;
153 :
154 : /// @brief the attributes already seen (including null values)
155 : SumoXMLAttrMask mySeenAttrs;
156 :
157 171021 : void checkAttr(const SumoXMLAttr attr) {
158 171021 : if (myCheckColumns && myMaxDepth == (int)myXMLStack.size()) {
159 8278 : mySeenAttrs.set(attr);
160 8278 : if (!myExpectedAttrs.test(attr)) {
161 0 : throw ProcessError(TLF("Unexpected attribute '%', this file format does not support Parquet output yet.", toString(attr)));
162 : }
163 : }
164 171021 : }
165 :
166 : template <class ATTR_TYPE, class BUILDER>
167 225353 : void checkBuilder(const ATTR_TYPE& attr, const std::shared_ptr<arrow::DataType>& (*dataType)()) {
168 225353 : myNeedsWrite = true;
169 225353 : if (!myWroteHeader) {
170 157920 : std::string fieldName = toString(attr);
171 : std::string prefix;
172 607800 : for (const auto& entry : myXMLStack) {
173 799444 : prefix += entry.first + "_";
174 : }
175 : const std::string fullHeaderName = prefix + fieldName;
176 208078 : if (myHeaderFormat == "none") {
177 : fieldName = "";
178 208078 : } else if (myHeaderFormat != "plain") {
179 624234 : fieldName = myXMLStack.back().first + "_" + fieldName;
180 : }
181 208078 : const auto colIt = std::find(myFullHeader.begin(), myFullHeader.end(), fullHeaderName);
182 208078 : if (colIt == myFullHeader.end()) {
183 3392 : mySchema = *mySchema->AddField(mySchema->num_fields(), arrow::field(fieldName, dataType()));
184 : auto builder = std::make_shared<BUILDER>();
185 848 : if (!myBuilders.empty()) {
186 791 : if (myBuilders.back()->length() > 0) {
187 283 : PARQUET_THROW_NOT_OK(builder->AppendNulls(myBuilders.back()->length()));
188 : }
189 1024 : while (myValues.size() < myBuilders.size()) {
190 466 : myValues.push_back(nullptr);
191 : }
192 : }
193 848 : myBuilders.push_back(builder);
194 848 : myFullHeader.emplace_back(fullHeaderName);
195 : } else {
196 207230 : myValues.resize(std::distance(myFullHeader.begin(), colIt));
197 : }
198 : }
199 225353 : }
200 : };
201 :
202 :
203 : // ===========================================================================
204 : // member method definitions
205 : // ===========================================================================
206 59 : ParquetFormatter::ParquetFormatter(const std::string& columnNames, const std::string& compression, const int batchSize)
207 59 : : OutputFormatter(OutputFormatterType::PARQUET), myImpl(std::make_unique<Impl>(columnNames, batchSize)) {
208 59 : if (compression == "snappy") {
209 0 : myImpl->myCompression = parquet::Compression::SNAPPY;
210 59 : } else if (compression == "gzip") {
211 0 : myImpl->myCompression = parquet::Compression::GZIP;
212 59 : } else if (compression == "brotli") {
213 0 : myImpl->myCompression = parquet::Compression::BROTLI;
214 59 : } else if (compression == "zstd") {
215 0 : myImpl->myCompression = parquet::Compression::ZSTD;
216 59 : } else if (compression == "lz4") {
217 0 : myImpl->myCompression = parquet::Compression::LZ4;
218 59 : } else if (compression == "bz2") {
219 0 : myImpl->myCompression = parquet::Compression::BZ2;
220 59 : } else if (compression != "" && compression != "uncompressed") {
221 0 : WRITE_ERRORF("Unknown compression: %", compression);
222 : }
223 59 : if (!arrow::util::Codec::IsAvailable(myImpl->myCompression)) {
224 0 : WRITE_WARNINGF("Compression '%' not available, falling back to uncompressed.", compression);
225 0 : myImpl->myCompression = parquet::Compression::UNCOMPRESSED;
226 : }
227 59 : }
228 :
229 :
230 118 : ParquetFormatter::~ParquetFormatter() = default;
231 :
232 :
233 : bool
234 103 : ParquetFormatter::writeXMLHeader(std::ostream& into, const std::string& rootElement,
235 : const std::map<SumoXMLAttr, std::string>& attrs, bool /* writeMetadata */,
236 : bool /* includeConfig */) {
237 103 : if (attrs.size() > 2) {
238 13 : myImpl->myHaveRootAttrs = true;
239 13 : openTag(into, rootElement);
240 52 : for (const auto& a : attrs) {
241 39 : if (a.first != SUMO_ATTR_XMLNS && a.first != SUMO_ATTR_SCHEMA_LOCATION) {
242 13 : writeAttr(into, a.first, a.second, false, false);
243 : }
244 : }
245 : return true;
246 : }
247 : return false;
248 : }
249 :
250 :
251 : void
252 6616 : ParquetFormatter::openTag(std::ostream& /* into */, const std::string& xmlElement) {
253 6616 : myImpl->myXMLStack.push_back({xmlElement, (int)myImpl->myValues.size()});
254 6616 : }
255 :
256 :
257 : void
258 13576 : ParquetFormatter::openTag(std::ostream& /* into */, const SumoXMLTag& xmlElement) {
259 13576 : myImpl->myXMLStack.push_back({toString(xmlElement), (int)myImpl->myValues.size()});
260 13576 : }
261 :
262 :
263 : bool
264 20251 : ParquetFormatter::closeTag(std::ostream& into, const std::string& /* comment */) {
265 20251 : if (myImpl->myMaxDepth == 0) {
266 : // the auto detection case: the first closed tag determines the depth
267 0 : myImpl->myMaxDepth = (int)myImpl->myXMLStack.size();
268 : }
269 20251 : const bool eof = myImpl->myXMLStack.empty() || (myImpl->myHaveRootAttrs && myImpl->myXMLStack.size() == 1);
270 20251 : if ((myImpl->myMaxDepth == (int)myImpl->myXMLStack.size() || eof) && !myImpl->myWroteHeader) {
271 : // we are at the correct depth or the document has ended (XML stack is empty)
272 : // so we should initialize the writer with the schema (if not done yet)
273 59 : if (!myImpl->myCheckColumns) {
274 80 : WRITE_WARNING("Column based formats are still experimental. Autodetection only works for homogeneous output.");
275 : }
276 59 : bool full = myImpl->myHeaderFormat == "full";
277 59 : if (myImpl->myHeaderFormat == "auto") {
278 : std::set<std::string> uniq;
279 905 : for (const auto& field : myImpl->mySchema->fields()) {
280 : const auto result = uniq.insert(field->name());
281 847 : if (!result.second) {
282 : full = true;
283 1 : break;
284 : }
285 : }
286 : }
287 59 : if (full) {
288 : arrow::FieldVector new_fields;
289 41 : for (const auto& field : myImpl->mySchema->fields()) {
290 80 : new_fields.push_back(field->WithName(myImpl->myFullHeader[new_fields.size()]));
291 : }
292 3 : myImpl->mySchema = arrow::schema(std::move(new_fields), myImpl->mySchema->metadata());
293 1 : }
294 : auto arrow_stream = std::make_shared<ArrowOStreamWrapper>(into);
295 59 : std::shared_ptr<parquet::WriterProperties> props = parquet::WriterProperties::Builder().compression(myImpl->myCompression)->build();
296 236 : myImpl->myParquetWriter = *parquet::arrow::FileWriter::Open(*myImpl->mySchema, arrow::default_memory_pool(), arrow_stream, props);
297 59 : myImpl->myWroteHeader = true;
298 : }
299 : bool writeBatch = false;
300 20251 : if (myImpl->myNeedsWrite) {
301 16907 : if (myImpl->myCheckColumns && (int)myImpl->myXMLStack.size() == myImpl->myMaxDepth && myImpl->myExpectedAttrs != myImpl->mySeenAttrs) {
302 1552 : for (int i = 0; i < (int)myImpl->myExpectedAttrs.size(); ++i) {
303 1536 : if (myImpl->myExpectedAttrs.test(i) && !myImpl->mySeenAttrs.test(i)) {
304 36 : WRITE_ERRORF("Incomplete attribute set, '%' is missing. This file format does not support Parquet output yet.",
305 : toString((SumoXMLAttr)i));
306 : }
307 : }
308 : }
309 : int index = 0;
310 321673 : for (auto& builder : myImpl->myBuilders) {
311 305612 : const auto val = index < (int)myImpl->myValues.size() ? myImpl->myValues[index] : nullptr;
312 305612 : arrow::Status s = val == nullptr ? builder->AppendNull() : builder->AppendScalar(*val);
313 305612 : if (!s.ok()) {
314 0 : throw ProcessError(TLF("Error writing attribute '%' (index: %, value: '%'): %",
315 0 : myImpl->myFullHeader[index], index, val == nullptr ? "nullptr" : val->ToString(), s.ToString()));
316 : }
317 305612 : index++;
318 : }
319 16061 : writeBatch = myImpl->myWroteHeader && myImpl->myBuilders.back()->length() >= myImpl->myBatchSize;
320 : myImpl->mySeenAttrs.reset();
321 16061 : myImpl->myNeedsWrite = false;
322 : }
323 20251 : if (writeBatch || (eof && !myImpl->myBuilders.empty())) {
324 : std::vector<std::shared_ptr<arrow::Array> > data;
325 1221 : for (auto& builder : myImpl->myBuilders) {
326 1151 : std::shared_ptr<arrow::Array> column;
327 1151 : PARQUET_THROW_NOT_OK(builder->Finish(&column));
328 1151 : data.push_back(column);
329 : // builder.reset();
330 : }
331 210 : auto batch = arrow::RecordBatch::Make(myImpl->mySchema, data.back()->length(), data);
332 70 : PARQUET_THROW_NOT_OK(myImpl->myParquetWriter->WriteRecordBatch(*batch));
333 70 : }
334 20251 : if (!myImpl->myXMLStack.empty()) {
335 20192 : if ((int)myImpl->myValues.size() > myImpl->myXMLStack.back().second) {
336 19992 : myImpl->myValues.resize(myImpl->myXMLStack.back().second);
337 : }
338 : myImpl->myXMLStack.pop_back();
339 20192 : return true;
340 : }
341 : return false;
342 : }
343 :
344 :
345 : void
346 42673 : ParquetFormatter::writeAttr(std::ostream& into, const SumoXMLAttr attr, const double& val, const bool isNull, const bool /* escape */) {
347 42673 : myImpl->checkAttr(attr);
348 42673 : if (attr == SUMO_ATTR_X || attr == SUMO_ATTR_Y || into.precision() > 2) {
349 4242 : myImpl->checkBuilder<SumoXMLAttr, arrow::DoubleBuilder>(attr, arrow::float64);
350 8484 : myImpl->myValues.push_back(isNull ? nullptr : std::make_shared<arrow::DoubleScalar>(val));
351 : } else {
352 38431 : myImpl->checkBuilder<SumoXMLAttr, arrow::FloatBuilder>(attr, arrow::float32);
353 76862 : myImpl->myValues.push_back(isNull ? nullptr : std::make_shared<arrow::FloatScalar>((float)val));
354 : }
355 42673 : }
356 :
357 :
358 : void
359 15884 : ParquetFormatter::writeAttr(std::ostream& /* into */, const SumoXMLAttr attr, const int& val, const bool isNull, const bool /* escape */) {
360 15884 : myImpl->checkAttr(attr);
361 15884 : myImpl->checkBuilder<SumoXMLAttr, arrow::Int32Builder>(attr, arrow::int32);
362 15884 : myImpl->myValues.push_back(isNull ? nullptr : std::make_shared<arrow::Int32Scalar>(val));
363 15884 : }
364 :
365 :
366 : void
367 7105 : ParquetFormatter::writeAttr(std::ostream& into, const std::string& attr, const double& val, const bool isNull, const bool /* escape */) {
368 : assert(!myImpl->myCheckColumns);
369 7105 : if (into.precision() > 2) {
370 0 : myImpl->checkBuilder<std::string, arrow::DoubleBuilder>(attr, arrow::float64);
371 0 : myImpl->myValues.push_back(isNull ? nullptr : std::make_shared<arrow::DoubleScalar>(val));
372 : } else {
373 7105 : myImpl->checkBuilder<std::string, arrow::FloatBuilder>(attr, arrow::float32);
374 14210 : myImpl->myValues.push_back(isNull ? nullptr : std::make_shared<arrow::FloatScalar>((float)val));
375 : }
376 7105 : }
377 :
378 :
379 : void
380 23108 : ParquetFormatter::writeAttr(std::ostream& /* into */, const std::string& attr, const int& val, const bool isNull, const bool /* escape */) {
381 : assert(!myImpl->myCheckColumns);
382 23108 : myImpl->checkBuilder<std::string, arrow::Int32Builder>(attr, arrow::int32);
383 23108 : myImpl->myValues.push_back(isNull ? nullptr : std::make_shared<arrow::Int32Scalar>(val));
384 23108 : }
385 :
386 :
387 : void
388 36882 : ParquetFormatter::writeStringAttr(const SumoXMLAttr attr, const std::string& val) {
389 36882 : myImpl->checkAttr(attr);
390 36882 : myImpl->checkBuilder<SumoXMLAttr, arrow::StringBuilder>(attr, arrow::utf8);
391 36882 : myImpl->myValues.push_back(std::make_shared<arrow::StringScalar>(val));
392 36882 : }
393 :
394 :
395 : void
396 22531 : ParquetFormatter::writeStringAttr(const std::string& attr, const std::string& val) {
397 : assert(!myImpl->myCheckColumns);
398 22531 : myImpl->checkBuilder<std::string, arrow::StringBuilder>(attr, arrow::utf8);
399 22531 : myImpl->myValues.push_back(std::make_shared<arrow::StringScalar>(val));
400 22531 : }
401 :
402 :
403 : void
404 75582 : ParquetFormatter::writeNullAttr(const SumoXMLAttr attr) {
405 75582 : myImpl->checkAttr(attr);
406 75582 : myImpl->checkBuilder<SumoXMLAttr, arrow::StringBuilder>(attr, arrow::utf8);
407 75582 : myImpl->myValues.push_back(nullptr);
408 75582 : }
409 :
410 :
411 : void
412 432 : ParquetFormatter::writeNullAttr(const std::string& attr) {
413 : assert(!myImpl->myCheckColumns);
414 432 : myImpl->checkBuilder<std::string, arrow::StringBuilder>(attr, arrow::utf8);
415 432 : myImpl->myValues.push_back(nullptr);
416 432 : }
417 :
418 :
419 : void
420 1156 : ParquetFormatter::writeTime(std::ostream& /* into */, const SumoXMLAttr attr, const SUMOTime val) {
421 1156 : if (!gHumanReadableTime) {
422 : // always float64 for machine-readable time, regardless of stream precision
423 1156 : myImpl->checkBuilder<SumoXMLAttr, arrow::DoubleBuilder>(attr, arrow::float64);
424 1156 : myImpl->myValues.push_back(std::make_shared<arrow::DoubleScalar>(STEPS2TIME(val)));
425 1156 : return;
426 : }
427 0 : writeStringAttr(attr, time2string(val));
428 : }
429 :
430 :
431 : bool
432 0 : ParquetFormatter::wroteHeader() const {
433 0 : return myImpl->myWroteHeader;
434 : }
435 :
436 :
437 : void
438 58 : ParquetFormatter::setExpectedAttributes(const SumoXMLAttrMask& expected, const int depth) {
439 58 : myImpl->myExpectedAttrs = expected;
440 58 : myImpl->myMaxDepth = depth;
441 58 : myImpl->myCheckColumns = expected.any();
442 58 : }
443 :
444 :
445 : /****************************************************************************/
|
