Project
Loading...
Searching...
No Matches
test_Root2ArrowTable.cxx
Go to the documentation of this file.
1// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
2// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3// All rights not expressly granted are reserved.
4//
5// This software is distributed under the terms of the GNU General Public
6// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7//
8// In applying this license CERN does not waive the privileges and immunities
9// granted to it by virtue of its status as an Intergovernmental Organization
10// or submit itself to any jurisdiction.
11
12#include <catch_amalgamated.hpp>
13
14#include "Framework/TableBuilder.h"
15#include "Framework/ASoA.h"
16#include "Framework/PluginManager.h"
17#include "../src/ArrowDebugHelpers.h"
18
19#include <ROOT/RDataFrame.hxx>
20#include <ROOT/RArrowDS.hxx>
21#include <TBufferFile.h>
22#include <TClass.h>
23#include <TDirectoryFile.h>
24#include <TMemFile.h>
25#include <TDirectory.h>
26#include <TTree.h>
27#include <TRandom.h>
28#include <TFile.h>
29#include <ROOT/RField.hxx>
30#include <ROOT/RNTuple.hxx>
31#include <ROOT/RNTupleDescriptor.hxx>
32#include <ROOT/RNTupleModel.hxx>
33#include <ROOT/RNTupleReader.hxx>
34#include <ROOT/RNTupleWriter.hxx>
35#include <memory>
36
37#include <arrow/array/array_primitive.h>
38#include <arrow/array/builder_primitive.h>
39#include <arrow/buffer.h>
40#include <arrow/dataset/scanner.h>
41#include <arrow/record_batch.h>
42#include <arrow/table.h>
43#include <arrow/ipc/writer.h>
44#include <arrow/io/memory.h>
45#include <arrow/ipc/writer.h>
46#include <arrow/ipc/reader.h>
47#include "Framework/RootArrowFilesystem.h"
48
49using namespace o2::framework;
50
51namespace o2::aod
52{
53namespace test
54{
55DECLARE_SOA_COLUMN_FULL(Px, px, float, "px");
56DECLARE_SOA_COLUMN_FULL(Py, py, float, "py");
57DECLARE_SOA_COLUMN_FULL(Pz, pz, float, "pz");
58DECLARE_SOA_COLUMN_FULL(Xyz, xyz, float[3], "xyz");
59DECLARE_SOA_COLUMN_FULL(Ij, ij, int[2], "ij");
60DECLARE_SOA_COLUMN_FULL(Random, random, double, "random");
61DECLARE_SOA_COLUMN_FULL(Ev, ev, int, "ev");
62} // namespace test
63
64DECLARE_SOA_TABLE(Test, "AOD", "ETAPHI",
65 test::Px, test::Py, test::Pz, test::Xyz, test::Ij,
66 test::Random, test::Ev);
67} // namespace o2::aod
68
69TEST_CASE("RootTree2Fragment")
70{
71 using namespace o2::framework;
73
75 auto* file = new TBufferFile(TBuffer::kWrite);
76
77 TTree t1("t1", "a simple Tree with simple variables");
78 Float_t xyz[3];
79 Int_t ij[2];
80 Float_t px = 0, py = 1, pz = 2;
81 Double_t random;
82 Int_t ev;
83 t1.Branch("px", &px, "px/F");
84 t1.Branch("py", &py, "py/F");
85 t1.Branch("pz", &pz, "pz/F");
86 t1.Branch("random", &random, "random/D");
87 t1.Branch("ev", &ev, "ev/I");
88 t1.Branch("xyz", xyz, "xyz[3]/F");
89 t1.Branch("ij", ij, "ij[2]/I");
90 // fill the tree
91 for (Int_t i = 0; i < 1000; i++) {
92 xyz[0] = 1;
93 xyz[1] = 2;
94 xyz[2] = 3;
95 gRandom->Rannor(px, py);
96 pz = px * px + py * py;
97 xyz[2] = i + 1;
98 ij[0] = i;
99 ij[1] = i + 1;
100 random = gRandom->Rndm();
101 ev = i + 1;
102 t1.Fill();
103 }
104 file->WriteObjectAny(&t1, t1.Class());
105 auto* fileRead = new TBufferFile(TBuffer::kRead, file->BufferSize(), file->Buffer(), false, nullptr);
106
107 std::vector<char const*> capabilitiesSpecs = {
108 "O2Framework:RNTupleObjectReadingCapability",
109 "O2Framework:TTreeObjectReadingCapability",
110 };
111
112 std::vector<LoadablePlugin> plugins;
113 for (auto spec : capabilitiesSpecs) {
114 auto morePlugins = PluginManager::parsePluginSpecString(spec);
115 for (auto& extra : morePlugins) {
116 plugins.push_back(extra);
117 }
118 }
119 REQUIRE(plugins.size() == 2);
120
121 RootObjectReadingFactory factory;
122 std::vector<char const*> configDiscoverySpec = {};
123 PluginManager::loadFromPlugin<RootObjectReadingCapability, RootObjectReadingCapabilityPlugin>(plugins, factory.capabilities);
124 REQUIRE(factory.capabilities.size() == 2);
125 REQUIRE(factory.capabilities[0].name == "rntuple");
126 REQUIRE(factory.capabilities[1].name == "ttree");
127
128 // Plugins are hardcoded for now...
129 auto format = factory.capabilities[1].factory().format();
130
131 auto fs = std::make_shared<TBufferFileFS>(fileRead, factory);
132
133 arrow::dataset::FileSource source("p", fs);
134 REQUIRE(format->IsSupported(source) == true);
135 auto schemaOpt = format->Inspect(source);
136 REQUIRE(schemaOpt.ok());
137 auto schema = *schemaOpt;
138 REQUIRE(schema->num_fields() == 7);
139 REQUIRE(schema->field(0)->type()->id() == arrow::float32()->id());
140 REQUIRE(schema->field(1)->type()->id() == arrow::float32()->id());
141 REQUIRE(schema->field(2)->type()->id() == arrow::float32()->id());
142 REQUIRE(schema->field(3)->type()->id() == arrow::float64()->id());
143 REQUIRE(schema->field(4)->type()->id() == arrow::int32()->id());
144 REQUIRE(schema->field(5)->type()->id() == arrow::fixed_size_list(arrow::float32(), 3)->id());
145 REQUIRE(schema->field(6)->type()->id() == arrow::fixed_size_list(arrow::int32(), 2)->id());
146 auto fragment = format->MakeFragment(source, {}, schema);
147 REQUIRE(fragment.ok());
148 auto options = std::make_shared<arrow::dataset::ScanOptions>();
149 options->dataset_schema = schema;
150 auto scanner = format->ScanBatchesAsync(options, *fragment);
151 REQUIRE(scanner.ok());
152 auto batches = (*scanner)();
153 auto result = batches.result();
154 REQUIRE(result.ok());
155 REQUIRE((*result)->columns().size() == 7);
156 REQUIRE((*result)->num_rows() == 1000);
157}
158
159bool validateContents(std::shared_ptr<arrow::RecordBatch> batch)
160{
161 {
162 auto int_array = std::static_pointer_cast<arrow::Int32Array>(batch->GetColumnByName("ev"));
163 REQUIRE(int_array->length() == 100);
164 for (int64_t j = 0; j < int_array->length(); j++) {
165 REQUIRE(int_array->Value(j) == j + 1);
166 }
167 }
168
169 {
170 auto list_array = std::static_pointer_cast<arrow::FixedSizeListArray>(batch->GetColumnByName("xyz"));
171
172 REQUIRE(list_array->length() == 100);
173 // Iterate over the FixedSizeListArray
174 for (int64_t i = 0; i < list_array->length(); i++) {
175 auto value_slice = list_array->value_slice(i);
176 auto float_array = std::static_pointer_cast<arrow::FloatArray>(value_slice);
177
178 REQUIRE(float_array->Value(0) == 1);
179 REQUIRE(float_array->Value(1) == 2);
180 REQUIRE(float_array->Value(2) == i + 1);
181 }
182 }
183
184 {
185 auto list_array = std::static_pointer_cast<arrow::FixedSizeListArray>(batch->GetColumnByName("ij"));
186
187 REQUIRE(list_array->length() == 100);
188 // Iterate over the FixedSizeListArray
189 for (int64_t i = 0; i < list_array->length(); i++) {
190 auto value_slice = list_array->value_slice(i);
191 auto int_array = std::static_pointer_cast<arrow::Int32Array>(value_slice);
192 REQUIRE(int_array->Value(0) == i);
193 REQUIRE(int_array->Value(1) == i + 1);
194 }
195 }
196
197 {
198 auto bool_array = std::static_pointer_cast<arrow::BooleanArray>(batch->GetColumnByName("bools"));
199
200 REQUIRE(bool_array->length() == 100);
201 for (int64_t j = 0; j < bool_array->length(); j++) {
202 REQUIRE(bool_array->Value(j) == (j % 3 == 0));
203 }
204 }
205
206 {
207 auto list_array = std::static_pointer_cast<arrow::FixedSizeListArray>(batch->GetColumnByName("manyBools"));
208
209 REQUIRE(list_array->length() == 100);
210 for (int64_t i = 0; i < list_array->length(); i++) {
211 auto value_slice = list_array->value_slice(i);
212 auto bool_array = std::static_pointer_cast<arrow::BooleanArray>(value_slice);
213 REQUIRE(bool_array->Value(0) == (i % 4 == 0));
214 REQUIRE(bool_array->Value(1) == (i % 5 == 0));
215 }
216 }
217
218 {
219 auto list_array = std::static_pointer_cast<arrow::ListArray>(batch->GetColumnByName("vla"));
220
221 REQUIRE(list_array->length() == 100);
222 for (int64_t i = 0; i < list_array->length(); i++) {
223 auto value_slice = list_array->value_slice(i);
224 REQUIRE(value_slice->length() == (i % 10));
225 auto int_array = std::static_pointer_cast<arrow::Int32Array>(value_slice);
226 for (size_t j = 0; j < value_slice->length(); j++) {
227 REQUIRE(int_array->Value(j) == j);
228 }
229 }
230 }
231 return true;
232}
233
234bool validateSchema(std::shared_ptr<arrow::Schema> schema)
235{
236 REQUIRE(schema->num_fields() == 11);
237 REQUIRE(schema->field(0)->type()->id() == arrow::float32()->id());
238 REQUIRE(schema->field(1)->type()->id() == arrow::float32()->id());
239 REQUIRE(schema->field(2)->type()->id() == arrow::float32()->id());
240 REQUIRE(schema->field(3)->type()->id() == arrow::float64()->id());
241 REQUIRE(schema->field(4)->type()->id() == arrow::int32()->id());
242 REQUIRE(schema->field(5)->type()->id() == arrow::fixed_size_list(arrow::float32(), 3)->id());
243 REQUIRE(schema->field(6)->type()->id() == arrow::fixed_size_list(arrow::int32(), 2)->id());
244 REQUIRE(schema->field(7)->type()->id() == arrow::boolean()->id());
245 REQUIRE(schema->field(8)->type()->id() == arrow::fixed_size_list(arrow::boolean(), 2)->id());
246 REQUIRE(schema->field(9)->type()->id() == arrow::list(arrow::int32())->id());
247 REQUIRE(schema->field(10)->type()->id() == arrow::int8()->id());
248 return true;
249}
250
251bool validatePhysicalSchema(std::shared_ptr<arrow::Schema> schema)
252{
253 REQUIRE(schema->num_fields() == 12);
254 REQUIRE(schema->field(0)->type()->id() == arrow::float32()->id());
255 REQUIRE(schema->field(0)->name() == "px");
256 REQUIRE(schema->field(1)->type()->id() == arrow::float32()->id());
257 REQUIRE(schema->field(2)->type()->id() == arrow::float32()->id());
258 REQUIRE(schema->field(3)->type()->id() == arrow::float64()->id());
259 REQUIRE(schema->field(4)->type()->id() == arrow::int32()->id());
260 REQUIRE(schema->field(5)->type()->id() == arrow::fixed_size_list(arrow::float32(), 3)->id());
261 REQUIRE(schema->field(6)->type()->id() == arrow::fixed_size_list(arrow::int32(), 2)->id());
262 REQUIRE(schema->field(7)->type()->id() == arrow::boolean()->id());
263 REQUIRE(schema->field(8)->type()->id() == arrow::fixed_size_list(arrow::boolean(), 2)->id());
264 REQUIRE(schema->field(9)->type()->id() == arrow::int32()->id());
265 REQUIRE(schema->field(10)->type()->id() == arrow::list(arrow::int32())->id());
266 REQUIRE(schema->field(11)->type()->id() == arrow::int8()->id());
267 return true;
268}
269
270TEST_CASE("RootTree2Dataset")
271{
272 using namespace o2::framework;
274 // auto *f = new TFile("Foo.root", "RECREATE");
275 auto* f = new TMemFile("foo", "RECREATE");
276 f->mkdir("DF_1");
277 f->mkdir("DF_2");
278
279 f->cd("DF_1");
280 auto* t = new TTree("tracks", "a simple Tree with simple variables");
281 {
282 Float_t xyz[3];
283 Int_t ij[2];
284 Float_t px = 0, py = 1, pz = 2;
285 Double_t random;
286 Int_t ev;
287 t->Branch("px", &px, "px/F");
288 t->Branch("py", &py, "py/F");
289 t->Branch("pz", &pz, "pz/F");
290 t->Branch("random", &random, "random/D");
291 t->Branch("ev", &ev, "ev/I");
292 t->Branch("xyz", xyz, "xyz[3]/F");
293 t->Branch("ij", ij, "ij[2]/I");
294 // fill the tree
295 for (Int_t i = 0; i < 1000; i++) {
296 xyz[0] = 1;
297 xyz[1] = 2;
298 xyz[2] = 3;
299 gRandom->Rannor(px, py);
300 pz = px * px + py * py;
301 xyz[2] = i + 1;
302 ij[0] = i;
303 ij[1] = i + 1;
304 random = gRandom->Rndm();
305 ev = i + 1;
306 t->Fill();
307 }
308 }
309
310 f->cd("DF_2");
311 t = new TTree("tracks", "a simple Tree with simple variables");
312 {
313 Float_t xyz[3];
314 Int_t ij[2];
315 Float_t px = 0, py = 1, pz = 2;
316 Double_t random;
317 Int_t ev;
318 bool oneBool;
319 bool manyBool[2];
320 int vla[10] = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
321 int vlaSize = 0;
322 char byte;
323
324 t->Branch("px", &px, "px/F");
325 t->Branch("py", &py, "py/F");
326 t->Branch("pz", &pz, "pz/F");
327 t->Branch("random", &random, "random/D");
328 t->Branch("ev", &ev, "ev/I");
329 t->Branch("xyz", xyz, "xyz[3]/F");
330 t->Branch("ij", ij, "ij[2]/I");
331 t->Branch("bools", &oneBool, "bools/O");
332 t->Branch("manyBools", &manyBool, "manyBools[2]/O");
333 t->Branch("vla_size", &vlaSize, "vla_size/I");
334 t->Branch("vla", vla, "vla[vla_size]/I");
335 t->Branch("byte", &byte, "byte/B");
336 // fill the tree
337 for (Int_t i = 0; i < 100; i++) {
338 xyz[0] = 1;
339 xyz[1] = 2;
340 xyz[2] = 3;
341 gRandom->Rannor(px, py);
342 pz = px * px + py * py;
343 xyz[2] = i + 1;
344 ij[0] = i;
345 ij[1] = i + 1;
346 random = gRandom->Rndm();
347 ev = i + 1;
348 oneBool = (i % 3 == 0);
349 manyBool[0] = (i % 4 == 0);
350 manyBool[1] = (i % 5 == 0);
351 vlaSize = i % 10;
352 byte = i;
353 t->Fill();
354 }
355 }
356 f->Write();
357
358 std::vector<char const*> capabilitiesSpecs = {
359 "O2Framework:RNTupleObjectReadingCapability",
360 "O2Framework:TTreeObjectReadingCapability",
361 };
362
363 RootObjectReadingFactory factory;
364
365 std::vector<LoadablePlugin> plugins;
366 for (auto spec : capabilitiesSpecs) {
367 auto morePlugins = PluginManager::parsePluginSpecString(spec);
368 for (auto& extra : morePlugins) {
369 plugins.push_back(extra);
370 }
371 }
372 REQUIRE(plugins.size() == 2);
373
374 PluginManager::loadFromPlugin<RootObjectReadingCapability, RootObjectReadingCapabilityPlugin>(plugins, factory.capabilities);
375
376 REQUIRE(factory.capabilities.size() == 2);
377 REQUIRE(factory.capabilities[0].name == "rntuple");
378 REQUIRE(factory.capabilities[1].name == "ttree");
379
380 // Plugins are hardcoded for now...
381 auto rNtupleFormat = factory.capabilities[0].factory().format();
382 auto format = factory.capabilities[1].factory().format();
383
384 auto fs = std::make_shared<TFileFileSystem>(f, 50 * 1024 * 1024, factory);
385
386 arrow::dataset::FileSource source("DF_2/tracks", fs);
387 REQUIRE(format->IsSupported(source) == true);
388 auto physicalSchema = format->Inspect(source);
389 REQUIRE(physicalSchema.ok());
390 REQUIRE(validatePhysicalSchema(*physicalSchema));
391 // Create the dataset schema rather than using the physical one
392 std::vector<std::shared_ptr<arrow::Field>> fields;
393 for (auto& field : (*(physicalSchema))->fields()) {
394 if (field->name().ends_with("_size")) {
395 continue;
396 }
397 fields.push_back(field);
398 }
399 std::shared_ptr<arrow::Schema> schema = std::make_shared<arrow::Schema>(fields);
400
401 validateSchema(schema);
402
403 auto fragment = format->MakeFragment(source, {}, *physicalSchema);
404 REQUIRE(fragment.ok());
405 auto options = std::make_shared<arrow::dataset::ScanOptions>();
406 options->dataset_schema = schema;
407 auto scanner = format->ScanBatchesAsync(options, *fragment);
408 REQUIRE(scanner.ok());
409
410 // This is batch has deferred contents. Therefore we need to use a DeferredOutputStream to
411 // write it to a real one and read it back with the BufferReader, which is hopefully zero copy
412 std::shared_ptr<arrow::RecordBatch> batch;
413
414 auto batches = (*scanner)();
415 auto result = batches.result();
416 REQUIRE(result.ok());
417 REQUIRE((*result)->columns().size() == 11);
418 REQUIRE((*result)->num_rows() == 100);
419 std::shared_ptr<arrow::ResizableBuffer> buffer = *arrow::AllocateResizableBuffer(1000, 64);
420 auto deferredWriterStream = factory.capabilities[1].factory().deferredOutputStreamer(*fragment, buffer);
421 auto outBatch = arrow::ipc::MakeStreamWriter(deferredWriterStream.get(), schema);
422 auto status = outBatch.ValueOrDie()->WriteRecordBatch(**result);
423 std::shared_ptr<arrow::io::InputStream> bufferReader = std::make_shared<arrow::io::BufferReader>(buffer);
424 auto readerResult = arrow::ipc::RecordBatchStreamReader::Open(bufferReader);
425 auto batchReader = readerResult.ValueOrDie();
426
427 auto next = batchReader->ReadNext(&batch);
428 REQUIRE(batch != nullptr);
429
430 validateContents(batch);
431
432 auto* output = new TMemFile("foo", "RECREATE");
433 auto outFs = std::make_shared<TFileFileSystem>(output, 0, factory);
434
435 // Open a stream at toplevel
436 auto destination = outFs->OpenOutputStream("/", {});
437 REQUIRE(destination.ok());
438
439 // Write to the /DF_3 tree at top level
440 arrow::fs::FileLocator locator{outFs, "/DF_3"};
441 auto writer = format->MakeWriter(*destination, schema, {}, locator);
442 auto success = writer->get()->Write(batch);
443 REQUIRE(batch->schema()->field(0)->name() == "px");
444 auto rootDestination = std::dynamic_pointer_cast<TDirectoryFileOutputStream>(*destination);
445
446 SECTION("Read tree")
447 {
448 REQUIRE(success.ok());
449 // Let's read it back...
450 auto tfileFs = std::dynamic_pointer_cast<TFileFileSystem>(outFs);
451 REQUIRE(tfileFs.get());
452 REQUIRE(tfileFs->GetFile());
453 auto* tree = (TTree*)tfileFs->GetFile()->GetObjectChecked("/DF_3", TClass::GetClass("TTree"));
454 REQUIRE(tree != nullptr);
455 REQUIRE(((TBranch*)tree->GetListOfBranches()->At(0))->GetEntries() == 100);
456 REQUIRE(((TBranch*)tree->GetListOfBranches()->At(0))->GetName() == std::string("px"));
457
458 arrow::dataset::FileSource source2("/DF_3", outFs);
459
460 REQUIRE(format->IsSupported(source2) == true);
461 tfileFs = std::dynamic_pointer_cast<TFileFileSystem>(source2.filesystem());
462 REQUIRE(tfileFs.get());
463 REQUIRE(tfileFs->GetFile());
464 REQUIRE(tfileFs->GetFile()->GetObjectChecked("/DF_3", TClass::GetClass("TTree")));
465
466 tree = (TTree*)tfileFs->GetFile()->GetObjectChecked("/DF_3", TClass::GetClass("TTree"));
467 REQUIRE(tree != nullptr);
468 REQUIRE(((TBranch*)tree->GetListOfBranches()->At(0))->GetEntries() == 100);
469
470 auto schemaOptWritten = format->Inspect(source2);
471 tfileFs = std::dynamic_pointer_cast<TFileFileSystem>(source2.filesystem());
472 REQUIRE(tfileFs.get());
473 REQUIRE(tfileFs->GetFile());
474 REQUIRE(tfileFs->GetFile()->GetObjectChecked("/DF_3", TClass::GetClass("TTree")));
475 REQUIRE(schemaOptWritten.ok());
476 auto schemaWritten = *schemaOptWritten;
477
478 tree = (TTree*)tfileFs->GetFile()->GetObjectChecked("/DF_3", TClass::GetClass("TTree"));
479 REQUIRE(tree != nullptr);
480 REQUIRE(((TBranch*)tree->GetListOfBranches()->At(0))->GetEntries() == 100);
481
482 REQUIRE(validatePhysicalSchema(schemaWritten));
483 std::vector<std::shared_ptr<arrow::Field>> fields;
484 for (auto& field : schemaWritten->fields()) {
485 if (field->name().ends_with("_size")) {
486 continue;
487 }
488 fields.push_back(field);
489 }
490 std::shared_ptr<arrow::Schema> schema = std::make_shared<arrow::Schema>(fields);
491 REQUIRE(validateSchema(schema));
492
493 auto fragmentWritten = format->MakeFragment(source2, {}, *physicalSchema);
494 REQUIRE(fragmentWritten.ok());
495 auto optionsWritten = std::make_shared<arrow::dataset::ScanOptions>();
496 optionsWritten->dataset_schema = schema;
497 auto scannerWritten = format->ScanBatchesAsync(optionsWritten, *fragmentWritten);
498 REQUIRE(scannerWritten.ok());
499 tree = (TTree*)tfileFs->GetFile()->GetObjectChecked("/DF_3", TClass::GetClass("TTree"));
500 REQUIRE(tree != nullptr);
501 REQUIRE(((TBranch*)tree->GetListOfBranches()->At(0))->GetEntries() == 100);
502 auto batchesWritten = (*scannerWritten)();
503 auto resultWritten = batchesWritten.result();
504 REQUIRE(resultWritten.ok());
505 REQUIRE((*resultWritten)->columns().size() == 11);
506 REQUIRE((*resultWritten)->num_rows() == 100);
507
508 std::shared_ptr<arrow::ResizableBuffer> buffer = *arrow::AllocateResizableBuffer(1000, 64);
509 auto deferredWriterStream2 = factory.capabilities[1].factory().deferredOutputStreamer(*fragmentWritten, buffer);
510 auto outBatch = arrow::ipc::MakeStreamWriter(deferredWriterStream2.get(), schema);
511 auto status = outBatch.ValueOrDie()->WriteRecordBatch(**resultWritten);
512 std::shared_ptr<arrow::io::InputStream> bufferReader = std::make_shared<arrow::io::BufferReader>(buffer);
513 auto readerResult = arrow::ipc::RecordBatchStreamReader::Open(bufferReader);
514 auto batchReader = readerResult.ValueOrDie();
515
516 auto next = batchReader->ReadNext(&batch);
517 REQUIRE(batch != nullptr);
518 validateContents(batch);
519 }
520
521#if __has_include(<ROOT/RFieldBase.hxx>)
522 arrow::fs::FileLocator rnTupleLocator{outFs, "rntuple"};
523#else
524 arrow::fs::FileLocator rnTupleLocator{outFs, "/rntuple"};
525#endif
526 // We write an RNTuple in the same TMemFile, using /rntuple as a location
527 auto rntupleDestination = std::dynamic_pointer_cast<TDirectoryFileOutputStream>(*destination);
528
529 {
530 auto rNtupleWriter = rNtupleFormat->MakeWriter(*destination, schema, {}, rnTupleLocator);
531 auto rNtupleSuccess = rNtupleWriter->get()->Write(batch);
532 REQUIRE(rNtupleSuccess.ok());
533 }
534
535 // And now we can read back the RNTuple into a RecordBatch
536#if __has_include(<ROOT/RFieldBase.hxx>)
537 arrow::dataset::FileSource writtenRntupleSource("rntuple", outFs);
538#else
539 arrow::dataset::FileSource writtenRntupleSource("/rntuple", outFs);
540#endif
541
542 REQUIRE(rNtupleFormat->IsSupported(writtenRntupleSource) == true);
543
544 auto rntupleSchemaOpt = rNtupleFormat->Inspect(writtenRntupleSource);
545 REQUIRE(rntupleSchemaOpt.ok());
546 auto rntupleSchemaWritten = *rntupleSchemaOpt;
547 REQUIRE(validateSchema(rntupleSchemaWritten));
548
549 auto rntupleFragmentWritten = rNtupleFormat->MakeFragment(writtenRntupleSource, {}, rntupleSchemaWritten);
550 REQUIRE(rntupleFragmentWritten.ok());
551 auto rntupleOptionsWritten = std::make_shared<arrow::dataset::ScanOptions>();
552 rntupleOptionsWritten->dataset_schema = rntupleSchemaWritten;
553 auto rntupleScannerWritten = rNtupleFormat->ScanBatchesAsync(rntupleOptionsWritten, *rntupleFragmentWritten);
554 REQUIRE(rntupleScannerWritten.ok());
555 auto rntupleBatchesWritten = (*rntupleScannerWritten)();
556 auto rntupleResultWritten = rntupleBatchesWritten.result();
557 REQUIRE(rntupleResultWritten.ok());
558 REQUIRE((*rntupleResultWritten)->columns().size() == 11);
559 REQUIRE(validateSchema((*rntupleResultWritten)->schema()));
560 REQUIRE((*rntupleResultWritten)->num_rows() == 100);
561 REQUIRE(validateContents(*rntupleResultWritten));
562}
DECLARE_SOA_TABLE
#define DECLARE_SOA_TABLE(_Name_, _Origin_, _Desc_,...)
Definition ASoA.h:3229
DECLARE_SOA_COLUMN_FULL
#define DECLARE_SOA_COLUMN_FULL(_Name_, _Getter_, _Type_, _Label_)
Definition ASoA.h:2346
schema
std::shared_ptr< arrow::Schema > schema
Definition ExpressionJSONHelpers.cxx:666
fields
std::vector< std::shared_ptr< arrow::Field > > fields
Definition ExpressionJSONHelpers.cxx:667
i
int32_t i
Definition GPUCommonAlgorithm.h:451
Test
Test
Definition Utils.h:55
output
void output(const std::map< std::string, ChannelStat > &channels)
Definition rawdump.cxx:197
j
uint32_t j
Definition RawData.h:0
Float_t
int64_t
result
GLuint64EXT * result
Definition glcorearb.h:5662
buffer
GLuint buffer
Definition glcorearb.h:655
f
GLdouble f
Definition glcorearb.h:310
source
GLsizei GLsizei GLchar * source
Definition glcorearb.h:798
format
GLint GLint GLsizei GLint GLenum format
Definition glcorearb.h:275
t1
GLuint GLfloat GLfloat GLfloat GLfloat GLfloat GLfloat GLfloat GLfloat GLfloat t1
Definition glcorearb.h:5034
o2::framework
Defining ITS Vertex explicitly as messageable.
Definition Cartesian.h:288
o2::framework::TEST_CASE
TEST_CASE("test_prepareArguments")
Definition test_DeviceSpecHelpers.cxx:102
test
FIXME: do not use data model tables.
Definition benchmark_TableToTree.cxx:25
o2::framework::PluginManager::parsePluginSpecString
static std::vector< LoadablePlugin > parsePluginSpecString(char const *str)
Parse a comma separated list of <library>:<plugin-name> plugin declarations.
Definition PluginManager.cxx:19
o2::framework::RootObjectReadingFactory::capabilities
std::vector< RootObjectReadingCapability > capabilities
Definition RootArrowFilesystem.h:141
validateContents
bool validateContents(std::shared_ptr< arrow::RecordBatch > batch)
Definition test_Root2ArrowTable.cxx:159
validateSchema
bool validateSchema(std::shared_ptr< arrow::Schema > schema)
Definition test_Root2ArrowTable.cxx:234
validatePhysicalSchema
bool validatePhysicalSchema(std::shared_ptr< arrow::Schema > schema)
Definition test_Root2ArrowTable.cxx:251
tree
std::unique_ptr< TTree > tree((TTree *) flIn.Get(std::string(o2::base::NameConf::CTFTREENAME).c_str()))