SPB Git

spb/wp40_uqo

Public MIT
5.4 KB · 70 lines cpp
1// Arrow/Parquet reader for FinData bars (schema v1). Column chunks are read with the Arrow C++ API; numeric columns2// are copied once into the BarSeries vectors (a true zero-copy view is only possible for a single contiguous chunk of3// float64 without nulls — the reader records whether that was the case in `metadata["zero_copy_possible"]`).4#include "raid/data/io.hpp"5#include "raid/core/error.hpp"6#include <arrow/api.h>7#include <arrow/io/api.h>8#include <arrow/ipc/api.h>9#include <parquet/arrow/reader.h>10#include <unordered_set>1112namespace raid::data {1314namespace {15ParquetBars table_to_bars(const std::shared_ptr<arrow::Table>& table, const std::vector<std::string>& symbols, Interval iv) {16  ParquetBars out;17  if (table->schema()->metadata()) for (int i = 0; i < table->schema()->metadata()->size(); ++i) out.metadata[table->schema()->metadata()->key(i)] = table->schema()->metadata()->value(i);18  auto col = [&](const char* name) { auto c = table->GetColumnByName(name); RAID_REQUIRE(c != nullptr, io, std::string("missing column ") + name); return c; };19  auto ts = col("timestamp_utc"); auto sym = col("symbol_as_reported"); auto iid = table->GetColumnByName("instrument_id");   // may be absent (projected read)20  auto o = col("open"), h = col("high"), l = col("low"), c = col("close"), v = col("volume");21  std::unordered_set<std::string> keep(symbols.begin(), symbols.end());22  bool contiguous = (c->num_chunks() == 1) && (c->null_count() == 0);23  out.metadata["zero_copy_possible"] = contiguous ? "true" : "false";24  // Fast path: symbols are stored in runs (one symbol = one contiguous block in FinData exports); avoid a map lookup per row.25  BarSeries* cur = nullptr; std::string cur_sym;26  for (int ch = 0; ch < ts->num_chunks(); ++ch) {27    auto ta = std::static_pointer_cast<arrow::TimestampArray>(ts->chunk(ch));28    auto sa = std::static_pointer_cast<arrow::StringArray>(sym->chunk(ch));29    std::shared_ptr<arrow::StringArray> ia = iid ? std::static_pointer_cast<arrow::StringArray>(iid->chunk(ch)) : nullptr;30    auto oa = std::static_pointer_cast<arrow::DoubleArray>(o->chunk(ch)), ha = std::static_pointer_cast<arrow::DoubleArray>(h->chunk(ch));31    auto la = std::static_pointer_cast<arrow::DoubleArray>(l->chunk(ch)), ca = std::static_pointer_cast<arrow::DoubleArray>(c->chunk(ch));32    auto va = std::static_pointer_cast<arrow::DoubleArray>(v->chunk(ch));33    const auto unit = std::static_pointer_cast<arrow::TimestampType>(ta->type())->unit();34    i64 mult = 1; if (unit == arrow::TimeUnit::SECOND) mult = 1'000'000; else if (unit == arrow::TimeUnit::MILLI) mult = 1000; else if (unit == arrow::TimeUnit::NANO) mult = -1000;35    for (int64_t i = 0; i < ta->length(); ++i) {36      const std::string_view sv = sa->GetView(i);37      if (!cur || sv != cur_sym) { cur_sym = std::string(sv); if (!keep.empty() && !keep.count(cur_sym)) { cur = nullptr; continue; } cur = &out.series[cur_sym]; if (cur->symbol.empty()) { cur->symbol = cur_sym; cur->instrument_id = ia ? ia->GetString(i) : ("firstrate:" + std::string(std::static_pointer_cast<arrow::StringArray>(col("asset_class")->chunk(ch))->GetView(i)) + ":" + cur_sym); cur->interval = iv; cur->adjustment = "raw"; cur->reserve(4096); } }38      else if (!cur) continue;39      const i64 raw = ta->Value(i); const i64 us = mult > 0 ? raw * mult : raw / 1000;40      cur->push_back(Timestamp(us), oa->Value(i), ha->Value(i), la->Value(i), ca->Value(i), va->Value(i));41      ++out.rows_read;42    }43  }44  for (auto& [k, bs] : out.series) { auto st = bs.normalize(); out.duplicates_dropped += st.duplicates_removed; }45  return out;46}47}  // namespace4849ParquetBars read_bars_parquet(const std::string& path, const std::vector<std::string>& symbols, Interval iv) {50  auto infile = arrow::io::ReadableFile::Open(path); RAID_REQUIRE(infile.ok(), io, "cannot open " + path + ": " + infile.status().ToString());51  auto reader_r = parquet::arrow::OpenFile(*infile, arrow::default_memory_pool()); RAID_REQUIRE(reader_r.ok(), io, "parquet open: " + reader_r.status().ToString());52  std::unique_ptr<parquet::arrow::FileReader> reader = std::move(*reader_r);53  // Column projection: only the 8 columns the estimators need (of 20) are decoded — the key I/O optimisation of RQ3.54  std::shared_ptr<arrow::Schema> schema; auto sst = reader->GetSchema(&schema); RAID_REQUIRE(sst.ok(), io, "parquet schema: " + sst.ToString());55  // (instrument_id is derived from symbol_as_reported for FinData exports: "firstrate:<class>:<SYMBOL>"; reading the string column costs more than every numeric column together)56  std::vector<int> idx; for (const char* name : {"symbol_as_reported", "asset_class", "timestamp_utc", "open", "high", "low", "close", "volume"}) { const int k = schema->GetFieldIndex(name); RAID_REQUIRE(k >= 0, io, std::string("missing column ") + name); idx.push_back(k); }57  reader->set_use_threads(true);58  std::shared_ptr<arrow::Table> table; auto st = reader->ReadTable(idx, &table); RAID_REQUIRE(st.ok(), io, "parquet read: " + st.ToString());59  return table_to_bars(table, symbols, iv);60}6162ParquetBars read_bars_arrow_ipc(const std::string& path, Interval iv) {63  auto infile = arrow::io::ReadableFile::Open(path); RAID_REQUIRE(infile.ok(), io, "cannot open " + path);64  auto reader = arrow::ipc::RecordBatchStreamReader::Open(*infile); RAID_REQUIRE(reader.ok(), io, "arrow ipc open: " + reader.status().ToString());65  auto table = (*reader)->ToTable(); RAID_REQUIRE(table.ok(), io, "arrow ipc read: " + table.status().ToString());66  return table_to_bars(*table, {}, iv);67}6869}  // namespace raid::data70
Files · ae5a4e6ff6df81cef716eacbb5d51143a008d45b
Menu

Theme