28982597416

Committed 08 Jul 2026 11:22PM UTC coverage: 50.709% (-1.9%) from 52.577%

Build # 28982597416

Build Type

Pull #87

github

Committed by

web-flow

Commit Message

Merge fdf73a826 into 4908d9921

Pull Request Pull Request #87: fix(comparator): fix wrong counting files

Coverage Stats

16259 of 43579 branches covered (37.31%)

Branch coverage included in aggregate %.

37 of 38 new or added lines in 4 files covered. (97.37%)

616 existing lines in 111 files now uncovered.

21634 of 31148 relevant lines covered (69.46%)

13117.41 hits per line

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

44.9

/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 {
37	std::string key;
38	std::size_t col_idx = 0;
39	ColumnType type = ColumnType::INT64;
40	bool valid = false;
41	};
42
43	inline std::size_t resolve_col_idx(RecordBatchBuilder& builder,	15,708✔
44	std::vector<ArrowKeyHint>& hints,
45	std::size_t pos, std::string_view key_sv,
46	ColumnType type) {
47	if (pos < hints.size()) {	15,708✔
48	auto& h = hints[pos];	15,087✔
49	if (h.valid && h.type == type && h.key.size() == key_sv.size() &&	30,192!
50	std::memcmp(h.key.data(), key_sv.data(), key_sv.size()) == 0) {	15,086✔
51	return h.col_idx;	15,080✔
52	}
53	}
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,815✔
61	if (i == pos) continue;	2,207!
62	auto& h = hints[i];	2,207✔
63	if (h.valid && h.type == type && h.key.size() == key_sv.size() &&	2,482!
64	std::memcmp(h.key.data(), key_sv.data(), key_sv.size()) == 0) {	276!
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;	×
71	}
72	return h.col_idx;	×
73	}
74	}
75	std::size_t idx = builder.add_or_get_column(key_sv, type);	599✔
76	if (pos >= hints.size()) hints.resize(pos + 1);	610!
77	auto& h = hints[pos];	605✔
78	h.key.assign(key_sv);	603✔
79	h.type = type;	607✔
80	h.col_idx = idx;	607✔
81	h.valid = true;	607✔
82	return idx;	607✔
83	}
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,652✔
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,652!
93	case simdjson::ondemand::json_type::number: {	8,040✔
94	auto num_r = val.get_number();	8,040✔
95	if (num_r.error()) break;	8,029!
96	auto num = num_r.value();	8,028!
97	if (num.is_int64()) {	8,035✔
98	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	7,941!
99	ColumnType::INT64);
100	builder.append_int64(idx, num.get_int64());	7,922!
101	} else if (num.is_uint64()) {	96!
102	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	×
103	ColumnType::UINT64);
104	builder.append_uint64(idx, num.get_uint64());	×
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,073✔
111	}
112	case simdjson::ondemand::json_type::string: {	6,314✔
113	auto str_r = val.get_string();	6,314✔
114	if (str_r.error()) break;	6,343!
115	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	6,346!
116	ColumnType::STRING);
117	builder.append_string(idx, str_r.value());	6,345!
118	break;	6,340✔
119	}
UNCOV 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	}
UNCOV 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:	1,361✔
135	case simdjson::ondemand::json_type::array: {
136	auto raw_r = val.raw_json();	1,361✔
137	auto idx = resolve_col_idx(builder, hints, pos++, key_sv,	1,350!
138	ColumnType::STRING);
139	if (!raw_r.error()) {	1,351!
140	auto sv = raw_r.value();	1,352!
141	builder.append_string(idx, sv);	1,348!
142	} else {
143	builder.append_null(idx);	×
144	}
145	break;	1,364✔
146	}
UNCOV 147	default:	×
148	++pos;	×
149	break;	×
150	}
151	}	15,714✔
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,911✔
159	std::vector<ArrowKeyHint>& hints,
160	simdjson::ondemand::document_reference doc,
161	bool flatten_objects = false) {
162	auto obj_result = doc.get_object();	1,911✔
163	if (obj_result.error()) return false;	1,910!
164	char key_buf[512];
165	std::size_t pos = 0;	1,912✔
166	for (auto field : obj_result.value()) {	17,181!
167	if (field.error()) continue;	15,495!
168	auto key_r = field.unescaped_key();	15,216✔
169	if (key_r.error()) continue;	15,204!
170	auto key_sv = key_r.value();	15,195!
171	auto val_r = field.value();	15,193✔
172	if (val_r.error()) continue;	15,150!
173	auto val = val_r.value();	15,141!
174	auto type_r = val.type();	15,148✔
175	if (type_r.error()) continue;	15,193!
176	auto type = type_r.value();	15,183!
177
178	if (flatten_objects && type == simdjson::ondemand::json_type::object) {	15,161✔
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,	793!
197	std::string_view(key_buf, needed), nv,
198	nt_r.value());	793!
199	}
200	continue;	281✔
201	}	281✔
202
203	append_scalar_or_json(builder, hints, pos, key_sv, val, type);	14,880!
204	}
205	builder.end_row();	1,888!
206	return true;	1,902✔
207	}
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');	×
279	}
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);	×
318	}
319	}
320	}
321
322	} // namespace
323
324	coro::AsyncGenerator<ArrowExportResult> TraceReader::read_arrow(	79!
325	ReadConfig config, std::size_t batch_size) {
326	std::optional<Query> query;
327	if (!config.query.empty()) {
328	auto parsed = Query::from_string(config.query);
329	if (!parsed) throw common::query::QueryParseError(parsed.error());
330	query = std::move(*parsed);
331	}
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;
338	if (query && config.chunk_prune_only) {
339	const auto& fset = query->fields();
340	presence_check_paths.assign(fset.begin(), fset.end());
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;
347	bool use_compiled = false;
348	if (query && !config.chunk_prune_only) {
349	if (auto p = try_compile_eq_probes(query->root())) {
350	compiled_probes = std::move(*p);
351	use_compiled = !compiled_probes.empty();
352	}
353	}
354
355	bool flatten = config.flatten_objects;
356
357	if (!has_index_) {
358	// Fallback: drive the per-line read_json path and build rows.
359	auto json_gen = read_json(config);
360	RecordBatchBuilder builder;
361	StringArena arena;
362	std::vector<ArrowKeyHint> hints;
363	builder.reserve(batch_size);
364	while (auto opt = co_await json_gen.next()) {
365	if (!arrow_row_from_doc(builder, hints,
366	simdjson::ondemand::document_reference(
367	opt->parser->raw_document()),
368	flatten))
369	continue;
370	if (builder.num_rows() >= batch_size) {
371	co_yield builder.finish();
372	arena.clear();
373	if (!builder.is_schema_locked()) builder.lock_schema();
374	builder.reset(true);
375	builder.reserve(batch_size);
376	}
377	}
378	if (builder.num_rows() > 0) {
379	co_yield builder.finish();
380	}
381	co_return;
382	}
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;
387	if (has_index_ && !index_path_.empty()) {
388	try {
389	db_keep_alive.emplace(index_path_,
390	rocksdb::RocksDatabase::OpenMode::ReadOnly);
391	} catch (...) {
392	}
393	}
394
395	auto reader = create_indexed_reader();
396	auto chunk_gen = read_chunks_indexed(
397	reader, index_path_, config_.file_path, config, query,
398	/extend_to_line_boundary=/config.end_at_checkpoint);
399
400	LinePrefilter prefilter = (query && !config.chunk_prune_only)
401	? build_prefilter(*query)
402	: LinePrefilter{};
403	bool have_line_prefilter = !prefilter.empty();
404
405	simdjson::ondemand::parser bulk_parser;
406	RecordBatchBuilder builder;
407	StringArena arena;
408	std::vector<ArrowKeyHint> hints;
409	builder.reserve(batch_size);
410
411	auto maybe_flush = [&builder, &arena, batch_size](	389✔
412	bool final) -> std::optional<ArrowExportResult> {	1,262✔
413	if (builder.num_rows() == 0) return std::nullopt;	389✔
414	if (!final && builder.num_rows() < batch_size) return std::nullopt;	386✔
415	auto result = builder.finish();	23!
416	arena.clear();	23!
417	if (!builder.is_schema_locked()) builder.lock_schema();	23✔
418	builder.reset(true);	23!
419	builder.reserve(batch_size);	23!
420	return result;	23!
421	};	23✔
422
423	bool first_chunk = true;
424	while (auto chunk_opt = co_await chunk_gen.next()) {
425	auto chunk = *chunk_opt;
426	if (chunk.empty()) continue;
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 &&
433	config.start_at_checkpoint) {
434	const char* nl = static_cast<const char*>(
435	std::memchr(chunk.data(), '\n', chunk.size()));
436	if (nl) {
437	std::size_t skip =
438	static_cast<std::size_t>(nl - chunk.data()) + 1;
439	if (skip < chunk.size()) {
440	chunk = chunk.subspan(skip);
441	} else {
442	first_chunk = false;
443	continue;
444	}
445	}
446	}
447	first_chunk = false;
448
449	simdjson::padded_string padded;
450	if (have_line_prefilter) {
451	auto collected = collect_matching_lines(chunk, prefilter);
452	if (collected.empty()) continue;
453	padded = simdjson::padded_string(std::move(collected));
454	} else {
455	auto trimmed = strip_ndjson_bookends(
456	std::string_view(chunk.data(), chunk.size()));
457	if (trimmed.empty()) continue;
458	padded = simdjson::padded_string(trimmed);
459	}
460
461	auto docs_r = bulk_parser.iterate_many(padded, 1 << 20,
462	/allow_comma_separated=/false);
463	if (docs_r.error()) continue;
464	auto& docs = docs_r.value();
465
466	for (auto it = docs.begin(); it != docs.end(); ++it) {
467	auto doc_result = *it;
468	if (doc_result.error()) continue;
469	auto& doc = doc_result.value();
470
471	if (query && !config.chunk_prune_only) {
472	if (use_compiled) {
473	if (!eval_compiled_eq(compiled_probes, doc)) continue;
474	} else {
475	common::query::ValueMap fields;
476	collect_query_fields(doc, *query, fields);
477	if (!query->evaluate(fields)) continue;
478	}
479	doc.rewind();
480	} else if (!presence_check_paths.empty()) {
481	bool all_present = true;
482	for (const auto& path : presence_check_paths) {
483	auto fld = doc.find_field_unordered(path);
484	if (fld.error()) {
485	all_present = false;
486	break;
487	}
488	}
489	if (!all_present) continue;
490	doc.rewind();
491	}
492
493	if (!arrow_row_from_doc(builder, hints, doc, flatten)) continue;
494
495	if (auto flushed = maybe_flush(/final=/false)) {
496	co_yield std::move(*flushed);
497	}
498	}
499	}
500
501	if (auto flushed = maybe_flush(/final=/true)) {
502	co_yield std::move(*flushed);
503	}
504	}	158!
505
506	} // namespace dftracer::utils::utilities::reader
507
508	#endif // DFTRACER_UTILS_ENABLE_ARROW

llnl / dftracer-utils / 28982597416

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