29811479573

Committed 21 Jul 2026 07:42AM UTC coverage: 52.899% (+0.2%) from 52.66%

Build # 29811479573

Build Type

Pull #99

github

Committed by

web-flow

Commit Message

Merge ee9af7c24 into 06bc84ec9

Pull Request Pull Request #99: Support CM time_metric (NS/MS/SEC/US) across reader and viz

Coverage Stats

41156 of 100801 branches covered (40.83%)

Branch coverage included in aggregate %.

1065 of 1185 new or added lines in 30 files covered. (89.87%)

23 existing lines in 5 files now uncovered.

36312 of 45644 relevant lines covered (79.55%)

23458.91 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

52.21

/src/dftracer/utils/utilities/reader/trace_reader_arrow.cpp

#include <dftracer/utils/utilities/reader/trace_reader.h>
#ifdef DFTRACER_UTILS_ENABLE_ARROW
#include <dftracer/utils/core/common/string_arena.h>
#include <dftracer/utils/utilities/common/arrow/column_builder.h>
#include <dftracer/utils/utilities/common/query/query.h>
#include <dftracer/utils/utilities/composites/dft/indexing/resolved_field_rewriter.h>
#include <dftracer/utils/utilities/indexer/index_database.h>
#include <dftracer/utils/utilities/reader/internal/reader.h>
#include <dftracer/utils/utilities/reader/internal/trace_reader_prefilter.h>
#include <dftracer/utils/utilities/reader/internal/trace_reader_shared.h>
#include <simdjson.h>

#include <cstring>
#include <optional>
#include <span>
#include <string>
#include <vector>

namespace dftracer::utils::utilities::reader {

namespace indexing = composites::dft::indexing;
using common::query::Query;
using internal::build_prefilter;
using internal::CompiledEqProbe;
using internal::eval_compiled_eq;
using internal::LinePrefilter;
using internal::ondemand_to_literal;
using internal::read_chunks_indexed;
using internal::strip_ndjson_bookends;
using internal::try_compile_eq_probes;

namespace {

using common::arrow::ArrowExportResult;
using common::arrow::ColumnType;
using common::arrow::RecordBatchBuilder;

struct ArrowKeyHint {
    std::string key;
    std::size_t col_idx = 0;
    ColumnType type = ColumnType::INT64;
    bool valid = false;
};

inline std::size_t resolve_col_idx(RecordBatchBuilder& builder,
                                   std::vector<ArrowKeyHint>& hints,
                                   std::size_t pos, std::string_view key_sv,
                                   ColumnType type) {
    if (pos < hints.size()) {
        auto& h = hints[pos];
        if (h.valid && h.type == type && h.key.size() == key_sv.size() &&
            std::memcmp(h.key.data(), key_sv.data(), key_sv.size()) == 0) {
            return h.col_idx;
        }
    }
    // Position-keyed miss. Variable-shape rows (e.g., open vs read events
    // with different args fields) push fields to different positions, so
    // the position cache misses constantly while the underlying schema is
    // small (~15 keys). A linear scan over the hint vector with a SIMD
    // memcmp beats RecordBatchBuilder's name_to_index_ hash lookup for this
    // size.
    for (std::size_t i = 0; i < hints.size(); ++i) {
        if (i == pos) continue;
        auto& h = hints[i];
        if (h.valid && h.type == type && h.key.size() == key_sv.size() &&
            std::memcmp(h.key.data(), key_sv.data(), key_sv.size()) == 0) {
            if (pos < hints.size()) {
                auto& slot = hints[pos];
                slot.key.assign(key_sv);
                slot.type = type;
                slot.col_idx = h.col_idx;
                slot.valid = true;
            }
            return h.col_idx;
        }
    }
    std::size_t idx = builder.add_or_get_column(key_sv, type);
    if (pos >= hints.size()) hints.resize(pos + 1);
    auto& h = hints[pos];
    h.key.assign(key_sv);
    h.type = type;
    h.col_idx = idx;
    h.valid = true;
    return idx;
}

// Append a typed scalar value under `key_sv`. Nested objects/arrays are
// always round-tripped as JSON strings (flattening is one level only).
void append_scalar_or_json(RecordBatchBuilder& builder,
                           std::vector<ArrowKeyHint>& hints, std::size_t& pos,
                           std::string_view key_sv,
                           simdjson::ondemand::value val,
                           simdjson::ondemand::json_type type) {
    switch (type) {
        case simdjson::ondemand::json_type::number: {
            auto num_r = val.get_number();
            if (num_r.error()) break;
            auto num = num_r.value();
            if (num.is_int64()) {
                auto idx = resolve_col_idx(builder, hints, pos++, key_sv,
                                           ColumnType::INT64);
                builder.append_int64(idx, num.get_int64());
            } else if (num.is_uint64()) {
                auto idx = resolve_col_idx(builder, hints, pos++, key_sv,
                                           ColumnType::UINT64);
                builder.append_uint64(idx, num.get_uint64());
            } else {
                auto idx = resolve_col_idx(builder, hints, pos++, key_sv,
                                           ColumnType::DOUBLE);
                builder.append_double(idx, num.get_double());
            }
            break;
        }
        case simdjson::ondemand::json_type::string: {
            auto str_r = val.get_string();
            if (str_r.error()) break;
            auto idx = resolve_col_idx(builder, hints, pos++, key_sv,
                                       ColumnType::STRING);
            builder.append_string(idx, str_r.value());
            break;
        }
        case simdjson::ondemand::json_type::boolean: {
            auto b_r = val.get_bool();
            if (b_r.error()) break;
            auto idx = resolve_col_idx(builder, hints, pos++, key_sv,
                                       ColumnType::BOOL);
            builder.append_bool(idx, b_r.value());
            break;
        }
        case simdjson::ondemand::json_type::null: {
            auto existing = builder.find_column(key_sv);
            if (existing) builder.append_null(*existing);
            ++pos;
            break;
        }
        case simdjson::ondemand::json_type::object:
        case simdjson::ondemand::json_type::array: {
            auto raw_r = val.raw_json();
            auto idx = resolve_col_idx(builder, hints, pos++, key_sv,
                                       ColumnType::STRING);
            if (!raw_r.error()) {
                auto sv = raw_r.value();
                builder.append_string(idx, sv);
            } else {
                builder.append_null(idx);
            }
            break;
        }
        default:
            ++pos;
            break;
    }
}

// Append one Arrow row from an already-parsed simdjson document.
// Dynamic schema: new columns appended as they appear. When flatten_objects
// is true, top-level object values are expanded one level into `parent.child`
// columns; deeper nesting still lands as a JSON string under the flattened
// key. Returns false on error paths so callers can skip the row.
bool arrow_row_from_doc(RecordBatchBuilder& builder,
                        std::vector<ArrowKeyHint>& hints,
                        simdjson::ondemand::document_reference doc,
                        bool flatten_objects = false) {
    auto obj_result = doc.get_object();
    if (obj_result.error()) return false;
    char key_buf[512];
    std::size_t pos = 0;
    for (auto field : obj_result.value()) {
        if (field.error()) continue;
        auto key_r = field.unescaped_key();
        if (key_r.error()) continue;
        auto key_sv = key_r.value();
        auto val_r = field.value();
        if (val_r.error()) continue;
        auto val = val_r.value();
        auto type_r = val.type();
        if (type_r.error()) continue;
        auto type = type_r.value();

        if (flatten_objects && type == simdjson::ondemand::json_type::object) {
            auto nested = val.get_object();
            if (nested.error()) continue;
            for (auto nf : nested.value()) {
                if (nf.error()) continue;
                auto nk_r = nf.unescaped_key();
                if (nk_r.error()) continue;
                auto nk = nk_r.value();
                auto nv_r = nf.value();
                if (nv_r.error()) continue;
                auto nv = nv_r.value();
                auto nt_r = nv.type();
                if (nt_r.error()) continue;
                std::size_t needed = key_sv.size() + 1 + nk.size();
                if (needed >= sizeof(key_buf)) continue;
                std::memcpy(key_buf, key_sv.data(), key_sv.size());
                key_buf[key_sv.size()] = '.';
                std::memcpy(key_buf + key_sv.size() + 1, nk.data(), nk.size());
                append_scalar_or_json(builder, hints, pos,
                                      std::string_view(key_buf, needed), nv,
                                      nt_r.value());
            }
            continue;
        }

        append_scalar_or_json(builder, hints, pos, key_sv, val, type);
    }
    builder.end_row();
    return true;
}

void collect_query_fields(simdjson::ondemand::document_reference doc,
                          const Query& query, bool check_dotted,
                          common::query::ValueMap& out);

// Run iterate_many over `padded`, build arrow rows, and emit completed
// batches via `yield_one`. Updates `carry` with the truncated tail (if any)
// for the caller to prepend to the next chunk.
template <typename Yield>
void parse_padded_into_arrow(simdjson::ondemand::parser& bulk_parser,
                             simdjson::padded_string& padded,
                             const std::optional<Query>& query, bool flatten,
                             bool query_has_dotted, RecordBatchBuilder& builder,
                             StringArena& arena,
                             std::vector<ArrowKeyHint>& hints,
                             std::size_t batch_size, std::string* carry,
                             Yield&& yield_one) {
    auto docs_r = bulk_parser.iterate_many(padded, 1 << 20, false);
    if (docs_r.error()) {
        if (carry) carry->clear();
        return;
    }
    auto& docs = docs_r.value();
    for (auto it = docs.begin(); it != docs.end(); ++it) {
        auto doc_result = *it;
        if (doc_result.error()) continue;
        auto& doc = doc_result.value();
        if (query) {
            common::query::ValueMap fields;
            collect_query_fields(doc, *query, query_has_dotted, fields);
            if (!query->evaluate(fields)) continue;
            doc.rewind();
        }
        if (!arrow_row_from_doc(builder, hints, doc, flatten)) continue;
        if (builder.num_rows() >= batch_size) {
            auto result = builder.finish();
            arena.clear();
            if (!builder.is_schema_locked()) builder.lock_schema();
            builder.reset(true);
            builder.reserve(batch_size);
            yield_one(std::move(result));
        }
    }
    if (carry) {
        std::size_t total = padded.size();
        std::size_t truncated = docs.truncated_bytes();
        if (truncated > 0 && truncated <= total) {
            carry->assign(padded.data() + total - truncated,
                          padded.data() + total);
        } else {
            carry->clear();
        }
    }
}

// Build a simdjson-padded buffer containing only the lines in `chunk` that
// pass the line-level prefilter. For queries with no useful prefilter, the
// caller should skip this and feed the raw chunk directly.
std::string collect_matching_lines(std::span<const char> chunk,
                                   const LinePrefilter& prefilter) {
    std::string out;
    out.reserve(chunk.size());
    const char* data = chunk.data();
    std::size_t len = chunk.size();
    std::size_t pos = 0;
    while (pos < len) {
        const void* nl = std::memchr(data + pos, '\n', len - pos);
        std::size_t end_pos = nl ? static_cast<const char*>(nl) - data : len;
        if (end_pos > pos) {
            std::string_view line(data + pos, end_pos - pos);
            if (prefilter.may_match(line)) {
                out.append(line);
                out.push_back('\n');
            }
        }
        pos = end_pos + 1;
    }
    return out;
}

// Extract fields referenced by the query into a ValueMap, walking one level
// of object nesting. Fields not referenced by the query are skipped.
void collect_query_fields(simdjson::ondemand::document_reference doc,
                          const Query& query, bool check_dotted,
                          common::query::ValueMap& out) {
    auto obj = doc.get_object();
    if (obj.error()) return;
    for (auto field : obj.value()) {
        if (field.error()) continue;
        auto key_r = field.unescaped_key();
        if (key_r.error()) continue;
        auto val_r = field.value();
        if (val_r.error()) continue;
        auto key = key_r.value();
        auto val = val_r.value();
        auto type_r = val.type();
        if (type_r.error()) continue;
        auto type = type_r.value();
        if (type == simdjson::ondemand::json_type::object) {
            auto nested = val.get_object();
            if (nested.error()) continue;
            for (auto nf : nested.value()) {
                if (nf.error()) continue;
                auto nk_r = nf.unescaped_key();
                if (nk_r.error()) continue;
                auto nv_r = nf.value();
                if (nv_r.error()) continue;
                internal::store_referenced_nested(out, query, check_dotted, key,
                                                  nk_r.value(), nv_r.value());
            }
        } else if (query.references(key)) {
            out[std::string(key)] = ondemand_to_literal(val);
        }
    }
}

}  // namespace

coro::AsyncGenerator<ArrowExportResult> TraceReader::read_arrow(
    ReadConfig config, std::size_t batch_size) {
    std::optional<Query> query;
    if (!config.query.empty()) {
        auto parsed = Query::from_string(config.query);
        if (!parsed) throw common::query::QueryParseError(parsed.error());
        query = std::move(*parsed);
    }

    // Resolve virtual fields (resolved.*/r.*) to concrete hash in-clauses via
    // the index before the query drives pruning or per-event evaluation.
    if (query && has_index_ && !index_path_.empty() &&
        indexing::has_resolved_fields(*query)) {
        try {
            indexer::IndexDatabase db(
                index_path_, rocksdb::RocksDatabase::OpenMode::ReadOnly);
            if (auto rewritten =
                    indexing::rewrite_resolved_fields(*query, db)) {
                query = std::move(*rewritten);
            }
        } catch (...) {
        }
    }

    // When chunk_prune_only is set, dim_stats already proved every event in
    // the chunk that has the predicate field matches the literal. We still
    // need to skip events that lack the field (e.g., metadata "ph":"M"
    // events lack pid), since the original predicate would reject them.
    std::vector<std::string> presence_check_paths;
    if (query && config.chunk_prune_only) {
        const auto& fset = query->fields();
        presence_check_paths.assign(fset.begin(), fset.end());
    }

    const bool query_has_dotted =
        query && internal::query_references_dotted(*query);

    // For AND-of-EQ predicates, evaluate directly against simdjson without
    // ValueMap (avoids wyhash + per-field std::string allocation per row).
    // Falls back to the ValueMap path on unsupported AST shapes.
    std::vector<CompiledEqProbe> compiled_probes;
    bool use_compiled = false;
    if (query && !config.chunk_prune_only) {
        if (auto p = try_compile_eq_probes(query->root())) {
            compiled_probes = std::move(*p);
            use_compiled = !compiled_probes.empty();
        }
    }

    bool flatten = config.flatten_objects;

    if (!has_index_) {
        // Fallback: drive the per-line read_json path and build rows.
        auto json_gen = read_json(config);
        RecordBatchBuilder builder;
        StringArena arena;
        std::vector<ArrowKeyHint> hints;
        builder.reserve(batch_size);
        while (auto opt = co_await json_gen.next()) {
            if (!arrow_row_from_doc(builder, hints,
                                    simdjson::ondemand::document_reference(
                                        opt->parser->raw_document()),
                                    flatten))
                continue;
            if (builder.num_rows() >= batch_size) {
                co_yield builder.finish();
                arena.clear();
                if (!builder.is_schema_locked()) builder.lock_schema();
                builder.reset(true);
                builder.reserve(batch_size);
            }
        }
        if (builder.num_rows() > 0) {
            co_yield builder.finish();
        }
        co_return;
    }

    // Keep RocksDB alive for the generator's lifetime so per-method opens
    // in GzipIndexer reuse DBManager's cached handle.
    std::optional<indexer::IndexDatabase> db_keep_alive;
    if (has_index_ && !index_path_.empty()) {
        try {
            db_keep_alive.emplace(index_path_,
                                  rocksdb::RocksDatabase::OpenMode::ReadOnly);
        } catch (...) {
        }
    }

    auto reader = create_indexed_reader();
    auto chunk_gen = read_chunks_indexed(
        reader, index_path_, config_.file_path, config, query,
        /*extend_to_line_boundary=*/config.end_at_checkpoint);

    LinePrefilter prefilter = (query && !config.chunk_prune_only)
                                  ? build_prefilter(*query)
                                  : LinePrefilter{};
    bool have_line_prefilter = !prefilter.empty();

    simdjson::ondemand::parser bulk_parser;
    RecordBatchBuilder builder;
    StringArena arena;
    std::vector<ArrowKeyHint> hints;
    builder.reserve(batch_size);

    auto maybe_flush = [&builder, &arena, batch_size](
                           bool final) -> std::optional<ArrowExportResult> {
        if (builder.num_rows() == 0) return std::nullopt;
        if (!final && builder.num_rows() < batch_size) return std::nullopt;
        auto result = builder.finish();
        arena.clear();
        if (!builder.is_schema_locked()) builder.lock_schema();
        builder.reset(true);
        builder.reserve(batch_size);
        return result;
    };

    bool first_chunk = true;
    while (auto chunk_opt = co_await chunk_gen.next()) {
        auto chunk = *chunk_opt;
        if (chunk.empty()) continue;

        // Work items with start_byte > 0 begin at a deflate-block boundary
        // that is typically mid-line; the previous worker emitted that
        // spanning line via its tail-flush, so drop bytes up to (and
        // including) the first newline in our first chunk.
        if (first_chunk && config.start_byte > 0 &&
            config.start_at_checkpoint) {
            const char* nl = static_cast<const char*>(
                std::memchr(chunk.data(), '\n', chunk.size()));
            if (nl) {
                std::size_t skip =
                    static_cast<std::size_t>(nl - chunk.data()) + 1;
                if (skip < chunk.size()) {
                    chunk = chunk.subspan(skip);
                } else {
                    first_chunk = false;
                    continue;
                }
            }
        }
        first_chunk = false;

        simdjson::padded_string padded;
        if (have_line_prefilter) {
            auto collected = collect_matching_lines(chunk, prefilter);
            if (collected.empty()) continue;
            padded = simdjson::padded_string(std::move(collected));
        } else {
            auto trimmed = strip_ndjson_bookends(
                std::string_view(chunk.data(), chunk.size()));
            if (trimmed.empty()) continue;
            padded = simdjson::padded_string(trimmed);
        }

        auto docs_r = bulk_parser.iterate_many(padded, 1 << 20,
                                               /*allow_comma_separated=*/false);
        if (docs_r.error()) continue;
        auto& docs = docs_r.value();

        for (auto it = docs.begin(); it != docs.end(); ++it) {
            auto doc_result = *it;
            if (doc_result.error()) continue;
            auto& doc = doc_result.value();

            if (query && !config.chunk_prune_only) {
                if (use_compiled) {
                    if (!eval_compiled_eq(compiled_probes, doc)) continue;
                } else {
                    common::query::ValueMap fields;
                    collect_query_fields(doc, *query, query_has_dotted, fields);
                    if (!query->evaluate(fields)) continue;
                }
                doc.rewind();
            } else if (!presence_check_paths.empty()) {
                bool all_present = true;
                for (const auto& path : presence_check_paths) {
                    auto fld = doc.find_field_unordered(path);
                    if (fld.error()) {
                        all_present = false;
                        break;
                    }
                }
                if (!all_present) continue;
                doc.rewind();
            }

            if (!arrow_row_from_doc(builder, hints, doc, flatten)) continue;

            if (auto flushed = maybe_flush(/*final=*/false)) {
                co_yield std::move(*flushed);
            }
        }
    }

    if (auto flushed = maybe_flush(/*final=*/true)) {
        co_yield std::move(*flushed);
    }
}

}  // namespace dftracer::utils::utilities::reader

#endif  // DFTRACER_UTILS_ENABLE_ARROW

1	#include <dftracer/utils/utilities/reader/trace_reader.h>
2	#ifdef DFTRACER_UTILS_ENABLE_ARROW
3	#include <dftracer/utils/core/common/string_arena.h>
4	#include <dftracer/utils/utilities/common/arrow/column_builder.h>
5	#include <dftracer/utils/utilities/common/query/query.h>
6	#include <dftracer/utils/utilities/composites/dft/indexing/resolved_field_rewriter.h>
7	#include <dftracer/utils/utilities/indexer/index_database.h>
8	#include <dftracer/utils/utilities/reader/internal/reader.h>
9	#include <dftracer/utils/utilities/reader/internal/trace_reader_prefilter.h>
10	#include <dftracer/utils/utilities/reader/internal/trace_reader_shared.h>
11	#include <simdjson.h>
12
13	#include <cstring>
14	#include <optional>
15	#include <span>
16	#include <string>
17	#include <vector>
18
19	namespace dftracer::utils::utilities::reader {
20
21	namespace indexing = composites::dft::indexing;
22	using common::query::Query;
23	using internal::build_prefilter;
24	using internal::CompiledEqProbe;
25	using internal::eval_compiled_eq;
26	using internal::LinePrefilter;
27	using internal::ondemand_to_literal;
28	using internal::read_chunks_indexed;
29	using internal::strip_ndjson_bookends;
30	using internal::try_compile_eq_probes;
31
32	namespace {
33
34	using common::arrow::ArrowExportResult;
35	using common::arrow::ColumnType;
36	using common::arrow::RecordBatchBuilder;
37
38	struct ArrowKeyHint {	618✔
39	std::string key;
40	std::size_t col_idx = 0;	618✔
41	ColumnType type = ColumnType::INT64;	618✔
42	bool valid = false;	618✔
43	};
44
45	inline std::size_t resolve_col_idx(RecordBatchBuilder& builder,	31,658✔
46	std::vector<ArrowKeyHint>& hints,
47	std::size_t pos, std::string_view key_sv,
48	ColumnType type) {
49	if (pos < hints.size()) {	31,658✔
50	auto& h = hints[pos];	30,421✔
51	if (h.valid && h.type == type && h.key.size() == key_sv.size() &&	45,577!
52	std::memcmp(h.key.data(), key_sv.data(), key_sv.size()) == 0) {	30,427✔
53	return h.col_idx;	30,415✔
54	}
55	}	6✔
56	// Position-keyed miss. Variable-shape rows (e.g., open vs read events
57	// with different args fields) push fields to different positions, so
58	// the position cache misses constantly while the underlying schema is
59	// small (~15 keys). A linear scan over the hint vector with a SIMD
60	// memcmp beats RecordBatchBuilder's name_to_index_ hash lookup for this
61	// size.
62	for (std::size_t i = 0; i < hints.size(); ++i) {	5,626✔
63	if (i == pos) continue;	4,401!
64	auto& h = hints[i];	4,401✔
65	if (h.valid && h.type == type && h.key.size() == key_sv.size() &&	4,687!
66	std::memcmp(h.key.data(), key_sv.data(), key_sv.size()) == 0) {	555!
67	if (pos < hints.size()) {	×
68	auto& slot = hints[pos];	×
69	slot.key.assign(key_sv);	×
70	slot.type = type;	×
71	slot.col_idx = h.col_idx;	×
72	slot.valid = true;	×
73	}
74	return h.col_idx;	×
75	}
76	}	2,254✔
77	std::size_t idx = builder.add_or_get_column(key_sv, type);	1,195✔
78	if (pos >= hints.size()) hints.resize(pos + 1);	1,223✔
79	auto& h = hints[pos];	1,217✔
80	h.key.assign(key_sv);	1,219✔
81	h.type = type;	1,221✔
82	h.col_idx = idx;	1,221✔
83	h.valid = true;	1,221✔
84	return idx;	1,221✔
85	}	15,926✔
86
87	// Append a typed scalar value under `key_sv`. Nested objects/arrays are
88	// always round-tripped as JSON strings (flattening is one level only).
89	void append_scalar_or_json(RecordBatchBuilder& builder,	31,602✔
90	std::vector<ArrowKeyHint>& hints, std::size_t& pos,
91	std::string_view key_sv,
92	simdjson::ondemand::value val,
93	simdjson::ondemand::json_type type) {
94	switch (type) {	31,602!
95	case simdjson::ondemand::json_type::number: {	8,046✔
96	auto num_r = val.get_number();	16,208✔
97	if (num_r.error()) break;	16,203!
98	auto num = num_r.value();	16,185!
99	if (num.is_int64()) {	16,212✔
100	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	16,021!
101	ColumnType::INT64);
102	builder.append_int64(idx, num.get_int64());	16,005!
103	} else if (num.is_uint64()) {	8,249!
104	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	×
105	ColumnType::UINT64);
106	builder.append_uint64(idx, num.get_uint64());	×
107	} else {
108	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	180!
109	ColumnType::DOUBLE);
110	builder.append_double(idx, num.get_double());	180!
111	}
112	break;	16,237✔
113	}
114	case simdjson::ondemand::json_type::string: {	6,315✔
115	auto str_r = val.get_string();	12,708✔
116	if (str_r.error()) break;	12,731!
117	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	12,731!
118	ColumnType::STRING);
119	builder.append_string(idx, str_r.value());	12,731!
120	break;	12,739✔
121	}
122	case simdjson::ondemand::json_type::boolean: {
123	auto b_r = val.get_bool();	×
124	if (b_r.error()) break;	×
125	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	×
126	ColumnType::BOOL);
127	builder.append_bool(idx, b_r.value());	×
128	break;	×
129	}
130	case simdjson::ondemand::json_type::null: {
131	auto existing = builder.find_column(key_sv);	×
132	if (existing) builder.append_null(*existing);	×
133	++pos;	×
134	break;	×
135	}
136	case simdjson::ondemand::json_type::object:	1,357✔
137	case simdjson::ondemand::json_type::array: {
138	auto raw_r = val.raw_json();	2,733✔
139	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	2,731!
140	ColumnType::STRING);
141	if (!raw_r.error()) {	2,727✔
142	auto sv = raw_r.value();	2,725!
143	builder.append_string(idx, sv);	2,726!
144	} else {	1,376✔
145	builder.append_null(idx);	×
146	}
147	break;	2,743✔
148	}
149	default:
150	++pos;	×
151	break;	×
152	}
153	}	31,672✔
154
155	// Append one Arrow row from an already-parsed simdjson document.
156	// Dynamic schema: new columns appended as they appear. When flatten_objects
157	// is true, top-level object values are expanded one level into `parent.child`
158	// columns; deeper nesting still lands as a JSON string under the flattened
159	// key. Returns false on error paths so callers can skip the row.
160	bool arrow_row_from_doc(RecordBatchBuilder& builder,	3,840✔
161	std::vector<ArrowKeyHint>& hints,
162	simdjson::ondemand::document_reference doc,
163	bool flatten_objects = false) {
164	auto obj_result = doc.get_object();	3,840✔
165	if (obj_result.error()) return false;	3,838!
166	char key_buf[512];
167	std::size_t pos = 0;	3,832✔
168	for (auto field : obj_result.value()) {	34,554✔
169	if (field.error()) continue;	30,865!
170	auto key_r = field.unescaped_key();	30,579✔
171	if (key_r.error()) continue;	30,617!
172	auto key_sv = key_r.value();	30,594!
173	auto val_r = field.value();	30,594✔
174	if (val_r.error()) continue;	30,567!
175	auto val = val_r.value();	30,571!
176	auto type_r = val.type();	30,570✔
177	if (type_r.error()) continue;	30,600!
178	auto type = type_r.value();	30,604!
179
180	if (flatten_objects && type == simdjson::ondemand::json_type::object) {	30,578✔
181	auto nested = val.get_object();	562✔
182	if (nested.error()) continue;	562!
183	for (auto nf : nested.value()) {	2,148✔
184	if (nf.error()) continue;	1,586!
185	auto nk_r = nf.unescaped_key();	1,586✔
186	if (nk_r.error()) continue;	1,586!
187	auto nk = nk_r.value();	1,586!
188	auto nv_r = nf.value();	1,586✔
189	if (nv_r.error()) continue;	1,586!
190	auto nv = nv_r.value();	1,586!
191	auto nt_r = nv.type();	1,586✔
192	if (nt_r.error()) continue;	1,586!
193	std::size_t needed = key_sv.size() + 1 + nk.size();	1,586✔
194	if (needed >= sizeof(key_buf)) continue;	1,586!
195	std::memcpy(key_buf, key_sv.data(), key_sv.size());	1,586✔
196	key_buf[key_sv.size()] = '.';	1,586✔
197	std::memcpy(key_buf + key_sv.size() + 1, nk.data(), nk.size());	1,586✔
198	append_scalar_or_json(builder, hints, pos,	2,379!
199	std::string_view(key_buf, needed), nv,	793✔
200	nt_r.value());	1,586!
201	}
202	continue;	562✔
203	}	281✔
204
205	append_scalar_or_json(builder, hints, pos, key_sv, val, type);	30,016!
206	}
207	builder.end_row();	3,805!
208	return true;	3,838✔
209	}	1,928✔
210
211	void collect_query_fields(simdjson::ondemand::document_reference doc,
212	const Query& query, bool check_dotted,
213	common::query::ValueMap& out);
214
215	// Run iterate_many over `padded`, build arrow rows, and emit completed
216	// batches via `yield_one`. Updates `carry` with the truncated tail (if any)
217	// for the caller to prepend to the next chunk.
218	template <typename Yield>
219	void parse_padded_into_arrow(simdjson::ondemand::parser& bulk_parser,
220	simdjson::padded_string& padded,
221	const std::optional<Query>& query, bool flatten,
222	bool query_has_dotted, RecordBatchBuilder& builder,
223	StringArena& arena,
224	std::vector<ArrowKeyHint>& hints,
225	std::size_t batch_size, std::string* carry,
226	Yield&& yield_one) {
227	auto docs_r = bulk_parser.iterate_many(padded, 1 << 20, false);
228	if (docs_r.error()) {
229	if (carry) carry->clear();
230	return;
231	}
232	auto& docs = docs_r.value();
233	for (auto it = docs.begin(); it != docs.end(); ++it) {
234	auto doc_result = *it;
235	if (doc_result.error()) continue;
236	auto& doc = doc_result.value();
237	if (query) {
238	common::query::ValueMap fields;
239	collect_query_fields(doc, *query, query_has_dotted, fields);
240	if (!query->evaluate(fields)) continue;
241	doc.rewind();
242	}
243	if (!arrow_row_from_doc(builder, hints, doc, flatten)) continue;
244	if (builder.num_rows() >= batch_size) {
245	auto result = builder.finish();
246	arena.clear();
247	if (!builder.is_schema_locked()) builder.lock_schema();
248	builder.reset(true);
249	builder.reserve(batch_size);
250	yield_one(std::move(result));
251	}
252	}
253	if (carry) {
254	std::size_t total = padded.size();
255	std::size_t truncated = docs.truncated_bytes();
256	if (truncated > 0 && truncated <= total) {
257	carry->assign(padded.data() + total - truncated,
258	padded.data() + total);
259	} else {
260	carry->clear();
261	}
262	}
263	}
264
265	// Build a simdjson-padded buffer containing only the lines in `chunk` that
266	// pass the line-level prefilter. For queries with no useful prefilter, the
267	// caller should skip this and feed the raw chunk directly.
268	std::string collect_matching_lines(std::span<const char> chunk,	×
269	const LinePrefilter& prefilter) {
270	std::string out;	×
271	out.reserve(chunk.size());	×
272	const char* data = chunk.data();	×
273	std::size_t len = chunk.size();	×
274	std::size_t pos = 0;	×
275	while (pos < len) {	×
276	const void* nl = std::memchr(data + pos, '\n', len - pos);	×
277	std::size_t end_pos = nl ? static_cast<const char*>(nl) - data : len;	×
278	if (end_pos > pos) {	×
279	std::string_view line(data + pos, end_pos - pos);	×
280	if (prefilter.may_match(line)) {	×
281	out.append(line);	×
282	out.push_back('\n');	×
283	}
284	}
285	pos = end_pos + 1;	×
286	}
287	return out;	×
288	}	×
289
290	// Extract fields referenced by the query into a ValueMap, walking one level
291	// of object nesting. Fields not referenced by the query are skipped.
292	void collect_query_fields(simdjson::ondemand::document_reference doc,	×
293	const Query& query, bool check_dotted,
294	common::query::ValueMap& out) {
295	auto obj = doc.get_object();	×
296	if (obj.error()) return;	×
297	for (auto field : obj.value()) {	×
298	if (field.error()) continue;	×
299	auto key_r = field.unescaped_key();	×
300	if (key_r.error()) continue;	×
301	auto val_r = field.value();	×
302	if (val_r.error()) continue;	×
303	auto key = key_r.value();	×
304	auto val = val_r.value();	×
305	auto type_r = val.type();	×
306	if (type_r.error()) continue;	×
307	auto type = type_r.value();	×
308	if (type == simdjson::ondemand::json_type::object) {	×
309	auto nested = val.get_object();	×
310	if (nested.error()) continue;	×
311	for (auto nf : nested.value()) {	×
312	if (nf.error()) continue;	×
313	auto nk_r = nf.unescaped_key();	×
314	if (nk_r.error()) continue;	×
315	auto nv_r = nf.value();	×
316	if (nv_r.error()) continue;	×
NEW 317	internal::store_referenced_nested(out, query, check_dotted, key,	×
NEW 318	nk_r.value(), nv_r.value());	×
319	}
320	} else if (query.references(key)) {	×
321	out[std::string(key)] = ondemand_to_literal(val);	×
322	}
323	}
324	}
325
326	} // namespace
327
328	coro::AsyncGenerator<ArrowExportResult> TraceReader::read_arrow(	3,813!
329	ReadConfig config, std::size_t batch_size) {	79!
330	std::optional<Query> query;	3,310✔
331	if (!config.query.empty()) {	3,310✔
332	auto parsed = Query::from_string(config.query);	7!
333	if (!parsed) throw common::query::QueryParseError(parsed.error());	7!
334	query = std::move(*parsed);	7!
335	}	7✔
336
337	// Resolve virtual fields (resolved./r.) to concrete hash in-clauses via
338	// the index before the query drives pruning or per-event evaluation.
339	if (query && has_index_ && !index_path_.empty() &&	3,310!
340	indexing::has_resolved_fields(*query)) {	×
341	try {
342	indexer::IndexDatabase db(	×
343	index_path_, rocksdb::RocksDatabase::OpenMode::ReadOnly);
344	if (auto rewritten =	×
345	indexing::rewrite_resolved_fields(*query, db)) {	×
346	query = std::move(*rewritten);	×
347	}
348	} catch (...) {
349	}	×
350	}
351
352	// When chunk_prune_only is set, dim_stats already proved every event in
353	// the chunk that has the predicate field matches the literal. We still
354	// need to skip events that lack the field (e.g., metadata "ph":"M"
355	// events lack pid), since the original predicate would reject them.
356	std::vector<std::string> presence_check_paths;	3,310✔
357	if (query && config.chunk_prune_only) {	3,310!
358	const auto& fset = query->fields();	×
359	presence_check_paths.assign(fset.begin(), fset.end());	×
360	}
361
362	const bool query_has_dotted =	3,317✔
363	query && internal::query_references_dotted(*query);	3,310!
364
365	// For AND-of-EQ predicates, evaluate directly against simdjson without
366	// ValueMap (avoids wyhash + per-field std::string allocation per row).
367	// Falls back to the ValueMap path on unsupported AST shapes.
368	std::vector<CompiledEqProbe> compiled_probes;	3,310✔
369	bool use_compiled = false;	3,310✔
370	if (query && !config.chunk_prune_only) {	3,310!
371	if (auto p = try_compile_eq_probes(query->root())) {	14!
372	compiled_probes = std::move(*p);	7✔
373	use_compiled = !compiled_probes.empty();	7✔
374	}	7✔
375	}	7✔
376
377	bool flatten = config.flatten_objects;	3,310✔
378
379	if (!has_index_) {	3,310✔
380	// Fallback: drive the per-line read_json path and build rows.
381	auto json_gen = read_json(config);	3,295!
382	RecordBatchBuilder builder;	3,295!
383	StringArena arena;	3,295!
384	std::vector<ArrowKeyHint> hints;	3,295✔
385	builder.reserve(batch_size);	3,295!
386	while (auto opt = co_await json_gen.next()) {	6,545!
387	if (!arrow_row_from_doc(builder, hints,	1,553!
388	simdjson::ondemand::document_reference(	1,553✔
389	opt->parser->raw_document()),	1,553!
390	flatten))	1,553✔
391	continue;
392	if (builder.num_rows() >= batch_size) {	1,553✔
393	co_yield builder.finish();	70!
394	arena.clear();	35!
395	if (!builder.is_schema_locked()) builder.lock_schema();	35✔
396	builder.reset(true);	35!
397	builder.reserve(batch_size);	35!
398	}	35✔
399	}	1,617✔
400	if (builder.num_rows() > 0) {	62✔
401	co_yield builder.finish();	110!
402	}	54✔
403	co_return;	62✔
404	}	65✔
405
406	// Keep RocksDB alive for the generator's lifetime so per-method opens
407	// in GzipIndexer reuse DBManager's cached handle.
408	std::optional<indexer::IndexDatabase> db_keep_alive;	15✔
409	if (has_index_ && !index_path_.empty()) {	15!
410	try {
411	db_keep_alive.emplace(index_path_,	30!
412	rocksdb::RocksDatabase::OpenMode::ReadOnly);	15✔
413	} catch (...) {	15✔
414	}	×
415	}	15✔
416
417	auto reader = create_indexed_reader();	15!
418	auto chunk_gen = read_chunks_indexed(	30!
419	reader, index_path_, config_.file_path, config, query,	15!
420	/extend_to_line_boundary=/config.end_at_checkpoint);	15✔
421
422	LinePrefilter prefilter = (query && !config.chunk_prune_only)	15!
423	? build_prefilter(*query)	×
424	: LinePrefilter{};	15✔
425	bool have_line_prefilter = !prefilter.empty();	15!
426
427	simdjson::ondemand::parser bulk_parser;	15✔
428	RecordBatchBuilder builder;	15!
429	StringArena arena;	15!
430	std::vector<ArrowKeyHint> hints;	15✔
431	builder.reserve(batch_size);	15!
432
433	auto maybe_flush = [&builder, &arena, batch_size](	795✔
434	bool final) -> std::optional<ArrowExportResult> {	1,269✔
435	if (builder.num_rows() == 0) return std::nullopt;	780✔
436	if (!final && builder.num_rows() < batch_size) return std::nullopt;	778✔
437	auto result = builder.finish();	46!
438	arena.clear();	46!
439	if (!builder.is_schema_locked()) builder.lock_schema();	46✔
440	builder.reset(true);	46!
441	builder.reserve(batch_size);	46!
442	return result;	46!
443	};	413✔
444
445	bool first_chunk = true;	15✔
446	while (auto chunk_opt = co_await chunk_gen.next()) {	116!
447	auto chunk = *chunk_opt;	14✔
448	if (chunk.empty()) continue;	14!
449
450	// Work items with start_byte > 0 begin at a deflate-block boundary
451	// that is typically mid-line; the previous worker emitted that
452	// spanning line via its tail-flush, so drop bytes up to (and
453	// including) the first newline in our first chunk.
454	if (first_chunk && config.start_byte > 0 &&	14!
455	config.start_at_checkpoint) {
456	const char* nl = static_cast<const char*>(
457	std::memchr(chunk.data(), '\n', chunk.size()));	×
458	if (nl) {	×
459	std::size_t skip =
460	static_cast<std::size_t>(nl - chunk.data()) + 1;
461	if (skip < chunk.size()) {	×
462	chunk = chunk.subspan(skip);
463	} else {
464	first_chunk = false;
465	continue;
466	}
467	}	×
468	}	×
469	first_chunk = false;	14✔
470
471	simdjson::padded_string padded;	14✔
472	if (have_line_prefilter) {	14!
473	auto collected = collect_matching_lines(chunk, prefilter);	×
474	if (collected.empty()) continue;	×
475	padded = simdjson::padded_string(std::move(collected));
476	} else {	×
477	auto trimmed = strip_ndjson_bookends(	28!
478	std::string_view(chunk.data(), chunk.size()));	14✔
479	if (trimmed.empty()) continue;	14!
480	padded = simdjson::padded_string(trimmed);	14✔
481	}	14!
482
483	auto docs_r = bulk_parser.iterate_many(padded, 1 << 20,	14✔
484	/allow_comma_separated=/false);
485	if (docs_r.error()) continue;	14!
486	auto& docs = docs_r.value();	14!
487
488	for (auto it = docs.begin(); it != docs.end(); ++it) {	389!
489	auto doc_result = *it;	375✔
490	if (doc_result.error()) continue;	375!
491	auto& doc = doc_result.value();	375!
492
493	if (query && !config.chunk_prune_only) {	375!
494	if (use_compiled) {	×
495	if (!eval_compiled_eq(compiled_probes, doc)) continue;	×
496	} else {
497	common::query::ValueMap fields;	×
498	collect_query_fields(doc, *query, query_has_dotted, fields);	×
499	if (!query->evaluate(fields)) continue;	×
500	}	×
501	doc.rewind();
502	} else if (!presence_check_paths.empty()) {	375!
503	bool all_present = true;
504	for (const auto& path : presence_check_paths) {	×
505	auto fld = doc.find_field_unordered(path);
506	if (fld.error()) {	×
507	all_present = false;
508	break;
509	}
510	}	×
511	if (!all_present) continue;	×
512	doc.rewind();
513	}	×
514
515	if (!arrow_row_from_doc(builder, hints, doc, flatten)) continue;	375!
516
517	if (auto flushed = maybe_flush(/final=/false)) {	384!
518	co_yield std::move(*flushed);	18!
519	}	9✔
520	}	375!
521	}	29!
522
523	if (auto flushed = maybe_flush(/final=/true)) {	29!
524	co_yield std::move(*flushed);	28!
525	}	14✔
526	}	6,874!
527
528	} // namespace dftracer::utils::utilities::reader
529
530	#endif // DFTRACER_UTILS_ENABLE_ARROW

llnl / dftracer-utils / 29811479573

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous