29787659139

Committed 20 Jul 2026 11:34PM UTC coverage: 51.578% (-1.1%) from 52.66%

Build # 29787659139

Build Type

Pull #99

github

Committed by

web-flow

Commit Message

Merge ee805adeb into 06bc84ec9

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

Coverage Stats

34832 of 86278 branches covered (40.37%)

Branch coverage included in aggregate %.

1034 of 1319 new or added lines in 30 files covered. (78.39%)

5193 existing lines in 197 files now uncovered.

35349 of 49791 relevant lines covered (70.99%)

9765.54 hits per line

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

50.37

/src/dftracer/utils/utilities/reader/internal/chunk_geometry.cpp

#include <dftracer/utils/utilities/reader/internal/chunk_geometry.h>

#ifdef DFTRACER_UTILS_ENABLE_ARROW

#include <dftracer/utils/core/common/filesystem.h>
#include <dftracer/utils/core/rocksdb/database.h>
#include <dftracer/utils/utilities/common/query/query.h>
#include <dftracer/utils/utilities/composites/dft/indexing/chunk_pruner_utility.h>
#include <dftracer/utils/utilities/composites/dft/internal/utils.h>
#include <dftracer/utils/utilities/indexer/index_database.h>
#include <dftracer/utils/utilities/indexer/internal/checkpoint.h>
#include <dftracer/utils/utilities/indexer/internal/helpers.h>

#include <algorithm>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

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

namespace {

// User-supplied byte/line clip applied to every emitted work item.
struct ClipRange {
    std::size_t start_byte = 0, end_byte = 0;
    std::size_t start_line = 0, end_line = 0;
    bool has_line_clip() const { return start_line > 0 || end_line > 0; }
    bool has_byte_clip() const { return end_byte > start_byte; }
};

// Geometry of the pruner chunks for one file, derived from its gzip recovery
// points. Pruner chunk_idx is 0-indexed over uncompressed slices: recovery
// point ckpts[k] sits at the START of chunk (k+1); chunk 0 has no recovery
// point at its start (decoded from gzip stream start). Total chunks =
// ckpts.size()+1.
struct ChunkGeometry {
    const std::vector<dftracer::utils::utilities::indexer::IndexerCheckpoint>
        &ckpts;

    std::size_t total_chunks() const { return ckpts.size() + 1; }
    std::size_t start_byte(std::uint64_t cidx) const {
        if (cidx == 0) return 0;
        return ckpts[cidx - 1].uc_offset;
    }
    std::size_t end_byte(std::uint64_t cidx) const {
        if (cidx == 0) return ckpts.empty() ? 0 : ckpts[0].uc_offset;
        std::size_t k = cidx - 1;
        return ckpts[k].uc_offset + ckpts[k].uc_size;
    }
    // Chunk 0 covers everything before the first recovery point; chunk k>=1
    // spans recovery point (k-1).
    std::size_t first_line(std::uint64_t cidx) const {
        if (cidx == 0) return 1;
        return ckpts[cidx - 1].first_line_num;
    }
    std::size_t last_line(std::uint64_t cidx) const {
        if (cidx == 0) {
            if (ckpts.empty()) return SIZE_MAX;
            return ckpts[0].first_line_num > 0 ? ckpts[0].first_line_num - 1
                                               : 0;
        }
        return ckpts[cidx - 1].last_line_num;
    }
};

// Extract AND-of-EQ leaves from a Query AST. Returns nullopt if the predicate
// shape is anything else (NE, range ops, IN, NOT, OR), in which case the
// uniform-match shortcut does not apply.
std::optional<std::vector<std::pair<std::string, std::string>>>
extract_eq_leaves(
    const dftracer::utils::utilities::common::query::QueryNode &node) {
    namespace q_ns = dftracer::utils::utilities::common::query;
    using LeafVec = std::vector<std::pair<std::string, std::string>>;

    auto literal_to_string = [](const q_ns::LiteralNode &lit) -> std::string {
        return std::visit(
            [](auto &&v) -> std::string {
                using T = std::decay_t<decltype(v)>;
                if constexpr (std::is_same_v<T, std::string>)
                    return v;
                else if constexpr (std::is_same_v<T, bool>)
                    return v ? "true" : "false";
                else if constexpr (std::is_same_v<T, std::int64_t>)
                    return std::to_string(v);
                else if constexpr (std::is_same_v<T, std::uint64_t>)
                    return std::to_string(v);
                else if constexpr (std::is_same_v<T, double>)
                    return std::to_string(v);
                else
                    return {};
            },
            lit.value);
    };

    return std::visit(
        [&](const auto &n) -> std::optional<LeafVec> {
            using T = std::decay_t<decltype(n)>;
            if constexpr (std::is_same_v<T, q_ns::CompareNode>) {
                if (n.op != q_ns::CompareOp::EQ) return std::nullopt;
                return LeafVec{{n.field.path, literal_to_string(n.value)}};
            } else if constexpr (std::is_same_v<T, q_ns::AndNode>) {
                auto l = extract_eq_leaves(*n.left);
                if (!l) return std::nullopt;
                auto r = extract_eq_leaves(*n.right);
                if (!r) return std::nullopt;
                l->insert(l->end(), r->begin(), r->end());
                return l;
            } else {
                return std::nullopt;
            }
        },
        node.data);
}

// True iff every checkpoint in `chunk_idxs` has dim_stats min == max == literal
// for every leaf. Empty leaves -> false (no shortcut). Missing dim_stats for
// any (chunk, leaf) -> false (we don't know, play safe).
bool all_chunks_uniform_match(
    const dftracer::utils::utilities::indexer::IndexDatabase &db, int fid,
    const std::vector<std::pair<std::string, std::string>> &leaves,
    const std::vector<std::uint64_t> &chunk_idxs) {
    if (leaves.empty() || chunk_idxs.empty()) return false;
    namespace indexing = dftracer::utils::utilities::composites::dft::indexing;

    for (const auto &[dim, val] : leaves) {
        auto rows = db.query_chunk_dimension_stats_for_dimension(fid, dim);
        if (rows.empty()) return false;
        std::unordered_map<std::uint64_t,
                           const indexing::ChunkDimensionStatsResult *>
            by_ckpt;
        by_ckpt.reserve(rows.size());
        for (const auto &r : rows) by_ckpt.emplace(r.checkpoint_idx, &r);
        for (auto cidx : chunk_idxs) {
            auto it = by_ckpt.find(cidx);
            if (it == by_ckpt.end()) return false;
            const auto &ds = *it->second;
            if (ds.min_value != val || ds.max_value != val) return false;
        }
    }
    return true;
}

// Select the chunk indices to scan for one file. Returns empty when the file
// is fully pruned or no chunk overlaps the line clip (caller skips the file).
std::vector<std::uint64_t> select_kept_chunks(
    const ChunkGeometry &geo, bool has_query,
    const dftracer::utils::utilities::composites::dft::indexing::
        ChunkPrunerOutput &pr,
    const ClipRange &clip) {
    const std::size_t total_chunks = geo.total_chunks();
    std::vector<std::uint64_t> keep_chunks;
    keep_chunks.reserve(total_chunks);
    if (has_query) {
        if (pr.success && !pr.file_may_match) {
            return keep_chunks;  // whole file pruned
        }
        if (pr.success && !pr.candidate_checkpoints.empty() &&
            pr.candidate_checkpoints.size() < pr.total_checkpoints) {
            for (auto cidx : pr.candidate_checkpoints) {
                if (cidx < total_chunks) keep_chunks.push_back(cidx);
            }
            std::sort(keep_chunks.begin(), keep_chunks.end());
            keep_chunks.erase(
                std::unique(keep_chunks.begin(), keep_chunks.end()),
                keep_chunks.end());
        } else {
            for (std::uint64_t c = 0; c < total_chunks; ++c)
                keep_chunks.push_back(c);
        }
    } else {
        for (std::uint64_t c = 0; c < total_chunks; ++c)
            keep_chunks.push_back(c);
    }

    // Intersect with the user's line range so workers only touch chunks that
    // actually overlap it. Each work item carries the sub-line-range;
    // LINE_RANGE on the read maps it back to bytes via the same checkpoint
    // table the gzip stream uses.
    if (clip.has_line_clip()) {
        std::size_t lo = clip.start_line > 0 ? clip.start_line : 1;
        std::size_t hi = clip.end_line > 0 ? clip.end_line : SIZE_MAX;
        std::vector<std::uint64_t> filtered;
        filtered.reserve(keep_chunks.size());
        for (auto c : keep_chunks) {
            std::size_t cf = geo.first_line(c);
            std::size_t cl = geo.last_line(c);
            if (cl < lo || cf > hi) continue;
            filtered.push_back(c);
        }
        keep_chunks = std::move(filtered);
    }
    return keep_chunks;
}

// Group the kept chunks into contiguous work ranges (one per worker slot) and
// append the resulting work items.
void emit_work_items(std::vector<ArrowWorkItem> &items, const std::string &fp,
                     const ChunkGeometry &geo,
                     const std::vector<std::uint64_t> &keep_chunks,
                     bool file_pure_match, std::size_t max_workers,
                     const ClipRange &clip) {
    std::size_t target_ranges = std::max<std::size_t>(1, max_workers);
    std::size_t per_range = std::max<std::size_t>(
        1, (keep_chunks.size() + target_ranges - 1) / target_ranges);

    std::size_t group_start = 0;
    while (group_start < keep_chunks.size()) {
        std::size_t group_end = group_start;
        std::size_t emitted = 0;
        while (group_end < keep_chunks.size() && emitted < per_range) {
            if (group_end > group_start &&
                keep_chunks[group_end] != keep_chunks[group_end - 1] + 1) {
                break;
            }
            ++group_end;
            ++emitted;
        }
        std::uint64_t scidx = keep_chunks[group_start];
        std::uint64_t ecidx = keep_chunks[group_end - 1];
        std::size_t start_byte = geo.start_byte(scidx);
        std::size_t end_byte = geo.end_byte(ecidx);
        // start_at_checkpoint: a gzip recovery point sits at start_byte (true
        // for any cidx>=1; false for the implicit chunk 0 which decodes from
        // stream start).
        bool start_at_checkpoint = (scidx >= 1);
        bool end_at_checkpoint = (group_end < keep_chunks.size());
        if (clip.has_line_clip()) {
            std::size_t lo = clip.start_line > 0 ? clip.start_line : 1;
            std::size_t hi = clip.end_line > 0 ? clip.end_line : SIZE_MAX;
            std::size_t cluster_first = geo.first_line(scidx);
            std::size_t cluster_last = geo.last_line(ecidx);
            std::size_t item_start = std::max<std::size_t>(lo, cluster_first);
            std::size_t item_end = std::min<std::size_t>(hi, cluster_last);
            if (item_start > item_end) {
                group_start = group_end;
                continue;
            }
            ArrowWorkItem item;
            item.file_path = fp;
            item.chunk_prune_only = file_pure_match;
            item.start_line = item_start;
            item.end_line = item_end;
            items.push_back(std::move(item));
            group_start = group_end;
            continue;
        }
        if (clip.has_byte_clip()) {
            if (start_byte < clip.start_byte) {
                start_byte = clip.start_byte;
                start_at_checkpoint = false;
            }
            if (end_byte > clip.end_byte) {
                end_byte = clip.end_byte;
                end_at_checkpoint = false;
            }
            if (start_byte >= end_byte) {
                group_start = group_end;
                continue;
            }
        }
        items.push_back({fp, start_byte, end_byte, start_at_checkpoint,
                         end_at_checkpoint, file_pure_match});
        group_start = group_end;
    }
}

}  // namespace

std::vector<ArrowWorkItem> enumerate_work_items(
    const std::vector<std::string> &files, const std::string &index_dir,
    const std::string &query_str, std::size_t max_workers,
    std::size_t clip_start_byte, std::size_t clip_end_byte,
    std::size_t clip_start_line, std::size_t clip_end_line) {
    namespace dft_internal =
        dftracer::utils::utilities::composites::dft::internal;
    namespace indexer_ns = dftracer::utils::utilities::indexer;
    namespace indexing = dftracer::utils::utilities::composites::dft::indexing;

    const ClipRange clip{clip_start_byte, clip_end_byte, clip_start_line,
                         clip_end_line};

    std::vector<ArrowWorkItem> items;
    items.reserve(files.size() * 4);

    auto push_unsplit = [&](const std::string &fp) {
        ArrowWorkItem item;
        item.file_path = fp;
        item.start_line = clip.start_line;
        item.end_line = clip.end_line;
        items.push_back(std::move(item));
    };

    // Parse the query once. Pruner input copies a Query, so we keep the
    // parsed form around to feed each ChunkPrunerInput without re-parsing.
    std::optional<dftracer::utils::utilities::common::query::Query> parsed;
    if (!query_str.empty()) {
        auto r = dftracer::utils::utilities::common::query::Query::from_string(
            query_str);
        if (r) parsed = std::move(*r);
    }

    // All files in a directory-mode scan share the same `.dftindex` root.
    // Group files by their resolved index path so we can open the RocksDB
    // once per index and reuse it to prune every file against that handle.
    std::unordered_map<std::string, std::vector<std::size_t>> by_index;
    for (std::size_t i = 0; i < files.size(); ++i) {
        std::string index_path =
            dft_internal::determine_index_path(files[i], index_dir);
        by_index[index_path].push_back(i);
    }

    for (auto &entry : by_index) {
        const auto &index_path = entry.first;
        const auto &file_idxs = entry.second;
        if (!fs::exists(index_path)) {
            for (auto i : file_idxs) push_unsplit(files[i]);
            continue;
        }
        std::unique_ptr<indexer_ns::IndexDatabase> idx_db;
        try {
            idx_db = std::make_unique<indexer_ns::IndexDatabase>(
                index_path,
                dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);
        } catch (...) {
            for (auto i : file_idxs) push_unsplit(files[i]);
            continue;
        }

        // Resolve fid + checkpoints per file (cheap queries).
        struct FileCtx {
            std::size_t file_idx;
            int fid;
            std::vector<indexer_ns::IndexerCheckpoint> ckpts;
        };
        std::vector<FileCtx> file_ctxs;
        file_ctxs.reserve(file_idxs.size());
        for (auto i : file_idxs) {
            FileCtx fc;
            fc.file_idx = i;
            fc.fid = idx_db->get_file_info_id(
                indexer_ns::internal::get_logical_path(files[i]));
            if (fc.fid < 0) {
                push_unsplit(files[i]);
                continue;
            }
            fc.ckpts = idx_db->query_checkpoints(fc.fid);
            if (fc.ckpts.empty()) {
                push_unsplit(files[i]);
                continue;
            }
            std::sort(fc.ckpts.begin(), fc.ckpts.end(),
                      [](const auto &a, const auto &b) {
                          return a.first_line_num < b.first_line_num;
                      });
            file_ctxs.push_back(std::move(fc));
        }

        // Batch-prune all files against the shared index: dim_stats and
        // chunk_statistics are loaded in one RocksDB scan each instead of
        // one scan per file.
        std::vector<indexing::ChunkPrunerOutput> pruner_outs(file_ctxs.size());
        if (parsed && !file_ctxs.empty()) {
            indexing::ChunkPrunerBatchInput batch_in;
            batch_in.index_path = index_path;
            batch_in.external_db = idx_db.get();
            batch_in.items.reserve(file_ctxs.size());
            for (auto &fc : file_ctxs) {
                batch_in.items.push_back({files[fc.file_idx], *parsed});
            }
            indexing::ChunkPrunerUtility pruner;
            auto batch_out = pruner.process_batch(batch_in);
            if (batch_out) {
                pruner_outs = std::move(batch_out->outputs);
            }
        }

        // For AND-of-EQ predicates, precompute uniform-match leaves once.
        // Per-file pure_match is checked inline below and lets workers skip
        // per-event predicate eval on chunks where dim_stats min == max ==
        // literal for every leaf.
        std::optional<std::vector<std::pair<std::string, std::string>>>
            eq_leaves;
        if (parsed) eq_leaves = extract_eq_leaves(parsed->root());

        for (std::size_t fc_idx = 0; fc_idx < file_ctxs.size(); ++fc_idx) {
            auto &fc = file_ctxs[fc_idx];
            const auto &fp = files[fc.file_idx];

            ChunkGeometry geo{fc.ckpts};
            std::vector<std::uint64_t> keep_chunks = select_kept_chunks(
                geo, parsed.has_value(), pruner_outs[fc_idx], clip);
            if (keep_chunks.empty()) continue;

            // All-or-nothing per file: if every kept chunk is uniform-matching
            // for every leaf, every work item from this file gets the
            // chunk_prune_only fast path. Mixed files fall back to per-event
            // eval to stay safe.
            bool file_pure_match = false;
            if (eq_leaves && !eq_leaves->empty() && idx_db) {
                file_pure_match = all_chunks_uniform_match(
                    *idx_db, fc.fid, *eq_leaves, keep_chunks);
            }

            emit_work_items(items, fp, geo, keep_chunks, file_pure_match,
                            max_workers, clip);
        }
    }
    return items;
}

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

#endif  // DFTRACER_UTILS_ENABLE_ARROW

1	#include <dftracer/utils/utilities/reader/internal/chunk_geometry.h>
2
3	#ifdef DFTRACER_UTILS_ENABLE_ARROW
4
5	#include <dftracer/utils/core/common/filesystem.h>
6	#include <dftracer/utils/core/rocksdb/database.h>
7	#include <dftracer/utils/utilities/common/query/query.h>
8	#include <dftracer/utils/utilities/composites/dft/indexing/chunk_pruner_utility.h>
9	#include <dftracer/utils/utilities/composites/dft/internal/utils.h>
10	#include <dftracer/utils/utilities/indexer/index_database.h>
11	#include <dftracer/utils/utilities/indexer/internal/checkpoint.h>
12	#include <dftracer/utils/utilities/indexer/internal/helpers.h>
13
14	#include <algorithm>
15	#include <cstdint>
16	#include <memory>
17	#include <optional>
18	#include <string>
19	#include <unordered_map>
20	#include <utility>
21	#include <vector>
22
23	namespace dftracer::utils::utilities::reader::internal {
24
25	namespace {
26
27	// User-supplied byte/line clip applied to every emitted work item.
28	struct ClipRange {
29	std::size_t start_byte = 0, end_byte = 0;
30	std::size_t start_line = 0, end_line = 0;
31	bool has_line_clip() const { return start_line > 0 \|\| end_line > 0; }	10✔
32	bool has_byte_clip() const { return end_byte > start_byte; }	4✔
33	};
34
35	// Geometry of the pruner chunks for one file, derived from its gzip recovery
36	// points. Pruner chunk_idx is 0-indexed over uncompressed slices: recovery
37	// point ckpts[k] sits at the START of chunk (k+1); chunk 0 has no recovery
38	// point at its start (decoded from gzip stream start). Total chunks =
39	// ckpts.size()+1.
40	struct ChunkGeometry {
41	const std::vector<dftracer::utils::utilities::indexer::IndexerCheckpoint>
42	&ckpts;
43
44	std::size_t total_chunks() const { return ckpts.size() + 1; }	4✔
45	std::size_t start_byte(std::uint64_t cidx) const {	6✔
46	if (cidx == 0) return 0;	6✔
47	return ckpts[cidx - 1].uc_offset;	4✔
48	}	6✔
49	std::size_t end_byte(std::uint64_t cidx) const {	6✔
50	if (cidx == 0) return ckpts.empty() ? 0 : ckpts[0].uc_offset;	6!
51	std::size_t k = cidx - 1;	4✔
52	return ckpts[k].uc_offset + ckpts[k].uc_size;	4✔
53	}	6✔
54	// Chunk 0 covers everything before the first recovery point; chunk k>=1
55	// spans recovery point (k-1).
56	std::size_t first_line(std::uint64_t cidx) const {	6✔
57	if (cidx == 0) return 1;	6✔
58	return ckpts[cidx - 1].first_line_num;	4✔
59	}	6✔
60	std::size_t last_line(std::uint64_t cidx) const {	6✔
61	if (cidx == 0) {	6✔
62	if (ckpts.empty()) return SIZE_MAX;	2!
63	return ckpts[0].first_line_num > 0 ? ckpts[0].first_line_num - 1	2!
64	: 0;
65	}
66	return ckpts[cidx - 1].last_line_num;	4✔
67	}	6✔
68	};
69
70	// Extract AND-of-EQ leaves from a Query AST. Returns nullopt if the predicate
71	// shape is anything else (NE, range ops, IN, NOT, OR), in which case the
72	// uniform-match shortcut does not apply.
73	std::optional<std::vector<std::pair<std::string, std::string>>>
74	extract_eq_leaves(	×
75	const dftracer::utils::utilities::common::query::QueryNode &node) {
76	namespace q_ns = dftracer::utils::utilities::common::query;
77	using LeafVec = std::vector<std::pair<std::string, std::string>>;
78
79	auto literal_to_string = [](const q_ns::LiteralNode &lit) -> std::string {	×
UNCOV 80	return std::visit(	×
81	[](auto &&v) -> std::string {	×
82	using T = std::decay_t<decltype(v)>;
83	if constexpr (std::is_same_v<T, std::string>)
84	return v;	×
85	else if constexpr (std::is_same_v<T, bool>)
86	return v ? "true" : "false";	×
87	else if constexpr (std::is_same_v<T, std::int64_t>)
88	return std::to_string(v);	×
89	else if constexpr (std::is_same_v<T, std::uint64_t>)
90	return std::to_string(v);	×
91	else if constexpr (std::is_same_v<T, double>)
92	return std::to_string(v);	×
93	else
94	return {};
95	},
96	lit.value);	×
97	};
98
UNCOV 99	return std::visit(	×
100	[&](const auto &n) -> std::optional<LeafVec> {	×
101	using T = std::decay_t<decltype(n)>;
102	if constexpr (std::is_same_v<T, q_ns::CompareNode>) {
103	if (n.op != q_ns::CompareOp::EQ) return std::nullopt;	×
104	return LeafVec{{n.field.path, literal_to_string(n.value)}};	×
105	} else if constexpr (std::is_same_v<T, q_ns::AndNode>) {
106	auto l = extract_eq_leaves(*n.left);	×
107	if (!l) return std::nullopt;	×
108	auto r = extract_eq_leaves(*n.right);	×
109	if (!r) return std::nullopt;	×
110	l->insert(l->end(), r->begin(), r->end());	×
111	return l;	×
112	} else {	×
113	return std::nullopt;	×
114	}
UNCOV 115	},	×
116	node.data);	×
117	}
118
119	// True iff every checkpoint in `chunk_idxs` has dim_stats min == max == literal
120	// for every leaf. Empty leaves -> false (no shortcut). Missing dim_stats for
121	// any (chunk, leaf) -> false (we don't know, play safe).
122	bool all_chunks_uniform_match(	×
123	const dftracer::utils::utilities::indexer::IndexDatabase &db, int fid,
124	const std::vector<std::pair<std::string, std::string>> &leaves,
125	const std::vector<std::uint64_t> &chunk_idxs) {
126	if (leaves.empty() \|\| chunk_idxs.empty()) return false;	×
127	namespace indexing = dftracer::utils::utilities::composites::dft::indexing;
128
129	for (const auto &[dim, val] : leaves) {	×
130	auto rows = db.query_chunk_dimension_stats_for_dimension(fid, dim);	×
131	if (rows.empty()) return false;	×
132	std::unordered_map<std::uint64_t,
133	const indexing::ChunkDimensionStatsResult *>
134	by_ckpt;	×
135	by_ckpt.reserve(rows.size());	×
136	for (const auto &r : rows) by_ckpt.emplace(r.checkpoint_idx, &r);	×
137	for (auto cidx : chunk_idxs) {	×
138	auto it = by_ckpt.find(cidx);	×
139	if (it == by_ckpt.end()) return false;	×
140	const auto &ds = *it->second;	×
141	if (ds.min_value != val \|\| ds.max_value != val) return false;	×
142	}
143	}	×
144	return true;	×
UNCOV 145	}	×
146
147	// Select the chunk indices to scan for one file. Returns empty when the file
148	// is fully pruned or no chunk overlaps the line clip (caller skips the file).
149	std::vector<std::uint64_t> select_kept_chunks(	4✔
150	const ChunkGeometry &geo, bool has_query,
151	const dftracer::utils::utilities::composites::dft::indexing::
152	ChunkPrunerOutput &pr,
153	const ClipRange &clip) {
154	const std::size_t total_chunks = geo.total_chunks();	4✔
155	std::vector<std::uint64_t> keep_chunks;	4✔
156	keep_chunks.reserve(total_chunks);	4!
157	if (has_query) {	4!
158	if (pr.success && !pr.file_may_match) {	×
159	return keep_chunks; // whole file pruned	×
160	}
161	if (pr.success && !pr.candidate_checkpoints.empty() &&	×
162	pr.candidate_checkpoints.size() < pr.total_checkpoints) {	×
163	for (auto cidx : pr.candidate_checkpoints) {	×
164	if (cidx < total_chunks) keep_chunks.push_back(cidx);	×
165	}
166	std::sort(keep_chunks.begin(), keep_chunks.end());	×
167	keep_chunks.erase(	×
168	std::unique(keep_chunks.begin(), keep_chunks.end()),	×
169	keep_chunks.end());	×
UNCOV 170	} else {	×
171	for (std::uint64_t c = 0; c < total_chunks; ++c)	×
172	keep_chunks.push_back(c);	×
173	}
UNCOV 174	} else {	×
175	for (std::uint64_t c = 0; c < total_chunks; ++c)	12✔
176	keep_chunks.push_back(c);	8!
177	}
178
179	// Intersect with the user's line range so workers only touch chunks that
180	// actually overlap it. Each work item carries the sub-line-range;
181	// LINE_RANGE on the read maps it back to bytes via the same checkpoint
182	// table the gzip stream uses.
183	if (clip.has_line_clip()) {	4!
184	std::size_t lo = clip.start_line > 0 ? clip.start_line : 1;	2!
185	std::size_t hi = clip.end_line > 0 ? clip.end_line : SIZE_MAX;	2!
186	std::vector<std::uint64_t> filtered;	2✔
187	filtered.reserve(keep_chunks.size());	2!
188	for (auto c : keep_chunks) {	6✔
189	std::size_t cf = geo.first_line(c);	4!
190	std::size_t cl = geo.last_line(c);	4!
191	if (cl < lo \|\| cf > hi) continue;	4!
192	filtered.push_back(c);	2!
193	}
194	keep_chunks = std::move(filtered);	2✔
195	}	2✔
196	return keep_chunks;	4✔
197	}	4!
198
199	// Group the kept chunks into contiguous work ranges (one per worker slot) and
200	// append the resulting work items.
201	void emit_work_items(std::vector<ArrowWorkItem> &items, const std::string &fp,	4✔
202	const ChunkGeometry &geo,
203	const std::vector<std::uint64_t> &keep_chunks,
204	bool file_pure_match, std::size_t max_workers,
205	const ClipRange &clip) {
206	std::size_t target_ranges = std::max<std::size_t>(1, max_workers);	4✔
207	std::size_t per_range = std::max<std::size_t>(	4✔
208	1, (keep_chunks.size() + target_ranges - 1) / target_ranges);	4✔
209
210	std::size_t group_start = 0;	4✔
211	while (group_start < keep_chunks.size()) {	10✔
212	std::size_t group_end = group_start;	6✔
213	std::size_t emitted = 0;	6✔
214	while (group_end < keep_chunks.size() && emitted < per_range) {	12✔
215	if (group_end > group_start &&	6!
216	keep_chunks[group_end] != keep_chunks[group_end - 1] + 1) {	×
217	break;	×
218	}
219	++group_end;	6✔
220	++emitted;	6✔
221	}
222	std::uint64_t scidx = keep_chunks[group_start];	6✔
223	std::uint64_t ecidx = keep_chunks[group_end - 1];	6✔
224	std::size_t start_byte = geo.start_byte(scidx);	6✔
225	std::size_t end_byte = geo.end_byte(ecidx);	6✔
226	// start_at_checkpoint: a gzip recovery point sits at start_byte (true
227	// for any cidx>=1; false for the implicit chunk 0 which decodes from
228	// stream start).
229	bool start_at_checkpoint = (scidx >= 1);	6✔
230	bool end_at_checkpoint = (group_end < keep_chunks.size());	6✔
231	if (clip.has_line_clip()) {	6✔
232	std::size_t lo = clip.start_line > 0 ? clip.start_line : 1;	2!
233	std::size_t hi = clip.end_line > 0 ? clip.end_line : SIZE_MAX;	2!
234	std::size_t cluster_first = geo.first_line(scidx);	2✔
235	std::size_t cluster_last = geo.last_line(ecidx);	2✔
236	std::size_t item_start = std::max<std::size_t>(lo, cluster_first);	2✔
237	std::size_t item_end = std::min<std::size_t>(hi, cluster_last);	2✔
238	if (item_start > item_end) {	2!
239	group_start = group_end;	×
240	continue;	×
241	}
242	ArrowWorkItem item;	2✔
243	item.file_path = fp;	2!
244	item.chunk_prune_only = file_pure_match;	2✔
245	item.start_line = item_start;	2✔
246	item.end_line = item_end;	2✔
247	items.push_back(std::move(item));	2!
248	group_start = group_end;	2✔
249	continue;
250	}	2✔
251	if (clip.has_byte_clip()) {	4✔
252	if (start_byte < clip.start_byte) {	2!
253	start_byte = clip.start_byte;	×
254	start_at_checkpoint = false;	×
UNCOV 255	}	×
256	if (end_byte > clip.end_byte) {	2!
257	end_byte = clip.end_byte;	2✔
258	end_at_checkpoint = false;	2✔
259	}	2✔
260	if (start_byte >= end_byte) {	2✔
261	group_start = group_end;	1✔
262	continue;	1✔
263	}
264	}	1✔
265	items.push_back({fp, start_byte, end_byte, start_at_checkpoint,	9!
266	end_at_checkpoint, file_pure_match});	6✔
267	group_start = group_end;	3✔
268	}
269	}	4✔
270
271	} // namespace
272
273	std::vector<ArrowWorkItem> enumerate_work_items(	36✔
274	const std::vector<std::string> &files, const std::string &index_dir,
275	const std::string &query_str, std::size_t max_workers,
276	std::size_t clip_start_byte, std::size_t clip_end_byte,
277	std::size_t clip_start_line, std::size_t clip_end_line) {
278	namespace dft_internal =
279	dftracer::utils::utilities::composites::dft::internal;
280	namespace indexer_ns = dftracer::utils::utilities::indexer;
281	namespace indexing = dftracer::utils::utilities::composites::dft::indexing;
282
283	const ClipRange clip{clip_start_byte, clip_end_byte, clip_start_line,	72✔
284	clip_end_line};	36✔
285
286	std::vector<ArrowWorkItem> items;	36✔
287	items.reserve(files.size() * 4);	36!
288
289	auto push_unsplit = [&](const std::string &fp) {	93✔
290	ArrowWorkItem item;	57✔
291	item.file_path = fp;	57!
292	item.start_line = clip.start_line;	57✔
293	item.end_line = clip.end_line;	57✔
294	items.push_back(std::move(item));	57!
295	};	57✔
296
297	// Parse the query once. Pruner input copies a Query, so we keep the
298	// parsed form around to feed each ChunkPrunerInput without re-parsing.
299	std::optional<dftracer::utils::utilities::common::query::Query> parsed;	36✔
300	if (!query_str.empty()) {	36✔
301	auto r = dftracer::utils::utilities::common::query::Query::from_string(	3!
302	query_str);	3✔
303	if (r) parsed = std::move(*r);	3!
304	}	3✔
305
306	// All files in a directory-mode scan share the same `.dftindex` root.
307	// Group files by their resolved index path so we can open the RocksDB
308	// once per index and reuse it to prune every file against that handle.
309	std::unordered_map<std::string, std::vector<std::size_t>> by_index;	36✔
310	for (std::size_t i = 0; i < files.size(); ++i) {	97✔
311	std::string index_path =
312	dft_internal::determine_index_path(files[i], index_dir);	61!
313	by_index[index_path].push_back(i);	61!
314	}	61✔
315
316	for (auto &entry : by_index) {	75!
317	const auto &index_path = entry.first;	39✔
318	const auto &file_idxs = entry.second;	39✔
319	if (!fs::exists(index_path)) {	39!
320	for (auto i : file_idxs) push_unsplit(files[i]);	78!
321	continue;	31✔
322	}
323	std::unique_ptr<indexer_ns::IndexDatabase> idx_db;	8✔
324	try {
325	idx_db = std::make_unique<indexer_ns::IndexDatabase>(	8!
326	index_path,	8✔
327	dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);	8✔
328	} catch (...) {	8✔
329	for (auto i : file_idxs) push_unsplit(files[i]);	×
330	continue;
331	}	×
332
333	// Resolve fid + checkpoints per file (cheap queries).
334	struct FileCtx {
335	std::size_t file_idx;
336	int fid;
337	std::vector<indexer_ns::IndexerCheckpoint> ckpts;
338	};
339	std::vector<FileCtx> file_ctxs;	8✔
340	file_ctxs.reserve(file_idxs.size());	8!
341	for (auto i : file_idxs) {	22✔
342	FileCtx fc;	14✔
343	fc.file_idx = i;	14✔
344	fc.fid = idx_db->get_file_info_id(	14!
345	indexer_ns::internal::get_logical_path(files[i]));	14!
346	if (fc.fid < 0) {	14!
347	push_unsplit(files[i]);	×
348	continue;	×
349	}
350	fc.ckpts = idx_db->query_checkpoints(fc.fid);	14!
351	if (fc.ckpts.empty()) {	14✔
352	push_unsplit(files[i]);	10!
353	continue;	10✔
354	}
355	std::sort(fc.ckpts.begin(), fc.ckpts.end(),	4!
356	[](const auto &a, const auto &b) {	×
357	return a.first_line_num < b.first_line_num;	×
358	});
359	file_ctxs.push_back(std::move(fc));	4!
360	}	14✔
361
362	// Batch-prune all files against the shared index: dim_stats and
363	// chunk_statistics are loaded in one RocksDB scan each instead of
364	// one scan per file.
365	std::vector<indexing::ChunkPrunerOutput> pruner_outs(file_ctxs.size());	8!
366	if (parsed && !file_ctxs.empty()) {	8!
367	indexing::ChunkPrunerBatchInput batch_in;	×
368	batch_in.index_path = index_path;	×
369	batch_in.external_db = idx_db.get();	×
370	batch_in.items.reserve(file_ctxs.size());	×
371	for (auto &fc : file_ctxs) {	×
372	batch_in.items.push_back({files[fc.file_idx], *parsed});	×
373	}
374	indexing::ChunkPrunerUtility pruner;	×
375	auto batch_out = pruner.process_batch(batch_in);	×
376	if (batch_out) {	×
377	pruner_outs = std::move(batch_out->outputs);	×
UNCOV 378	}	×
379	}	×
380
381	// For AND-of-EQ predicates, precompute uniform-match leaves once.
382	// Per-file pure_match is checked inline below and lets workers skip
383	// per-event predicate eval on chunks where dim_stats min == max ==
384	// literal for every leaf.
385	std::optional<std::vector<std::pair<std::string, std::string>>>
386	eq_leaves;	8✔
387	if (parsed) eq_leaves = extract_eq_leaves(parsed->root());	8!
388
389	for (std::size_t fc_idx = 0; fc_idx < file_ctxs.size(); ++fc_idx) {	12✔
390	auto &fc = file_ctxs[fc_idx];	4✔
391	const auto &fp = files[fc.file_idx];	4✔
392
393	ChunkGeometry geo{fc.ckpts};	4✔
394	std::vector<std::uint64_t> keep_chunks = select_kept_chunks(	4!
395	geo, parsed.has_value(), pruner_outs[fc_idx], clip);	4✔
396	if (keep_chunks.empty()) continue;	4!
397
398	// All-or-nothing per file: if every kept chunk is uniform-matching
399	// for every leaf, every work item from this file gets the
400	// chunk_prune_only fast path. Mixed files fall back to per-event
401	// eval to stay safe.
402	bool file_pure_match = false;	4✔
403	if (eq_leaves && !eq_leaves->empty() && idx_db) {	4!
404	file_pure_match = all_chunks_uniform_match(	×
405	idx_db, fc.fid, eq_leaves, keep_chunks);	×
UNCOV 406	}	×
407
408	emit_work_items(items, fp, geo, keep_chunks, file_pure_match,	8!
409	max_workers, clip);	4✔
410	}	4!
411	}	8!
412	return items;	36✔
413	}	36!
414
415	} // namespace dftracer::utils::utilities::reader::internal
416
417	#endif // DFTRACER_UTILS_ENABLE_ARROW

llnl / dftracer-utils / 29787659139

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