29187058583

Committed 12 Jul 2026 09:11AM UTC coverage: 51.324% (-1.4%) from 52.754%

Build # 29187058583

Build Type

Pull #96

github

Committed by

web-flow

Commit Message

Merge 7e90b76dd into 056c79287

Pull Request Pull Request #96: fix: consolidate stale-index rebuilds across consumers and fix multi-run breaks

Coverage Stats

33807 of 84432 branches covered (40.04%)

Branch coverage included in aggregate %.

145 of 152 new or added lines in 16 files covered. (95.39%)

5186 existing lines in 197 files now uncovered.

34558 of 48770 relevant lines covered (70.86%)

10797.15 hits per line

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

53.13

/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/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 {

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, 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,
                             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, 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, 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;
                if (!query.references(nk_r.value())) continue;
                out[std::string(nk_r.value())] =
                    ondemand_to_literal(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);
    }

    // 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());
    }

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

llnl / dftracer-utils / 29187058583

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