• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

llnl / dftracer-utils / 28281886195

27 Jun 2026 06:59AM UTC coverage: 52.356% (+0.2%) from 52.111%
28281886195

push

github

hariharan-devarajan
feat(comparator): add dlio preset and consolidate constants

37518 of 93029 branches covered (40.33%)

Branch coverage included in aggregate %.

249 of 307 new or added lines in 8 files covered. (81.11%)

32 existing lines in 10 files now uncovered.

33715 of 43025 relevant lines covered (78.36%)

20293.3 hits per line

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

47.49
/src/dftracer/utils/python/trace_reader.cpp
1
#define PY_SSIZE_T_CLEAN
2
#include <Python.h>
3
#include <dftracer/utils/core/common/config.h>
4
#include <dftracer/utils/core/common/filesystem.h>
5
#include <dftracer/utils/core/common/memory_budget.h>
6
#include <dftracer/utils/core/coro/channel.h>
7
#include <dftracer/utils/core/coro/task.h>
8
#include <dftracer/utils/core/coro/when_all.h>
9
#include <dftracer/utils/core/tasks/coro_scope.h>
10
#include <dftracer/utils/core/utils/string.h>
11
#include <dftracer/utils/python/arrow_helpers.h>
12
#include <dftracer/utils/python/batch_byte_size.h>
13
#include <dftracer/utils/python/json.h>
14
#include <dftracer/utils/python/runtime.h>
15
#include <dftracer/utils/python/trace_reader.h>
16
#include <dftracer/utils/python/trace_reader_iterator.h>
17
#include <dftracer/utils/utilities/common/query/query.h>
18
#include <dftracer/utils/utilities/composites/dft/indexing/chunk_pruner_utility.h>
19
#include <dftracer/utils/utilities/composites/dft/internal/utils.h>
20
#include <dftracer/utils/utilities/filesystem/pattern_directory_scanner_utility.h>
21
#include <dftracer/utils/utilities/indexer/index_database.h>
22
#include <dftracer/utils/utilities/indexer/internal/helpers.h>
23
#include <dftracer/utils/utilities/reader/trace_reader.h>
24

25
#include <algorithm>
26
#include <cctype>
27
#include <cstddef>
28
#include <cstdio>
29
#include <cstring>
30
#include <exception>
31
#include <memory>
32
#include <string>
33
#include <unordered_map>
34
#include <vector>
35
#ifdef DFTRACER_UTILS_ENABLE_ARROW
36
#include <dftracer/utils/python/arrow_stream_capsule.h>
37
#include <dftracer/utils/utilities/common/arrow/column_builder.h>
38
#include <dftracer/utils/utilities/common/json/parser.h>
39
#endif
40
#ifdef DFTRACER_UTILS_ENABLE_ARROW_IPC
41
#include <dftracer/utils/utilities/common/arrow/ipc_writer.h>
42
#include <dftracer/utils/utilities/common/arrow/partition_writer.h>
43
#include <dftracer/utils/utilities/composites/dft/internal/utils.h>
44
#include <dftracer/utils/utilities/composites/dft/metadata_collector_utility.h>
45
#include <dftracer/utils/utilities/composites/dft/views/view_builder_utility.h>
46
#include <dftracer/utils/utilities/composites/dft/views/view_definition.h>
47
#include <dftracer/utils/utilities/composites/dft/views/view_reader_utility.h>
48
#endif
49

50
namespace {
51

52
using dftracer::utils::CoroScope;
53
using dftracer::utils::Runtime;
54
using dftracer::utils::coro::CoroTask;
55
using dftracer::utils::coro::when_all;
56
using dftracer::utils::utilities::filesystem::PatternDirectoryScannerUtility;
57
using dftracer::utils::utilities::filesystem::
58
    PatternDirectoryScannerUtilityInput;
59
using dftracer::utils::utilities::reader::ReadConfig;
60
using dftracer::utils::utilities::reader::TraceReader;
61
using dftracer::utils::utilities::reader::TraceReaderConfig;
62
#ifdef DFTRACER_UTILS_ENABLE_ARROW
63
using dftracer::utils::utilities::common::arrow::ColumnType;
64
using dftracer::utils::utilities::common::arrow::RecordBatchBuilder;
65
using dftracer::utils::utilities::common::json::JsonParser;
66
using dftracer::utils::utilities::common::json::JsonValueHelper;
67
#endif
68
#ifdef DFTRACER_UTILS_ENABLE_ARROW_IPC
69
using dftracer::utils::utilities::common::arrow::IpcCompression;
70
using dftracer::utils::utilities::common::arrow::PartitionWriter;
71
using dftracer::utils::utilities::common::arrow::PartitionWriteStats;
72
using dftracer::utils::utilities::composites::dft::MetadataCollectorUtility;
73
using dftracer::utils::utilities::composites::dft::
74
    MetadataCollectorUtilityInput;
75
using dftracer::utils::utilities::composites::dft::views::ViewBuilderInput;
76
using dftracer::utils::utilities::composites::dft::views::ViewBuilderUtility;
77
using dftracer::utils::utilities::composites::dft::views::ViewDefinition;
78
using dftracer::utils::utilities::composites::dft::views::ViewReaderInput;
79
using dftracer::utils::utilities::composites::dft::views::ViewReaderUtility;
80
#endif
81

82
using dftracer::utils::python::MemoryViewBatchData;
83
using dftracer::utils::python::MemoryViewBatchIteratorState;
84

85
CoroTask<void> produce_lines_batched(
4,536!
86
    std::shared_ptr<MemoryViewBatchIteratorState> state,
87
    dftracer::utils::coro::ChannelProducer<MemoryViewBatchData> producer,
88
    TraceReaderConfig cfg, ReadConfig rc, std::size_t batch_size) {
56!
89
    auto guard = producer.guard();
56!
90
    try {
91
        TraceReader reader(std::move(cfg));
56!
92
        auto gen = reader.read_lines(rc);
56!
93
        MemoryViewBatchData batch;
56✔
94
        std::size_t count = 0;
56✔
95

96
        while (auto opt = co_await gen.next()) {
4,062!
97
            if (state->cancelled.load(std::memory_order_acquire)) break;
946!
98
            auto sv = opt->content;
946✔
99
            Py_ssize_t offset = static_cast<Py_ssize_t>(batch.buffer.size());
946✔
100
            batch.buffer.insert(batch.buffer.end(), sv.begin(), sv.end());
946!
101
            batch.offsets.push_back(offset);
946!
102
            batch.lengths.push_back(static_cast<Py_ssize_t>(sv.size()));
946!
103
            ++count;
946✔
104

105
            if (count >= batch_size) {
946!
106
                auto batch_bytes = dftracer::utils::python::byte_size(batch);
×
107
                state->bytes_in_queue.fetch_add(batch_bytes,
108
                                                std::memory_order_acq_rel);
109
                if (!co_await producer.send(std::move(batch))) break;
×
110
                batch = MemoryViewBatchData{};
111
                count = 0;
112
            }
×
113
        }
1,002!
114
        if (count > 0 && !state->cancelled.load(std::memory_order_acquire)) {
153✔
115
            auto batch_bytes = dftracer::utils::python::byte_size(batch);
156!
116
            state->bytes_in_queue.fetch_add(batch_bytes,
156✔
117
                                            std::memory_order_acq_rel);
118
            co_await producer.send(std::move(batch));
208!
119
        }
52!
120
    } catch (...) {
2,358✔
121
        state->set_error(std::current_exception());
1!
122
    }
1!
123
}
6,481!
124

125
CoroTask<void> produce_raw_batched(
100!
126
    std::shared_ptr<MemoryViewBatchIteratorState> state,
127
    dftracer::utils::coro::ChannelProducer<MemoryViewBatchData> producer,
128
    TraceReaderConfig cfg, ReadConfig rc) {
10!
129
    auto guard = producer.guard();
10!
130
    try {
131
        TraceReader reader(std::move(cfg));
10!
132
        auto gen = reader.read_raw(rc);
10!
133
        while (auto opt = co_await gen.next()) {
642!
134
            if (state->cancelled.load(std::memory_order_acquire)) break;
445!
135
            MemoryViewBatchData batch;
445✔
136
            batch.buffer.assign(opt->data(), opt->data() + opt->size());
445!
137
            batch.offsets.push_back(0);
445!
138
            batch.lengths.push_back(static_cast<Py_ssize_t>(opt->size()));
445!
139
            auto batch_bytes = dftracer::utils::python::byte_size(batch);
445✔
140
            state->bytes_in_queue.fetch_add(batch_bytes,
445✔
141
                                            std::memory_order_acq_rel);
142
            if (!co_await producer.send(std::move(batch))) break;
593!
143
        }
454✔
144
    } catch (...) {
30✔
145
        state->set_error(std::current_exception());
×
146
    }
×
147
}
1,274!
148

149
using dftracer::utils::utilities::common::json::JsonParser;
150
using dftracer::utils::utilities::common::json::JsonValueHelper;
151

152
static constexpr std::size_t ESTIMATED_BYTES_PER_LINE = 256;
153
static constexpr std::size_t ESTIMATED_BYTES_PER_RAW_CHUNK = 4 * 1024 * 1024;
154
static constexpr std::size_t ESTIMATED_BYTES_PER_JSON_EVENT = 512;
155
static constexpr std::size_t ESTIMATED_BYTES_PER_ARROW_ROW = 1024;
156

157
static void insert_simdjson_value(ArgsMap &map, std::string_view key,
4,288✔
158
                                  simdjson::ondemand::value val) {
159
    auto type = val.type();
4,288✔
160
    if (type.error()) return;
4,320!
161
    switch (type.value_unsafe()) {
4,320!
162
        case simdjson::ondemand::json_type::string: {
960✔
163
            auto r = val.get_string();
1,920✔
164
            if (!r.error()) map.insert(key, std::string(r.value_unsafe()));
1,910!
165
            break;
1,914✔
166
        }
167
        case simdjson::ondemand::json_type::number: {
1,200✔
168
            auto ri = val.get_int64();
2,400✔
169
            if (!ri.error()) {
2,396✔
170
                auto v = ri.value_unsafe();
2,396✔
171
                if (v >= 0)
2,396✔
172
                    map.insert(key, static_cast<std::uint64_t>(v));
2,395!
173
                else
174
                    map.insert(key, v);
1!
175
            } else {
1,200✔
176
                auto rd = val.get_double();
×
177
                if (!rd.error()) map.insert(key, rd.value_unsafe());
×
178
            }
179
            break;
2,395✔
180
        }
181
        case simdjson::ondemand::json_type::boolean: {
182
            auto r = val.get_bool();
×
183
            if (!r.error()) map.insert(key, r.value_unsafe());
×
184
            break;
×
185
        }
186
        default:
187
            break;
×
188
    }
189
}
2,160✔
190

191
static void parse_json_to_event(JsonParser &parser, JsonDictEvent &ev) {
476✔
192
    ev.top.set_valid(true);
476✔
193
    parser.for_each_field(
714!
194
        [&](std::string_view key, simdjson::ondemand::value val) {
4,046✔
195
            if (key == "args") {
3,808✔
196
                auto obj = val.get_object();
479✔
197
                if (!obj.error()) {
480✔
198
                    ev.args.set_valid(true);
479✔
199
                    for (auto field : obj.value_unsafe()) {
1,435✔
200
                        if (field.error()) continue;
957!
201
                        auto fkey = field.unescaped_key();
957✔
202
                        if (fkey.error()) continue;
952!
203
                        auto fval = field.value();
952✔
204
                        if (fval.error()) continue;
953!
205
                        insert_simdjson_value(ev.args, fkey.value_unsafe(),
1,433!
206
                                              fval.value_unsafe());
953✔
207
                    }
208
                }
239✔
209
            } else {
240✔
210
                insert_simdjson_value(ev.top, key, val);
3,334✔
211
            }
212
        });
3,825✔
213
}
478✔
214

215
CoroTask<void> produce_json_dicts(
134!
216
    std::shared_ptr<JsonDictIteratorState> state,
217
    dftracer::utils::coro::ChannelProducer<JsonDictBatch> producer,
218
    TraceReaderConfig cfg, ReadConfig rc, std::size_t batch_size) {
1!
219
    auto guard = producer.guard();
1!
220
    try {
221
        TraceReader reader(std::move(cfg));
1!
222
        auto gen = reader.read_json(rc);
1!
223
        JsonDictBatch batch;
1✔
224
        batch.events.reserve(batch_size);
1!
225

226
        while (auto opt = co_await gen.next()) {
125!
227
            if (state->cancelled.load(std::memory_order_acquire)) break;
30!
228

229
            JsonDictEvent ev;
30!
230
            parse_json_to_event(*opt->parser, ev);
30!
231
            batch.events.push_back(std::move(ev));
30!
232

233
            if (batch.events.size() >= batch_size) {
30!
234
                auto batch_bytes = dftracer::utils::python::byte_size(batch);
×
235
                state->bytes_in_queue.fetch_add(batch_bytes,
236
                                                std::memory_order_acq_rel);
237
                if (!co_await producer.send(std::move(batch))) break;
×
238
                batch = JsonDictBatch{};
239
                batch.events.reserve(batch_size);
×
240
            }
×
241
        }
31!
242
        if (!batch.events.empty() &&
3✔
243
            !state->cancelled.load(std::memory_order_acquire)) {
1✔
244
            auto batch_bytes = dftracer::utils::python::byte_size(batch);
3!
245
            state->bytes_in_queue.fetch_add(batch_bytes,
3✔
246
                                            std::memory_order_acq_rel);
247
            co_await producer.send(std::move(batch));
4!
248
        }
1!
249
    } catch (...) {
69✔
250
        state->set_error(std::current_exception());
×
251
    }
×
252
}
195!
253

254
static CoroTask<void> send_files_to_channel(
144!
255
    std::shared_ptr<dftracer::utils::coro::Channel<std::string>> file_chan,
256
    const std::vector<std::string> *files, std::atomic<bool> *cancelled) {
8!
257
    for (const auto &fp : *files) {
128✔
258
        if (cancelled->load(std::memory_order_acquire)) break;
72!
259
        if (!co_await file_chan->send(fp)) break;
104!
260
    }
24✔
261
    file_chan->close();
8!
262
    co_return;
8✔
263
}
112!
264

265
static CoroTask<void> json_dict_file_worker(
970!
266
    std::shared_ptr<dftracer::utils::coro::Channel<std::string>> file_chan,
267
    dftracer::utils::coro::Channel<JsonDictBatch> *out_chan,
268
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
269
    ReadConfig rc, std::size_t batch_size, std::atomic<bool> *cancelled) {
9!
270
    dftracer::utils::coro::ChannelProducer<JsonDictBatch> producer(out_chan);
9!
271
    auto guard = producer.guard();
9!
272

273
    while (auto file_path = co_await file_chan->receive()) {
56!
274
        if (cancelled->load(std::memory_order_acquire)) co_return;
14!
275
        TraceReaderConfig cfg;
14✔
276
        cfg.file_path = std::move(*file_path);
14✔
277
        cfg.index_dir = index_dir;
14!
278
        cfg.checkpoint_size = checkpoint_size;
14✔
279
        cfg.auto_build_index = auto_build_index;
14✔
280

281
        TraceReader reader(std::move(cfg));
14!
282
        auto gen = reader.read_json(rc);
14!
283
        JsonDictBatch batch;
14✔
284
        batch.events.reserve(batch_size);
14!
285

286
        while (auto opt = co_await gen.next()) {
896!
287
            if (cancelled->load(std::memory_order_acquire)) co_return;
210!
288
            JsonDictEvent ev;
210!
289
            parse_json_to_event(*opt->parser, ev);
210!
290
            batch.events.push_back(std::move(ev));
210!
291
            if (batch.events.size() >= batch_size) {
210!
292
                if (!co_await producer.send(std::move(batch))) co_return;
×
293
                batch = JsonDictBatch{};
294
                batch.events.reserve(batch_size);
×
295
            }
296
        }
224!
297
        if (!batch.events.empty()) {
40!
298
            if (!co_await producer.send(std::move(batch))) co_return;
53!
299
        }
14✔
300
    }
497✔
301
    co_return;
9✔
302
}
1,469!
303

304
static CoroTask<void> spawn_json_dict_producers(
30!
305
    CoroScope &child, dftracer::utils::coro::Channel<JsonDictBatch> *out_chan,
306
    const std::vector<std::string> *files, const std::string *index_dir,
307
    std::size_t checkpoint_size, bool auto_build_index, const ReadConfig *rc,
308
    std::size_t batch_size, std::atomic<bool> *cancelled_ptr,
309
    std::size_t max_workers) {
5!
310
    std::size_t num_workers = std::min(files->size(), max_workers);
5!
311
    auto file_chan =
5✔
312
        dftracer::utils::coro::make_channel<std::string>(num_workers);
5!
313

314
    for (std::size_t i = 0; i < num_workers; ++i) {
14✔
315
        child.spawn([out_chan, fc = file_chan, idx = *index_dir,
27!
316
                     checkpoint_size, auto_build_index, r = *rc, batch_size,
27!
317
                     cancelled_ptr](CoroScope &) {
36✔
318
            return json_dict_file_worker(fc, out_chan, idx, checkpoint_size,
27!
319
                                         auto_build_index, r, batch_size,
18!
320
                                         cancelled_ptr);
18!
321
        });
322
    }
9✔
323

324
    child.spawn([fc = file_chan, files, cancelled_ptr](CoroScope &) {
15!
325
        return send_files_to_channel(fc, files, cancelled_ptr);
10!
326
    });
327
    co_return;
10✔
328
}
15!
329

330
static CoroTask<void> produce_json_dicts_parallel(
68!
331
    CoroScope &scope, JsonDictIteratorState *sp, std::string dir_path,
332
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
333
    ReadConfig rc, std::size_t batch_size, std::size_t max_workers) {
6!
334
    try {
335
        PatternDirectoryScannerUtility scanner;
18!
336
        auto scan_input = PatternDirectoryScannerUtilityInput(
36!
337
            dir_path, {".pfw", ".pfw.gz"}, true, false);
18!
338
        auto entries = co_await scope.spawn(scanner, scan_input);
30!
339

340
        std::vector<std::string> files;
16✔
341
        files.reserve(entries.size());
16✔
342
        for (auto &e : entries) files.push_back(e.path.string());
20!
343
        std::sort(files.begin(), files.end());
6✔
344

345
        if (files.empty()) {
16✔
346
            sp->channel->close();
1!
347
            co_return;
1✔
348
        }
349

350
        auto *chan_ptr = sp->channel.get();
15✔
351
        auto *cancelled_ptr = &sp->cancelled;
15✔
352

353
        co_await scope.scope([chan_ptr, &files, &index_dir, checkpoint_size,
120!
354
                              auto_build_index, &rc, batch_size, cancelled_ptr,
45✔
355
                              max_workers](CoroScope &child) -> CoroTask<void> {
20!
356
            co_await spawn_json_dict_producers(
40!
357
                child, chan_ptr, &files, &index_dir, checkpoint_size,
15✔
358
                auto_build_index, &rc, batch_size, cancelled_ptr, max_workers);
15✔
359
        });
20!
360
    } catch (...) {
16✔
361
        sp->set_error(std::current_exception());
×
362
    }
×
363
}
78!
364

365
static CoroTask<void> lines_file_worker(
544!
366
    std::shared_ptr<dftracer::utils::coro::Channel<std::string>> file_chan,
367
    dftracer::utils::coro::Channel<MemoryViewBatchData> *out_chan,
368
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
369
    ReadConfig rc, std::size_t batch_size, std::atomic<bool> *cancelled) {
4!
370
    dftracer::utils::coro::ChannelProducer<MemoryViewBatchData> producer(
8!
371
        out_chan);
4✔
372
    auto guard = producer.guard();
4!
373

374
    while (auto file_path = co_await file_chan->receive()) {
27!
375
        if (cancelled->load(std::memory_order_acquire)) co_return;
7!
376
        TraceReaderConfig cfg;
7✔
377
        cfg.file_path = std::move(*file_path);
7✔
378
        cfg.index_dir = index_dir;
7!
379
        cfg.checkpoint_size = checkpoint_size;
7✔
380
        cfg.auto_build_index = auto_build_index;
7✔
381

382
        TraceReader reader(std::move(cfg));
7!
383
        auto gen = reader.read_lines(rc);
7!
384
        MemoryViewBatchData batch;
7✔
385
        std::size_t count = 0;
7✔
386

387
        while (auto opt = co_await gen.next()) {
508!
388
            if (cancelled->load(std::memory_order_acquire)) co_return;
120!
389
            auto sv = opt->content;
120✔
390
            Py_ssize_t offset = static_cast<Py_ssize_t>(batch.buffer.size());
120✔
391
            batch.buffer.insert(batch.buffer.end(), sv.begin(), sv.end());
120!
392
            batch.offsets.push_back(offset);
120!
393
            batch.lengths.push_back(static_cast<Py_ssize_t>(sv.size()));
120!
394
            ++count;
120✔
395
            if (count >= batch_size) {
120!
396
                if (!co_await producer.send(std::move(batch))) co_return;
×
397
                batch = MemoryViewBatchData{};
398
                count = 0;
399
            }
400
        }
127!
401
        if (count > 0) {
21!
402
            if (!co_await producer.send(std::move(batch))) co_return;
28!
403
        }
7✔
404
    }
279✔
405
    co_return;
4✔
406
}
824!
407

408
static CoroTask<void> spawn_lines_producers(
12!
409
    CoroScope &child,
410
    dftracer::utils::coro::Channel<MemoryViewBatchData> *out_chan,
411
    const std::vector<std::string> *files, const std::string *index_dir,
412
    std::size_t checkpoint_size, bool auto_build_index, const ReadConfig *rc,
413
    std::size_t batch_size, std::atomic<bool> *cancelled_ptr,
414
    std::size_t max_workers) {
2!
415
    std::size_t num_workers = std::min(files->size(), max_workers);
2!
416
    auto file_chan =
2✔
417
        dftracer::utils::coro::make_channel<std::string>(num_workers);
2!
418

419
    for (std::size_t i = 0; i < num_workers; ++i) {
6✔
420
        child.spawn([out_chan, fc = file_chan, idx = *index_dir,
12!
421
                     checkpoint_size, auto_build_index, r = *rc, batch_size,
12!
422
                     cancelled_ptr](CoroScope &) {
16✔
423
            return lines_file_worker(fc, out_chan, idx, checkpoint_size,
12!
424
                                     auto_build_index, r, batch_size,
8!
425
                                     cancelled_ptr);
8!
426
        });
427
    }
4✔
428

429
    child.spawn([fc = file_chan, files, cancelled_ptr](CoroScope &) {
6!
430
        return send_files_to_channel(fc, files, cancelled_ptr);
4!
431
    });
432
    co_return;
4✔
433
}
6!
434

435
static CoroTask<void> produce_lines_parallel(
32!
436
    CoroScope &scope, MemoryViewBatchIteratorState *sp, std::string dir_path,
437
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
438
    ReadConfig rc, std::size_t batch_size, std::size_t max_workers) {
3!
439
    try {
440
        PatternDirectoryScannerUtility scanner;
9!
441
        auto scan_input = PatternDirectoryScannerUtilityInput(
18!
442
            dir_path, {".pfw", ".pfw.gz"}, true, false);
9!
443
        auto entries = co_await scope.spawn(scanner, scan_input);
15!
444

445
        std::vector<std::string> files;
7✔
446
        files.reserve(entries.size());
7✔
447
        for (auto &e : entries) files.push_back(e.path.string());
10!
448
        std::sort(files.begin(), files.end());
3✔
449

450
        if (files.empty()) {
7✔
451
            sp->channel->close();
1!
452
            co_return;
1✔
453
        }
454

455
        auto *chan_ptr = sp->channel.get();
6✔
456
        auto *cancelled_ptr = &sp->cancelled;
6✔
457

458
        co_await scope.scope([chan_ptr, &files, &index_dir, checkpoint_size,
48!
459
                              auto_build_index, &rc, batch_size, cancelled_ptr,
18✔
460
                              max_workers](CoroScope &child) -> CoroTask<void> {
8!
461
            co_await spawn_lines_producers(
16!
462
                child, chan_ptr, &files, &index_dir, checkpoint_size,
6✔
463
                auto_build_index, &rc, batch_size, cancelled_ptr, max_workers);
6✔
464
        });
8!
465
    } catch (...) {
7✔
466
        sp->set_error(std::current_exception());
×
467
    }
×
468
}
36!
469

470
static CoroTask<void> raw_file_worker(
12!
471
    std::shared_ptr<dftracer::utils::coro::Channel<std::string>> file_chan,
472
    dftracer::utils::coro::Channel<MemoryViewBatchData> *out_chan,
473
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
474
    ReadConfig rc, std::atomic<bool> *cancelled) {
2!
475
    dftracer::utils::coro::ChannelProducer<MemoryViewBatchData> producer(
4!
476
        out_chan);
2✔
477
    auto guard = producer.guard();
2!
478

479
    while (auto file_path = co_await file_chan->receive()) {
13!
480
        if (cancelled->load(std::memory_order_acquire)) co_return;
3!
481
        TraceReaderConfig cfg;
3✔
482
        cfg.file_path = std::move(*file_path);
3✔
483
        cfg.index_dir = index_dir;
3!
484
        cfg.checkpoint_size = checkpoint_size;
3✔
485
        cfg.auto_build_index = auto_build_index;
3✔
486

487
        TraceReader reader(std::move(cfg));
3!
488
        auto gen = reader.read_raw(rc);
3!
489
        while (auto opt = co_await gen.next()) {
250!
490
            if (cancelled->load(std::memory_order_acquire)) co_return;
180!
491
            MemoryViewBatchData batch;
180✔
492
            batch.buffer.assign(opt->data(), opt->data() + opt->size());
180!
493
            batch.offsets.push_back(0);
180!
494
            batch.lengths.push_back(static_cast<Py_ssize_t>(opt->size()));
180!
495
            if (!co_await producer.send(std::move(batch))) co_return;
240!
496
        }
183!
497
    }
9!
498
    co_return;
2✔
499
}
504!
500

501
static CoroTask<void> spawn_raw_producers(
6!
502
    CoroScope &child,
503
    dftracer::utils::coro::Channel<MemoryViewBatchData> *out_chan,
504
    const std::vector<std::string> *files, const std::string *index_dir,
505
    std::size_t checkpoint_size, bool auto_build_index, const ReadConfig *rc,
506
    std::atomic<bool> *cancelled_ptr, std::size_t max_workers) {
1!
507
    std::size_t num_workers = std::min(files->size(), max_workers);
1!
508
    auto file_chan =
1✔
509
        dftracer::utils::coro::make_channel<std::string>(num_workers);
1!
510

511
    for (std::size_t i = 0; i < num_workers; ++i) {
3✔
512
        child.spawn([out_chan, fc = file_chan, idx = *index_dir,
6!
513
                     checkpoint_size, auto_build_index, r = *rc,
6!
514
                     cancelled_ptr](CoroScope &) {
8✔
515
            return raw_file_worker(fc, out_chan, idx, checkpoint_size,
6!
516
                                   auto_build_index, r, cancelled_ptr);
4!
517
        });
518
    }
2✔
519

520
    child.spawn([fc = file_chan, files, cancelled_ptr](CoroScope &) {
3!
521
        return send_files_to_channel(fc, files, cancelled_ptr);
2!
522
    });
523
    co_return;
2✔
524
}
3!
525

526
static CoroTask<void> produce_raw_parallel(
20!
527
    CoroScope &scope, MemoryViewBatchIteratorState *sp, std::string dir_path,
528
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
529
    ReadConfig rc, std::size_t max_workers) {
2!
530
    try {
531
        PatternDirectoryScannerUtility scanner;
6!
532
        auto scan_input = PatternDirectoryScannerUtilityInput(
12!
533
            dir_path, {".pfw", ".pfw.gz"}, true, false);
6!
534
        auto entries = co_await scope.spawn(scanner, scan_input);
10!
535

536
        std::vector<std::string> files;
4✔
537
        files.reserve(entries.size());
4✔
538
        for (auto &e : entries) files.push_back(e.path.string());
5!
539
        std::sort(files.begin(), files.end());
2✔
540

541
        if (files.empty()) {
4✔
542
            sp->channel->close();
1!
543
            co_return;
1✔
544
        }
545

546
        auto *chan_ptr = sp->channel.get();
3✔
547
        auto *cancelled_ptr = &sp->cancelled;
3✔
548

549
        co_await scope.scope([chan_ptr, &files, &index_dir, checkpoint_size,
21!
550
                              auto_build_index, &rc, cancelled_ptr,
6✔
551
                              max_workers](CoroScope &child) -> CoroTask<void> {
4!
552
            co_await spawn_raw_producers(child, chan_ptr, &files, &index_dir,
8!
553
                                         checkpoint_size, auto_build_index, &rc,
3✔
554
                                         cancelled_ptr, max_workers);
3✔
555
        });
4!
556
    } catch (...) {
4✔
557
        sp->set_error(std::current_exception());
×
558
    }
×
559
}
22!
560

561
#ifdef DFTRACER_UTILS_ENABLE_ARROW
562

563
using dftracer::utils::utilities::common::arrow::ArrowExportResult;
564
using dftracer::utils::utilities::common::arrow::ColumnType;
565
using dftracer::utils::utilities::common::arrow::RecordBatchBuilder;
566

567
// Bump arena for string_views that must survive until builder.finish().
568
struct StringArena {
569
    static constexpr std::size_t BLOCK_SIZE = 64 * 1024;
570
    std::vector<std::vector<char>> blocks;
571
    std::size_t pos = 0;
572

573
    StringArena() { blocks.emplace_back(BLOCK_SIZE); }
×
574

575
    std::string_view push(const char *data, std::size_t len) {
×
576
        if (pos + len > blocks.back().size()) {
×
577
            blocks.emplace_back(std::max(BLOCK_SIZE, len));
×
578
            pos = 0;
×
579
        }
580
        char *dst = blocks.back().data() + pos;
×
581
        std::memcpy(dst, data, len);
×
582
        pos += len;
×
583
        return {dst, len};
×
584
    }
585

586
    void clear() {
×
587
        if (blocks.size() > 1) blocks.resize(1);
×
588
        pos = 0;
×
589
    }
×
590
};
591

592
// --- Row type constants (must match Python TYPE_* constants) ---
593
enum RowType : int8_t {
594
    ROW_EVENT = 0,
595
    ROW_FILE_HASH = 1,
596
    ROW_HOST_HASH = 2,
597
    ROW_STRING_HASH = 3,
598
    ROW_METADATA = 4,
599
    ROW_PROC_METADATA = 5,
600
    ROW_PROFILE = 6,
601
    ROW_SYSTEM = 7,
602
};
603

604
// --- IO category constants (must match Python IOCategory values) ---
605
enum IOCat : int8_t {
606
    IO_READ = 1,
607
    IO_WRITE = 2,
608
    IO_METADATA = 3,
609
    IO_PCTL = 4,
610
    IO_IPC = 5,
611
    IO_OTHER = 6,
612
    IO_SYNC = 7,
613
};
614

615
static int8_t get_io_cat(std::string_view func) {
×
616
    using namespace dftracer::utils::utilities::composites::dft::internal;
NEW
617
    for (auto op : posix_ops::READ)
×
NEW
618
        if (op == func) return IO_READ;
×
NEW
619
    for (auto op : posix_ops::WRITE)
×
NEW
620
        if (op == func) return IO_WRITE;
×
NEW
621
    for (auto op : posix_ops::SYNC)
×
NEW
622
        if (op == func) return IO_SYNC;
×
NEW
623
    for (auto op : posix_ops::PCTL)
×
NEW
624
        if (op == func) return IO_PCTL;
×
NEW
625
    for (auto op : posix_ops::IPC)
×
NEW
626
        if (op == func) return IO_IPC;
×
NEW
627
    for (auto op : posix_ops::METADATA)
×
NEW
628
        if (op == func) return IO_METADATA;
×
UNCOV
629
    return IO_OTHER;
×
630
}
631

632
static bool str_iequal(std::string_view a, const char *b) {
×
633
    std::size_t len = std::strlen(b);
×
634
    if (a.size() != len) return false;
×
635
    for (std::size_t i = 0; i < len; ++i) {
×
636
        if (std::tolower(static_cast<unsigned char>(a[i])) !=
×
637
            static_cast<unsigned char>(b[i]))
×
638
            return false;
×
639
    }
640
    return true;
×
641
}
642

643
static bool str_contains_lower(std::string_view s, const char *needle) {
×
644
    std::size_t nlen = std::strlen(needle);
×
645
    if (s.size() < nlen) return false;
×
646
    for (std::size_t i = 0; i <= s.size() - nlen; ++i) {
×
647
        bool match = true;
×
648
        for (std::size_t j = 0; j < nlen; ++j) {
×
649
            if (std::tolower(static_cast<unsigned char>(s[i + j])) !=
×
650
                static_cast<unsigned char>(needle[j])) {
×
651
                match = false;
×
652
                break;
×
653
            }
654
        }
655
        if (match) return true;
×
656
    }
657
    return false;
×
658
}
659

660
// Normalize a raw JSON row (parsed with simdjson) into the semantic
661
// output schema.  Appends one row to `builder` with the full set of output
662
// columns.  Returns false if the row should be skipped (no valid name).
663
static bool normalize_row(RecordBatchBuilder &builder, StringArena &arena,
×
664
                          JsonParser &parser) {
665
    using SVH = JsonValueHelper;
666

667
    // --- Extract top-level fields ---
668
    auto ph = parser.get_string("ph").value_or(std::string_view{});
×
669
    auto name_sv = parser.get_string("name").value_or(std::string_view{});
×
670
    auto cat_sv = parser.get_string("cat").value_or(std::string_view{});
×
671
    auto pid_opt = parser.get_int64("pid");
×
672
    auto tid_opt = parser.get_int64("tid");
×
673
    auto ts_opt = parser.get_int64("ts");
×
674
    auto dur_opt = parser.get_int64("dur");
×
675

676
    // Helper lambdas to access args fields (need to rewind after each access)
677
    // We'll do a single pass over args instead
678
    std::optional<std::string_view> args_name, args_value, args_hhash,
×
679
        args_fhash;
×
680
    std::optional<int64_t> args_epoch, args_step, args_size_sum, args_ret;
×
681
    std::optional<int64_t> args_offset, args_image_idx, args_image_size;
×
682
    std::unordered_map<std::string, int64_t> args_int_map;
×
683
    std::unordered_map<std::string, double> args_float_map;
×
684

685
    parser.rewind();
×
686
    parser.for_each_field(
×
687
        "args", [&](std::string_view key, simdjson::ondemand::value val) {
×
688
            if (key == "name") {
×
689
                if (auto s = SVH::get_string(val)) args_name = s;
×
690
            } else if (key == "value") {
×
691
                if (auto s = SVH::get_string(val)) args_value = s;
×
692
            } else if (key == "hhash") {
×
693
                if (auto s = SVH::get_string(val)) args_hhash = s;
×
694
            } else if (key == "fhash") {
×
695
                if (auto s = SVH::get_string(val)) args_fhash = s;
×
696
            } else if (key == "epoch") {
×
697
                if (auto i = SVH::get_int64(val)) args_epoch = i;
×
698
            } else if (key == "step") {
×
699
                if (auto i = SVH::get_int64(val)) args_step = i;
×
700
            } else if (key == "size_sum") {
×
701
                if (auto i = SVH::get_int64(val)) args_size_sum = i;
×
702
            } else if (key == "ret") {
×
703
                if (auto i = SVH::get_int64(val)) args_ret = i;
×
704
            } else if (key == "offset") {
×
705
                if (auto i = SVH::get_int64(val)) args_offset = i;
×
706
            } else if (key == "image_idx") {
×
707
                if (auto i = SVH::get_int64(val)) args_image_idx = i;
×
708
            } else if (key == "image_size") {
×
709
                if (auto i = SVH::get_int64(val)) args_image_size = i;
×
710
            } else {
711
                // Store other int/float args for profile/sys columns
712
                if (auto i = SVH::get_int64(val)) {
×
713
                    args_int_map[std::string(key)] = *i;
×
714
                } else if (auto d = SVH::get_double(val)) {
×
715
                    args_float_map[std::string(key)] = *d;
×
716
                }
717
            }
718
        });
×
719

720
    // --- Type classification ---
721
    bool is_M = (ph == "M");
×
722
    bool is_C = (ph == "C");
×
723
    bool is_event = !is_M && !is_C;
×
724

725
    int8_t row_type = ROW_EVENT;
×
726
    if (is_M) {
×
727
        if (name_sv == "FH")
×
728
            row_type = ROW_FILE_HASH;
×
729
        else if (name_sv == "HH")
×
730
            row_type = ROW_HOST_HASH;
×
731
        else if (name_sv == "SH")
×
732
            row_type = ROW_STRING_HASH;
×
733
        else if (name_sv == "PR")
×
734
            row_type = ROW_PROC_METADATA;
×
735
        else
736
            row_type = ROW_METADATA;
×
737
    } else if (is_C) {
×
738
        row_type = str_iequal(cat_sv, "sys") ? ROW_SYSTEM : ROW_PROFILE;
×
739
    }
740
    bool is_hash = (row_type >= ROW_FILE_HASH && row_type <= ROW_STRING_HASH) ||
×
741
                   row_type == ROW_PROC_METADATA;
742
    bool is_profile = (row_type == ROW_PROFILE);
×
743
    bool is_sys = (row_type == ROW_SYSTEM);
×
744

745
    // Name: metadata rows use args.name if available
746
    std::string_view out_name = name_sv;
×
747
    if (is_M && args_name && !args_name->empty()) {
×
748
        out_name = *args_name;
×
749
    }
750
    if (out_name.empty()) return false;  // skip rows without name
×
751

752
    // --- Declare all output columns ---
753
    auto ci_type = builder.add_or_get_column("type", ColumnType::INT64);
×
754
    auto ci_cat = builder.add_or_get_column("cat", ColumnType::STRING);
×
755
    auto ci_name = builder.add_or_get_column("name", ColumnType::STRING);
×
756
    auto ci_pid = builder.add_or_get_column("pid", ColumnType::INT64);
×
757
    auto ci_tid = builder.add_or_get_column("tid", ColumnType::INT64);
×
758
    auto ci_hash = builder.add_or_get_column("hash", ColumnType::STRING);
×
759
    auto ci_value = builder.add_or_get_column("value", ColumnType::STRING);
×
760
    auto ci_host_hash =
761
        builder.add_or_get_column("host_hash", ColumnType::STRING);
×
762
    auto ci_file_hash =
763
        builder.add_or_get_column("file_hash", ColumnType::STRING);
×
764
    auto ci_epoch = builder.add_or_get_column("epoch", ColumnType::INT64);
×
765
    auto ci_step = builder.add_or_get_column("step", ColumnType::INT64);
×
766
    auto ci_ts = builder.add_or_get_column("ts", ColumnType::INT64);
×
767
    auto ci_dur = builder.add_or_get_column("dur", ColumnType::INT64);
×
768
    auto ci_te = builder.add_or_get_column("te", ColumnType::INT64);
×
769
    [[maybe_unused]] auto ci_trange =
770
        builder.add_or_get_column("trange", ColumnType::INT64);
×
771
    auto ci_io_cat = builder.add_or_get_column("io_cat", ColumnType::INT64);
×
772
    auto ci_size = builder.add_or_get_column("size", ColumnType::INT64);
×
773
    auto ci_offset = builder.add_or_get_column("offset", ColumnType::INT64);
×
774
    auto ci_image_id = builder.add_or_get_column("image_id", ColumnType::INT64);
×
775

776
    // --- Populate core columns ---
777
    builder.append_int64(ci_type, row_type);
×
778

779
    // cat (lowercased) - write into arena
780
    if (!cat_sv.empty()) {
×
781
        char lbuf[256];
782
        std::size_t clen = std::min(cat_sv.size(), sizeof(lbuf));
×
783
        for (std::size_t i = 0; i < clen; ++i)
×
784
            lbuf[i] = static_cast<char>(
×
785
                std::tolower(static_cast<unsigned char>(cat_sv[i])));
×
786
        builder.append_string(ci_cat, arena.push(lbuf, clen));
×
787
    } else {
788
        builder.append_null(ci_cat);
×
789
    }
790

791
    builder.append_string(ci_name, out_name);
×
792

793
    if (pid_opt) builder.append_int64(ci_pid, *pid_opt);
×
794
    if (tid_opt) builder.append_int64(ci_tid, *tid_opt);
×
795

796
    // hash / value
797
    if (is_hash && args_value && !args_value->empty())
×
798
        builder.append_string(ci_hash, *args_value);
×
799
    if (row_type == ROW_METADATA && args_value && !args_value->empty())
×
800
        builder.append_string(ci_value, *args_value);
×
801

802
    // host_hash / file_hash
803
    if (args_hhash && !args_hhash->empty())
×
804
        builder.append_string(ci_host_hash, *args_hhash);
×
805
    if (args_fhash && !args_fhash->empty())
×
806
        builder.append_string(ci_file_hash, *args_fhash);
×
807

808
    // epoch / step
809
    if (args_epoch && *args_epoch >= 0)
×
810
        builder.append_int64(ci_epoch, *args_epoch);
×
811
    if (args_step && *args_step >= 0) builder.append_int64(ci_step, *args_step);
×
812

813
    // --- Temporal ---
814
    bool has_ts = (is_event || is_C) && ts_opt.has_value();
×
815
    bool has_dur = dur_opt.has_value();
×
816
    int64_t ts_val = 0, dur_val = 0;
×
817
    if (has_ts) {
×
818
        ts_val = *ts_opt;
×
819
        builder.append_int64(ci_ts, ts_val);
×
820
    }
821
    if (is_event && has_ts && has_dur) {
×
822
        dur_val = *dur_opt;
×
823
        builder.append_int64(ci_dur, dur_val);
×
824
        builder.append_int64(ci_te, ts_val + dur_val);
×
825
    }
826

827
    // --- IO columns (events only) ---
828
    if (is_event) {
×
829
        bool is_posix_stdio =
830
            str_iequal(cat_sv, "posix") || str_iequal(cat_sv, "stdio");
×
831
        int8_t io_cat = IO_OTHER;
×
832

833
        // size priority: size_sum > POSIX ret > image_size
834
        if (args_size_sum) {
×
835
            builder.append_int64(ci_size, *args_size_sum);
×
836
            if (is_posix_stdio) io_cat = get_io_cat(out_name);
×
837
        } else if (is_posix_stdio) {
×
838
            io_cat = get_io_cat(out_name);
×
839
            if (args_ret && *args_ret > 0 &&
×
840
                (io_cat == IO_READ || io_cat == IO_WRITE))
×
841
                builder.append_int64(ci_size, *args_ret);
×
842
            if (args_offset && *args_offset >= 0)
×
843
                builder.append_int64(ci_offset, *args_offset);
×
844
        } else {
845
            if (args_image_idx && *args_image_idx > 0)
×
846
                builder.append_int64(ci_image_id, *args_image_idx);
×
847
            if (args_image_size && *args_image_size > 0 &&
×
848
                !str_contains_lower(out_name, "open"))
×
849
                builder.append_int64(ci_size, *args_image_size);
×
850
        }
851
        builder.append_int64(ci_io_cat, io_cat);
×
852
    }
853

854
    // --- Profile columns ---
855
    if (is_profile) {
×
856
        bool is_posix_stdio =
857
            str_iequal(cat_sv, "posix") || str_iequal(cat_sv, "stdio");
×
858
        int8_t io_cat = is_posix_stdio ? get_io_cat(out_name) : IO_OTHER;
×
859
        builder.append_int64(ci_io_cat, io_cat);
×
860

861
        static const char *profile_keys[] = {
862
            "count",      "count_max",  "count_min",  "count_sum",
863
            "dft_cnt",    "dur",        "dur_max",    "dur_min",
864
            "dur_sum",    "epoch",      "flags",      "offset",
865
            "offset_max", "offset_min", "offset_sum", "ret",
866
            "ret_max",    "ret_min",    "ret_sum",    "whence",
867
            "whence_max", "whence_min", "whence_sum", nullptr};
868
        for (const char **pk = profile_keys; *pk; ++pk) {
×
869
            auto it = args_int_map.find(*pk);
×
870
            if (it != args_int_map.end()) {
×
871
                auto idx = builder.add_or_get_column(*pk, ColumnType::INT64);
×
872
                builder.append_int64(idx, it->second);
×
873
            }
874
        }
875
    }
876

877
    // --- System columns ---
878
    if (is_sys) {
×
879
        static const char *sys_keys[] = {
880
            "user_pct", "system_pct",  "iowait_pct",   "idle_pct",
881
            "irq_pct",  "softirq_pct", "MemAvailable", "MemFree",
882
            "Cached",   "Dirty",       "Active",       nullptr};
883
        for (const char **sk = sys_keys; *sk; ++sk) {
×
884
            auto it = args_float_map.find(*sk);
×
885
            if (it != args_float_map.end()) {
×
886
                auto idx = builder.add_or_get_column(*sk, ColumnType::DOUBLE);
×
887
                builder.append_double(idx, it->second);
×
888
            }
889
        }
890
    }
891

892
    builder.end_row();
×
893
    return true;
×
894
}
×
895

896
// Flatten a simdjson object into "prefix.key" columns using native types.
897
// On type mismatch (same key, different type across rows), appends null.
898
static void flatten_object_into(RecordBatchBuilder &builder, StringArena &arena,
899
                                std::string_view prefix,
900
                                simdjson::ondemand::object obj) {
901
    using SVH = JsonValueHelper;
902
    char key_buf[512];
903

904
    for (auto field : obj) {
×
905
        if (field.error()) continue;
×
906

907
        auto key_result = field.unescaped_key();
908
        if (key_result.error()) continue;
×
909
        std::string_view sk = key_result.value_unsafe();
910

911
        auto val_result = field.value();
912
        if (val_result.error()) continue;
×
913
        auto sub_val = val_result.value_unsafe();
914

915
        std::size_t needed = prefix.size() + 1 + sk.size();
916
        if (needed >= sizeof(key_buf)) continue;
×
917
        std::memcpy(key_buf, prefix.data(), prefix.size());
918
        key_buf[prefix.size()] = '.';
919
        std::memcpy(key_buf + prefix.size() + 1, sk.data(), sk.size());
920
        std::string_view full_key(key_buf, needed);
921

922
        auto type_result = sub_val.type();
923
        if (type_result.error()) continue;
×
924
        auto json_type = type_result.value_unsafe();
925

926
        switch (json_type) {
×
927
            case simdjson::ondemand::json_type::number: {
928
                auto num_result = sub_val.get_number();
929
                if (num_result.error()) break;
×
930
                auto num = num_result.value_unsafe();
931
                if (num.is_int64()) {
×
932
                    auto idx =
933
                        builder.add_or_get_column(full_key, ColumnType::INT64);
×
934
                    if (builder.column_type(idx) == ColumnType::INT64)
×
935
                        builder.append_int64(idx, num.get_int64());
×
936
                    else
937
                        builder.append_null(idx);
×
938
                } else if (num.is_uint64()) {
×
939
                    auto idx =
940
                        builder.add_or_get_column(full_key, ColumnType::UINT64);
×
941
                    if (builder.column_type(idx) == ColumnType::UINT64)
×
942
                        builder.append_uint64(idx, num.get_uint64());
×
943
                    else
944
                        builder.append_null(idx);
×
945
                } else {
946
                    auto idx =
947
                        builder.add_or_get_column(full_key, ColumnType::DOUBLE);
×
948
                    if (builder.column_type(idx) == ColumnType::DOUBLE)
×
949
                        builder.append_double(idx, num.get_double());
×
950
                    else
951
                        builder.append_null(idx);
×
952
                }
953
                break;
954
            }
955
            case simdjson::ondemand::json_type::string: {
956
                auto str_result = sub_val.get_string();
957
                if (str_result.error()) break;
×
958
                auto str = str_result.value_unsafe();
959
                auto idx =
960
                    builder.add_or_get_column(full_key, ColumnType::STRING);
×
961
                if (builder.column_type(idx) == ColumnType::STRING)
×
962
                    builder.append_string(idx, str);
×
963
                else
964
                    builder.append_null(idx);
×
965
                break;
966
            }
967
            case simdjson::ondemand::json_type::boolean: {
968
                auto bool_result = sub_val.get_bool();
969
                if (bool_result.error()) break;
×
970
                auto b = bool_result.value_unsafe();
971
                auto idx =
972
                    builder.add_or_get_column(full_key, ColumnType::BOOL);
×
973
                if (builder.column_type(idx) == ColumnType::BOOL)
×
974
                    builder.append_bool(idx, b);
×
975
                else
976
                    builder.append_null(idx);
×
977
                break;
978
            }
979
            case simdjson::ondemand::json_type::null: {
980
                auto existing = builder.find_column(full_key);
×
981
                if (existing) builder.append_null(*existing);
×
982
                break;
983
            }
984
            case simdjson::ondemand::json_type::object:
985
            case simdjson::ondemand::json_type::array: {
986
                // Serialize nested object/array to JSON string
987
                auto json_str = SVH::to_json_string(sub_val);
×
988
                auto idx =
989
                    builder.add_or_get_column(full_key, ColumnType::STRING);
×
990
                if (json_str) {
×
991
                    builder.append_string(
×
992
                        idx, arena.push(json_str->data(), json_str->size()));
993
                } else {
994
                    builder.append_null(idx);
×
995
                }
996
                break;
997
            }
998
            default:
999
                break;
1000
        }
1001
    }
1002
}
1003

1004
static bool build_arrow_row(RecordBatchBuilder &builder, JsonParser &parser,
×
1005
                            StringArena &arena, bool normalize) {
1006
    if (normalize) return normalize_row(builder, arena, parser);
×
1007

1008
    using SVH = JsonValueHelper;
1009
    parser.for_each_field([&](std::string_view key_sv,
×
1010
                              simdjson::ondemand::value val) {
1011
        auto type_result = val.type();
×
1012
        if (type_result.error()) return;
×
1013
        auto json_type = type_result.value_unsafe();
×
1014
        switch (json_type) {
×
1015
            case simdjson::ondemand::json_type::number: {
1016
                auto num_result = val.get_number();
×
1017
                if (num_result.error()) break;
×
1018
                auto num = num_result.value_unsafe();
×
1019
                if (num.is_int64()) {
×
1020
                    std::size_t idx =
1021
                        builder.add_or_get_column(key_sv, ColumnType::INT64);
×
1022
                    builder.append_int64(idx, num.get_int64());
×
1023
                } else if (num.is_uint64()) {
×
1024
                    std::size_t idx =
1025
                        builder.add_or_get_column(key_sv, ColumnType::UINT64);
×
1026
                    builder.append_uint64(idx, num.get_uint64());
×
1027
                } else {
1028
                    std::size_t idx =
1029
                        builder.add_or_get_column(key_sv, ColumnType::DOUBLE);
×
1030
                    builder.append_double(idx, num.get_double());
×
1031
                }
1032
                break;
×
1033
            }
1034
            case simdjson::ondemand::json_type::string: {
1035
                auto str_result = val.get_string();
×
1036
                if (str_result.error()) break;
×
1037
                auto str = str_result.value_unsafe();
×
1038
                std::size_t idx =
1039
                    builder.add_or_get_column(key_sv, ColumnType::STRING);
×
1040
                builder.append_string(idx, str);
×
1041
                break;
×
1042
            }
1043
            case simdjson::ondemand::json_type::boolean: {
1044
                auto bool_result = val.get_bool();
×
1045
                if (bool_result.error()) break;
×
1046
                auto b = bool_result.value_unsafe();
×
1047
                std::size_t idx =
1048
                    builder.add_or_get_column(key_sv, ColumnType::BOOL);
×
1049
                builder.append_bool(idx, b);
×
1050
                break;
×
1051
            }
1052
            case simdjson::ondemand::json_type::null: {
1053
                auto existing = builder.find_column(key_sv);
×
1054
                if (existing) builder.append_null(*existing);
×
1055
                break;
×
1056
            }
1057
            case simdjson::ondemand::json_type::object:
1058
            case simdjson::ondemand::json_type::array: {
1059
                auto json_str = SVH::to_json_string(val);
×
1060
                std::size_t idx =
1061
                    builder.add_or_get_column(key_sv, ColumnType::STRING);
×
1062
                if (json_str) {
×
1063
                    builder.append_string(
×
1064
                        idx, arena.push(json_str->data(), json_str->size()));
×
1065
                } else {
1066
                    builder.append_null(idx);
×
1067
                }
1068
                break;
1069
            }
×
1070
            default:
1071
                break;
×
1072
        }
1073
    });
1074
    builder.end_row();
×
1075
    return true;
×
1076
}
1077

1078
static bool process_json_line(RecordBatchBuilder &builder, JsonParser &parser,
1079
                              StringArena &arena, std::string_view content,
1080
                              bool normalize) {
1081
    const char *trimmed;
1082
    std::size_t trimmed_length;
1083
    if (!dftracer::utils::json_trim_and_validate_with_comma(
×
1084
            content.data(), content.size(), trimmed, trimmed_length))
1085
        return false;
1086
    if (!parser.parse(std::string_view(trimmed, trimmed_length))) return false;
×
1087
    return build_arrow_row(builder, parser, arena, normalize);
×
1088
}
1089

1090
static CoroTask<void> produce_arrow_for_file(
×
1091
    dftracer::utils::coro::Channel<ArrowExportResult> *chan,
1092
    std::string file_path, std::string index_dir, std::size_t checkpoint_size,
1093
    bool auto_build_index, ReadConfig rc, std::size_t batch_size,
1094
    bool normalize, std::atomic<bool> *cancelled) {
1095
    dftracer::utils::coro::ChannelProducer<ArrowExportResult> producer(chan);
1096
    auto guard = producer.guard();
1097

1098
    TraceReaderConfig cfg;
1099
    cfg.file_path = std::move(file_path);
1100
    cfg.index_dir = std::move(index_dir);
1101
    cfg.checkpoint_size = checkpoint_size;
1102
    cfg.auto_build_index = auto_build_index;
1103

1104
    TraceReader reader(std::move(cfg));
1105

1106
    // Fast path: non-normalized Arrow build happens inside TraceReader.
1107
    // Normalize still goes through read_json + build_arrow_row for the
1108
    // richer schema derivation.
1109
    if (!normalize) {
1110
        auto batch_gen = reader.read_arrow(rc, batch_size);
1111
        while (auto batch_opt = co_await batch_gen.next()) {
1112
            if (cancelled->load(std::memory_order_acquire)) co_return;
1113
            if (!co_await producer.send(std::move(*batch_opt))) co_return;
1114
        }
1115
        co_return;
1116
    }
1117

1118
    auto gen = reader.read_json(rc);
1119
    RecordBatchBuilder builder;
1120
    builder.reserve(batch_size);
1121
    StringArena arena;
1122

1123
    while (auto opt = co_await gen.next()) {
1124
        if (cancelled->load(std::memory_order_acquire)) co_return;
1125
        if (!build_arrow_row(builder, *opt->parser, arena, normalize)) continue;
1126
        if (builder.num_rows() >= batch_size) {
1127
            auto result = builder.finish();
1128
            arena.clear();
1129
            if (!co_await producer.send(std::move(result))) co_return;
1130
            if (!builder.is_schema_locked()) builder.lock_schema();
1131
            builder.reset(true);
1132
            builder.reserve(batch_size);
1133
        }
1134
    }
1135
    if (builder.num_rows() > 0) {
1136
        co_await producer.send(builder.finish());
1137
    }
1138
    co_return;
1139
}
×
1140

1141
static CoroTask<void> file_worker(
×
1142
    std::shared_ptr<dftracer::utils::coro::Channel<std::string>> file_chan,
1143
    dftracer::utils::coro::Channel<ArrowExportResult> *out_chan,
1144
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
1145
    ReadConfig rc, std::size_t batch_size, bool normalize,
1146
    std::atomic<bool> *cancelled) {
1147
    dftracer::utils::coro::ChannelProducer<ArrowExportResult> producer(
1148
        out_chan);
1149
    auto guard = producer.guard();
1150

1151
    while (auto file_path = co_await file_chan->receive()) {
1152
        if (cancelled->load(std::memory_order_acquire)) co_return;
1153
        TraceReaderConfig cfg;
1154
        cfg.file_path = std::move(*file_path);
1155
        cfg.index_dir = index_dir;
1156
        cfg.checkpoint_size = checkpoint_size;
1157
        cfg.auto_build_index = auto_build_index;
1158

1159
        TraceReader reader(std::move(cfg));
1160

1161
        if (!normalize) {
1162
            auto batch_gen = reader.read_arrow(rc, batch_size);
1163
            while (auto batch_opt = co_await batch_gen.next()) {
1164
                if (cancelled->load(std::memory_order_acquire)) co_return;
1165
                if (!co_await producer.send(std::move(*batch_opt))) co_return;
1166
            }
1167
            continue;
1168
        }
1169

1170
        auto gen = reader.read_json(rc);
1171
        RecordBatchBuilder builder;
1172
        builder.reserve(batch_size);
1173
        StringArena arena;
1174

1175
        while (auto opt = co_await gen.next()) {
1176
            if (cancelled->load(std::memory_order_acquire)) co_return;
1177
            if (!build_arrow_row(builder, *opt->parser, arena, normalize))
1178
                continue;
1179
            if (builder.num_rows() >= batch_size) {
1180
                auto result = builder.finish();
1181
                arena.clear();
1182
                if (!co_await producer.send(std::move(result))) co_return;
1183
                if (!builder.is_schema_locked()) builder.lock_schema();
1184
                builder.reset(true);
1185
                builder.reserve(batch_size);
1186
            }
1187
        }
1188
        if (builder.num_rows() > 0) {
1189
            if (!co_await producer.send(builder.finish())) co_return;
1190
        }
1191
    }
1192
    co_return;
1193
}
×
1194

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

1204
    auto literal_to_string = [](const q_ns::LiteralNode &lit) -> std::string {
×
1205
        return std::visit(
1206
            [](auto &&v) -> std::string {
×
1207
                using T = std::decay_t<decltype(v)>;
1208
                if constexpr (std::is_same_v<T, std::string>)
1209
                    return v;
×
1210
                else if constexpr (std::is_same_v<T, bool>)
1211
                    return v ? "true" : "false";
×
1212
                else if constexpr (std::is_same_v<T, int64_t>)
1213
                    return std::to_string(v);
×
1214
                else if constexpr (std::is_same_v<T, uint64_t>)
1215
                    return std::to_string(v);
×
1216
                else if constexpr (std::is_same_v<T, double>)
1217
                    return std::to_string(v);
×
1218
                else
1219
                    return {};
1220
            },
1221
            lit.value);
×
1222
    };
1223

1224
    return std::visit(
1225
        [&](const auto &n) -> std::optional<LeafVec> {
×
1226
            using T = std::decay_t<decltype(n)>;
1227
            if constexpr (std::is_same_v<T, q_ns::CompareNode>) {
1228
                if (n.op != q_ns::CompareOp::EQ) return std::nullopt;
×
1229
                return LeafVec{{n.field.path, literal_to_string(n.value)}};
×
1230
            } else if constexpr (std::is_same_v<T, q_ns::AndNode>) {
1231
                auto l = extract_eq_leaves(*n.left);
×
1232
                if (!l) return std::nullopt;
×
1233
                auto r = extract_eq_leaves(*n.right);
×
1234
                if (!r) return std::nullopt;
×
1235
                l->insert(l->end(), r->begin(), r->end());
×
1236
                return l;
×
1237
            } else {
×
1238
                return std::nullopt;
×
1239
            }
1240
        },
1241
        node.data);
×
1242
}
1243

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

1254
    for (const auto &[dim, val] : leaves) {
×
1255
        auto rows = db.query_chunk_dimension_stats_for_dimension(fid, dim);
×
1256
        if (rows.empty()) return false;
×
1257
        std::unordered_map<std::uint64_t,
1258
                           const indexing::ChunkDimensionStatsResult *>
1259
            by_ckpt;
×
1260
        by_ckpt.reserve(rows.size());
×
1261
        for (const auto &r : rows) by_ckpt.emplace(r.checkpoint_idx, &r);
×
1262
        for (auto cidx : chunk_idxs) {
×
1263
            auto it = by_ckpt.find(cidx);
×
1264
            if (it == by_ckpt.end()) return false;
×
1265
            const auto &ds = *it->second;
×
1266
            if (ds.min_value != val || ds.max_value != val) return false;
×
1267
        }
1268
    }
×
1269
    return true;
×
1270
}
1271

1272
// Byte-range work unit for checkpoint-level parallelism. Each unit covers
1273
// one or more consecutive checkpoints from a single file. Decompression of
1274
// a single gz file is sequential per gzip stream, so splitting at
1275
// checkpoint-aligned byte offsets is what lets multiple workers share the
1276
// decode work for one file.
1277
struct ArrowWorkItem {
220✔
1278
    std::string file_path;
1279
    std::size_t start_byte = 0;
220✔
1280
    std::size_t end_byte = 0;
220✔
1281
    bool start_at_checkpoint = false;
220✔
1282
    bool end_at_checkpoint = false;
220✔
1283
    // When true, every kept chunk for this byte range is uniform-matching
1284
    // (dim_stats min == max == predicate literal for every AND-of-EQ leaf),
1285
    // so per-event predicate eval is skippable.
1286
    bool chunk_prune_only = false;
220✔
1287
    // Line-range work items override byte ranges: the worker passes these
1288
    // down as LINE_RANGE on the read, and the gzip stream resolves them to
1289
    // byte offsets via the checkpoint index. 0 = no line constraint.
1290
    std::size_t start_line = 0;
220✔
1291
    std::size_t end_line = 0;
220✔
1292
};
1293

1294
static std::vector<ArrowWorkItem> enumerate_work_items(
72✔
1295
    const std::vector<std::string> &files, const std::string &index_dir,
1296
    const std::string &query_str, std::size_t max_workers,
1297
    std::size_t clip_start_byte = 0, std::size_t clip_end_byte = 0,
1298
    std::size_t clip_start_line = 0, std::size_t clip_end_line = 0) {
1299
    namespace dft_internal =
1300
        dftracer::utils::utilities::composites::dft::internal;
1301
    namespace indexer_ns = dftracer::utils::utilities::indexer;
1302
    namespace indexing = dftracer::utils::utilities::composites::dft::indexing;
1303

1304
    std::vector<ArrowWorkItem> items;
72✔
1305
    items.reserve(files.size() * 4);
72!
1306

1307
    const bool has_line_clip = (clip_start_line > 0 || clip_end_line > 0);
72!
1308
    auto push_unsplit = [&](const std::string &fp) {
150✔
1309
        ArrowWorkItem item;
114✔
1310
        item.file_path = fp;
114!
1311
        item.start_line = clip_start_line;
114✔
1312
        item.end_line = clip_end_line;
114✔
1313
        items.push_back(std::move(item));
114!
1314
    };
114✔
1315

1316
    // Parse the query once. Pruner input copies a Query, so we keep the
1317
    // parsed form around to feed each ChunkPrunerInput without re-parsing.
1318
    std::optional<dftracer::utils::utilities::common::query::Query> parsed;
72✔
1319
    if (!query_str.empty()) {
72✔
1320
        auto r = dftracer::utils::utilities::common::query::Query::from_string(
3!
1321
            query_str);
6!
1322
        if (r) parsed = std::move(*r);
6!
1323
    }
6✔
1324

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

1335
    for (auto &entry : by_index) {
150!
1336
        const auto &index_path = entry.first;
78✔
1337
        const auto &file_idxs = entry.second;
78✔
1338
        if (!fs::exists(index_path)) {
78!
1339
            for (auto i : file_idxs) push_unsplit(files[i]);
156✔
1340
            continue;
62✔
1341
        }
31✔
1342
        std::unique_ptr<indexer_ns::IndexDatabase> idx_db;
16✔
1343
        try {
1344
            idx_db = std::make_unique<indexer_ns::IndexDatabase>(
16!
1345
                index_path,
8✔
1346
                dftracer::utils::rocksdb::RocksDatabase::OpenMode::ReadOnly);
16!
1347
        } catch (...) {
8✔
1348
            for (auto i : file_idxs) push_unsplit(files[i]);
×
1349
            continue;
1350
        }
×
1351

1352
        // Resolve fid + checkpoints per file (cheap queries).
1353
        struct FileCtx {
1354
            std::size_t file_idx;
1355
            int fid;
1356
            std::vector<indexer_ns::IndexerCheckpoint> ckpts;
1357
        };
1358
        std::vector<FileCtx> file_ctxs;
16✔
1359
        file_ctxs.reserve(file_idxs.size());
16!
1360
        for (auto i : file_idxs) {
44✔
1361
            FileCtx fc;
28✔
1362
            fc.file_idx = i;
28✔
1363
            fc.fid = idx_db->get_file_info_id(
42✔
1364
                indexer_ns::internal::get_logical_path(files[i]));
42!
1365
            if (fc.fid < 0) {
28!
1366
                push_unsplit(files[i]);
×
1367
                continue;
×
1368
            }
1369
            fc.ckpts = idx_db->query_checkpoints(fc.fid);
28!
1370
            if (fc.ckpts.empty()) {
28✔
1371
                push_unsplit(files[i]);
20!
1372
                continue;
20✔
1373
            }
1374
            std::sort(fc.ckpts.begin(), fc.ckpts.end(),
8!
1375
                      [](const auto &a, const auto &b) {
×
1376
                          return a.first_line_num < b.first_line_num;
×
1377
                      });
1378
            file_ctxs.push_back(std::move(fc));
8!
1379
        }
28✔
1380

1381
        // Batch-prune all files against the shared index: dim_stats and
1382
        // chunk_statistics are loaded in one RocksDB scan each instead of
1383
        // one scan per file.
1384
        std::vector<indexing::ChunkPrunerOutput> pruner_outs(file_ctxs.size());
16!
1385
        if (parsed && !file_ctxs.empty()) {
16!
1386
            indexing::ChunkPrunerBatchInput batch_in;
×
1387
            batch_in.index_path = index_path;
×
1388
            batch_in.external_db = idx_db.get();
×
1389
            batch_in.items.reserve(file_ctxs.size());
×
1390
            for (auto &fc : file_ctxs) {
×
1391
                batch_in.items.push_back({files[fc.file_idx], *parsed});
×
1392
            }
1393
            indexing::ChunkPrunerUtility pruner;
×
1394
            auto batch_out = pruner.process_batch(batch_in);
×
1395
            if (batch_out.success) {
×
1396
                pruner_outs = std::move(batch_out.outputs);
×
1397
            }
1398
        }
×
1399

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

1408
        for (std::size_t fc_idx = 0; fc_idx < file_ctxs.size(); ++fc_idx) {
24✔
1409
            auto &fc = file_ctxs[fc_idx];
8✔
1410
            const auto &fp = files[fc.file_idx];
8✔
1411

1412
            // Pruner chunk_idx semantics: 0-indexed over uncompressed
1413
            // slices. fc.ckpts holds gzip recovery points; recovery point
1414
            // fc.ckpts[k] sits at the START of pruner chunk (k+1). Pruner
1415
            // chunk 0 has no recovery point at its start (decoded from
1416
            // gzip stream start). Total pruner chunks = fc.ckpts.size()+1.
1417
            const std::size_t total_chunks = fc.ckpts.size() + 1;
8✔
1418
            auto chunk_start_byte = [&](std::uint64_t cidx) -> std::size_t {
16✔
1419
                if (cidx == 0) return 0;
12✔
1420
                return fc.ckpts[cidx - 1].uc_offset;
8✔
1421
            };
10✔
1422
            auto chunk_end_byte = [&](std::uint64_t cidx) -> std::size_t {
16✔
1423
                if (cidx == 0)
12✔
1424
                    return fc.ckpts.empty() ? 0 : fc.ckpts[0].uc_offset;
4!
1425
                std::size_t k = cidx - 1;
8✔
1426
                return fc.ckpts[k].uc_offset + fc.ckpts[k].uc_size;
8✔
1427
            };
10✔
1428
            // Line ranges for a chunk. Chunk 0 covers everything before the
1429
            // first recovery point; chunk k>=1 spans recovery point (k-1).
1430
            auto chunk_first_line = [&](std::uint64_t cidx) -> std::size_t {
16✔
1431
                if (cidx == 0) return 1;
12✔
1432
                return fc.ckpts[cidx - 1].first_line_num;
8✔
1433
            };
10✔
1434
            auto chunk_last_line = [&](std::uint64_t cidx) -> std::size_t {
16✔
1435
                if (cidx == 0) {
12✔
1436
                    if (fc.ckpts.empty()) return SIZE_MAX;
4!
1437
                    return fc.ckpts[0].first_line_num > 0
4!
1438
                               ? fc.ckpts[0].first_line_num - 1
4!
1439
                               : 0;
2✔
1440
                }
1441
                return fc.ckpts[cidx - 1].last_line_num;
8✔
1442
            };
10✔
1443

1444
            std::vector<std::uint64_t> keep_chunks;
8✔
1445
            keep_chunks.reserve(total_chunks);
8!
1446
            if (parsed) {
8!
1447
                const auto &pr = pruner_outs[fc_idx];
×
1448
                if (pr.success && !pr.file_may_match) {
×
1449
                    continue;  // whole file pruned
×
1450
                }
1451
                if (pr.success && !pr.candidate_checkpoints.empty() &&
×
1452
                    pr.candidate_checkpoints.size() < pr.total_checkpoints) {
×
1453
                    for (auto cidx : pr.candidate_checkpoints) {
×
1454
                        if (cidx < total_chunks) keep_chunks.push_back(cidx);
×
1455
                    }
1456
                    std::sort(keep_chunks.begin(), keep_chunks.end());
×
1457
                    keep_chunks.erase(
×
1458
                        std::unique(keep_chunks.begin(), keep_chunks.end()),
×
1459
                        keep_chunks.end());
×
1460
                } else {
1461
                    for (std::uint64_t c = 0; c < total_chunks; ++c)
×
1462
                        keep_chunks.push_back(c);
×
1463
                }
1464
            } else {
1465
                for (std::uint64_t c = 0; c < total_chunks; ++c)
24✔
1466
                    keep_chunks.push_back(c);
16!
1467
            }
1468

1469
            // Intersect with the user's line range so workers only touch
1470
            // chunks that actually overlap it. Each work item carries the
1471
            // sub-line-range; LINE_RANGE on the read maps it back to bytes
1472
            // via the same checkpoint table the gzip stream uses.
1473
            if (has_line_clip) {
8✔
1474
                std::size_t lo = clip_start_line > 0 ? clip_start_line : 1;
4!
1475
                std::size_t hi = clip_end_line > 0 ? clip_end_line : SIZE_MAX;
4!
1476
                std::vector<std::uint64_t> filtered;
4✔
1477
                filtered.reserve(keep_chunks.size());
4!
1478
                for (auto c : keep_chunks) {
12✔
1479
                    std::size_t cf = chunk_first_line(c);
8!
1480
                    std::size_t cl = chunk_last_line(c);
8!
1481
                    if (cl < lo || cf > hi) continue;
8✔
1482
                    filtered.push_back(c);
4!
1483
                }
1484
                keep_chunks = std::move(filtered);
4✔
1485
            }
4✔
1486

1487
            if (keep_chunks.empty()) continue;
8✔
1488

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

1499
            std::size_t target_ranges = std::max<std::size_t>(1, max_workers);
8!
1500
            std::size_t per_range = std::max<std::size_t>(
8!
1501
                1, (keep_chunks.size() + target_ranges - 1) / target_ranges);
8✔
1502

1503
            std::size_t group_start = 0;
8✔
1504
            while (group_start < keep_chunks.size()) {
20✔
1505
                std::size_t group_end = group_start;
12✔
1506
                std::size_t emitted = 0;
12✔
1507
                while (group_end < keep_chunks.size() && emitted < per_range) {
24✔
1508
                    if (group_end > group_start &&
12!
1509
                        keep_chunks[group_end] !=
×
1510
                            keep_chunks[group_end - 1] + 1) {
×
1511
                        break;
×
1512
                    }
1513
                    ++group_end;
12✔
1514
                    ++emitted;
12✔
1515
                }
1516
                std::uint64_t scidx = keep_chunks[group_start];
12✔
1517
                std::uint64_t ecidx = keep_chunks[group_end - 1];
12✔
1518
                std::size_t start_byte = chunk_start_byte(scidx);
12!
1519
                std::size_t end_byte = chunk_end_byte(ecidx);
12!
1520
                // start_at_checkpoint: a gzip recovery point sits at
1521
                // start_byte (true for any cidx>=1; false for the implicit
1522
                // chunk 0 which decodes from stream start).
1523
                bool start_at_checkpoint = (scidx >= 1);
12✔
1524
                bool end_at_checkpoint = (group_end < keep_chunks.size());
12✔
1525
                if (has_line_clip) {
12✔
1526
                    std::size_t lo = clip_start_line > 0 ? clip_start_line : 1;
4!
1527
                    std::size_t hi =
4✔
1528
                        clip_end_line > 0 ? clip_end_line : SIZE_MAX;
4!
1529
                    std::size_t cluster_first = chunk_first_line(scidx);
4!
1530
                    std::size_t cluster_last = chunk_last_line(ecidx);
4!
1531
                    std::size_t item_start =
2✔
1532
                        std::max<std::size_t>(lo, cluster_first);
4!
1533
                    std::size_t item_end =
2✔
1534
                        std::min<std::size_t>(hi, cluster_last);
4!
1535
                    if (item_start > item_end) {
4!
1536
                        group_start = group_end;
×
1537
                        continue;
×
1538
                    }
1539
                    ArrowWorkItem item;
4✔
1540
                    item.file_path = fp;
4!
1541
                    item.chunk_prune_only = file_pure_match;
4✔
1542
                    item.start_line = item_start;
4✔
1543
                    item.end_line = item_end;
4✔
1544
                    items.push_back(std::move(item));
4!
1545
                    group_start = group_end;
4✔
1546
                    continue;
2✔
1547
                }
4✔
1548
                if (clip_end_byte > clip_start_byte) {
8✔
1549
                    if (start_byte < clip_start_byte) {
4✔
1550
                        start_byte = clip_start_byte;
×
1551
                        start_at_checkpoint = false;
×
1552
                    }
1553
                    if (end_byte > clip_end_byte) {
4✔
1554
                        end_byte = clip_end_byte;
4✔
1555
                        end_at_checkpoint = false;
4✔
1556
                    }
2✔
1557
                    if (start_byte >= end_byte) {
4✔
1558
                        group_start = group_end;
2✔
1559
                        continue;
2✔
1560
                    }
1561
                }
1✔
1562
                items.push_back({fp, start_byte, end_byte, start_at_checkpoint,
6!
1563
                                 end_at_checkpoint, file_pure_match});
6✔
1564
                group_start = group_end;
6✔
1565
            }
1566
        }
8✔
1567
    }
16✔
1568
    return items;
108✔
1569
}
72!
1570

1571
static CoroTask<void> send_work_items_to_channel(
464!
1572
    std::shared_ptr<dftracer::utils::coro::Channel<ArrowWorkItem>> chan,
1573
    const std::vector<ArrowWorkItem> *items, std::atomic<bool> *cancelled) {
36!
1574
    for (const auto &it : *items) {
346✔
1575
        if (cancelled->load(std::memory_order_acquire)) break;
186!
1576
        if (!co_await chan->send(it)) break;
284!
1577
    }
62✔
1578
    chan->close();
36!
1579
    co_return;
36✔
1580
}
320!
1581

1582
static CoroTask<void> checkpoint_worker(
310!
1583
    std::shared_ptr<dftracer::utils::coro::Channel<ArrowWorkItem>> work_chan,
1584
    dftracer::utils::coro::Channel<ArrowExportResult> *out_chan,
1585
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
1586
    ReadConfig rc, std::size_t batch_size, bool normalize,
1587
    std::atomic<bool> *cancelled) {
52!
1588
    dftracer::utils::coro::ChannelProducer<ArrowExportResult> producer(
104!
1589
        out_chan);
52✔
1590
    auto guard = producer.guard();
52!
1591

1592
    // Cache readers keyed by file path so we don't re-probe the same file
1593
    // when successive work items land on it.
1594
    std::unordered_map<std::string, std::shared_ptr<TraceReader>> readers;
52✔
1595

1596
    while (auto item = co_await work_chan->receive()) {
349!
1597
        if (cancelled->load(std::memory_order_acquire)) co_return;
184!
1598

1599
        auto &reader_ptr = readers[item->file_path];
184!
1600
        if (!reader_ptr) {
184✔
1601
            TraceReaderConfig cfg;
62✔
1602
            cfg.file_path = item->file_path;
62!
1603
            cfg.index_dir = index_dir;
62!
1604
            cfg.checkpoint_size = checkpoint_size;
62✔
1605
            cfg.auto_build_index = auto_build_index;
62✔
1606
            reader_ptr = std::make_shared<TraceReader>(std::move(cfg));
62!
1607
        }
62✔
1608

1609
        ReadConfig local_rc = rc;
184!
1610
        if (item->start_line > 0 || item->end_line > 0) {
184!
1611
            // Line-range work items: the read drives off LINE_RANGE; the
1612
            // gzip stream resolves it back to byte offsets via checkpoints.
1613
            local_rc.start_line = item->start_line;
125✔
1614
            local_rc.end_line = item->end_line;
125✔
1615
            local_rc.start_byte = 0;
125✔
1616
            local_rc.end_byte = 0;
125✔
1617
            local_rc.start_at_checkpoint = false;
125✔
1618
            local_rc.end_at_checkpoint = false;
125✔
1619
        } else {
125✔
1620
            local_rc.start_byte = item->start_byte;
59✔
1621
            local_rc.end_byte = item->end_byte;
59✔
1622
            local_rc.start_at_checkpoint = item->start_at_checkpoint;
59✔
1623
            local_rc.end_at_checkpoint = item->end_at_checkpoint;
59✔
1624
        }
1625
        // Pruning already happened at enumeration time; avoid the per-
1626
        // work-item RocksDB opens that would otherwise dwarf the actual
1627
        // read cost at directory scale (256 files * N ranges).
1628
        local_rc.skip_pruning = true;
184✔
1629
        // chunks pre-classified as uniform-matching skip per-event eval.
1630
        if (item->chunk_prune_only) local_rc.chunk_prune_only = true;
184!
1631

1632
        if (!normalize) {
184!
1633
            auto batch_gen = reader_ptr->read_arrow(local_rc, batch_size);
184!
1634
            while (auto batch_opt = co_await batch_gen.next()) {
567!
1635
                if (cancelled->load(std::memory_order_acquire)) co_return;
241!
1636
                if (!co_await producer.send(std::move(*batch_opt))) co_return;
321!
1637
            }
303!
1638
            continue;
62✔
1639
        }
62✔
1640

1641
        auto gen = reader_ptr->read_json(local_rc);
×
1642
        RecordBatchBuilder builder;
×
1643
        builder.reserve(batch_size);
×
1644
        StringArena arena;
×
1645

1646
        while (auto opt = co_await gen.next()) {
×
1647
            if (cancelled->load(std::memory_order_acquire)) co_return;
×
1648
            if (!build_arrow_row(builder, *opt->parser, arena, normalize))
×
1649
                continue;
1650
            if (builder.num_rows() >= batch_size) {
×
1651
                auto result = builder.finish();
×
1652
                arena.clear();
×
1653
                if (!co_await producer.send(std::move(result))) co_return;
×
1654
                if (!builder.is_schema_locked()) builder.lock_schema();
×
1655
                builder.reset(true);
×
1656
                builder.reserve(batch_size);
×
1657
            }
×
1658
        }
×
1659
        if (builder.num_rows() > 0) {
×
1660
            if (!co_await producer.send(builder.finish())) co_return;
×
1661
        }
1662
    }
236!
1663
    co_return;
52✔
1664
}
1,280!
1665

1666
static CoroTask<void> spawn_arrow_producers(
216!
1667
    CoroScope &child,
1668
    dftracer::utils::coro::Channel<ArrowExportResult> *out_chan,
1669
    const std::vector<ArrowWorkItem> *work_items, const std::string *index_dir,
1670
    std::size_t checkpoint_size, bool auto_build_index, const ReadConfig *rc,
1671
    std::size_t batch_size, bool normalize, std::atomic<bool> *cancelled_ptr,
1672
    std::size_t max_workers) {
36!
1673
    std::size_t num_workers = std::min(work_items->size(), max_workers);
36!
1674
    if (num_workers == 0) num_workers = 1;
36!
1675
    auto work_chan =
36✔
1676
        dftracer::utils::coro::make_channel<ArrowWorkItem>(num_workers);
36!
1677

1678
    for (std::size_t i = 0; i < num_workers; ++i) {
88✔
1679
        child.spawn([out_chan, wc = work_chan, idx = *index_dir,
153!
1680
                     checkpoint_size, auto_build_index, r = *rc, batch_size,
156!
1681
                     normalize, cancelled_ptr](CoroScope &) {
255✔
1682
            return checkpoint_worker(wc, out_chan, idx, checkpoint_size,
154!
1683
                                     auto_build_index, r, batch_size, normalize,
101!
1684
                                     cancelled_ptr);
101!
1685
        });
1686
    }
52✔
1687

1688
    child.spawn([wc = work_chan, work_items, cancelled_ptr](CoroScope &) {
108!
1689
        return send_work_items_to_channel(wc, work_items, cancelled_ptr);
72!
1690
    });
1691
    co_return;
72✔
1692
}
108!
1693

1694
static CoroTask<void> produce_arrow_batches_for_files(
294!
1695
    CoroScope &scope, ArrowIteratorState *sp, std::vector<std::string> files,
1696
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
1697
    ReadConfig rc, std::size_t batch_size, bool normalize,
1698
    std::size_t max_workers) {
37!
1699
    try {
1700
        if (files.empty()) {
109✔
1701
            sp->channel->close();
1!
1702
            co_return;
38✔
1703
        }
1704

1705
        auto work_items = enumerate_work_items(
216!
1706
            files, index_dir, rc.query, max_workers, rc.start_byte, rc.end_byte,
108✔
1707
            rc.start_line, rc.end_line);
108✔
1708
        if (work_items.empty()) {
108!
1709
            sp->channel->close();
×
1710
            co_return;
1711
        }
1712

1713
        auto *chan_ptr = sp->channel.get();
108✔
1714
        auto *cancelled_ptr = &sp->cancelled;
108✔
1715

1716
        co_await scope.scope([chan_ptr, &work_items, &index_dir,
1,080!
1717
                              checkpoint_size, auto_build_index, &rc,
216✔
1718
                              batch_size, normalize, cancelled_ptr,
324✔
1719
                              max_workers](CoroScope &child) -> CoroTask<void> {
144!
1720
            co_await spawn_arrow_producers(
288!
1721
                child, chan_ptr, &work_items, &index_dir, checkpoint_size,
108✔
1722
                auto_build_index, &rc, batch_size, normalize, cancelled_ptr,
108✔
1723
                max_workers);
108✔
1724
        });
144!
1725
    } catch (...) {
36!
1726
        sp->set_error(std::current_exception());
×
1727
    }
×
1728
}
146!
1729

1730
static CoroTask<void> produce_arrow_batches_parallel(
180!
1731
    CoroScope &scope, ArrowIteratorState *sp, std::string dir_path,
1732
    std::string index_dir, std::size_t checkpoint_size, bool auto_build_index,
1733
    ReadConfig rc, std::size_t batch_size, bool normalize,
1734
    std::size_t max_workers) {
15!
1735
    try {
1736
        PatternDirectoryScannerUtility scanner;
45!
1737
        auto scan_input = PatternDirectoryScannerUtilityInput(
90!
1738
            dir_path, {".pfw", ".pfw.gz"}, true, false);
45!
1739
        auto entries = co_await scope.spawn(scanner, scan_input);
75!
1740

1741
        std::vector<std::string> files;
45✔
1742
        files.reserve(entries.size());
45✔
1743
        for (auto &e : entries) files.push_back(e.path.string());
54!
1744
        std::sort(files.begin(), files.end());
15✔
1745

1746
        co_await produce_arrow_batches_for_files(
105!
1747
            scope, sp, std::move(files), std::move(index_dir), checkpoint_size,
45✔
1748
            auto_build_index, std::move(rc), batch_size, normalize,
45✔
1749
            max_workers);
45✔
1750
    } catch (...) {
45✔
1751
        sp->set_error(std::current_exception());
×
1752
    }
×
1753
}
210!
1754

1755
CoroTask<void> produce_arrow_batches(
136!
1756
    std::shared_ptr<ArrowIteratorState> state,
1757
    dftracer::utils::coro::ChannelProducer<ArrowExportResult> producer,
1758
    TraceReaderConfig cfg, ReadConfig rc, std::size_t batch_size,
1759
    bool flatten_objects = false, bool normalize = false) {
17!
1760
    (void)flatten_objects;
1761

1762
    auto guard = producer.guard();
51!
1763
    try {
1764
        TraceReader reader(std::move(cfg));
51!
1765

1766
        if (!normalize) {
51!
1767
            auto batch_gen = reader.read_arrow(rc, batch_size);
51!
1768
            while (auto batch_opt = co_await batch_gen.next()) {
213!
1769
                if (state->cancelled.load(std::memory_order_acquire)) break;
97✔
1770
                auto result_bytes =
96✔
1771
                    dftracer::utils::python::byte_size(*batch_opt);
96!
1772
                state->bytes_in_queue.fetch_add(result_bytes,
96✔
1773
                                                std::memory_order_acq_rel);
1774
                if (!co_await producer.send(std::move(*batch_opt))) break;
128!
1775
            }
113!
1776
            co_return;
17✔
1777
        }
17✔
1778

1779
        auto gen = reader.read_json(rc);
×
1780
        RecordBatchBuilder builder;
×
1781
        builder.reserve(batch_size);
×
1782

1783
        StringArena arena;
×
1784

1785
        while (auto opt = co_await gen.next()) {
×
1786
            if (state->cancelled.load(std::memory_order_acquire)) break;
×
1787
            if (!build_arrow_row(builder, *opt->parser, arena, normalize))
×
1788
                continue;
1789

1790
            if (builder.num_rows() >= batch_size) {
×
1791
                auto result = builder.finish();
×
1792
                arena.clear();
×
1793
                auto result_bytes = dftracer::utils::python::byte_size(result);
×
1794
                state->bytes_in_queue.fetch_add(result_bytes,
1795
                                                std::memory_order_acq_rel);
1796
                if (!co_await producer.send(std::move(result))) break;
×
1797
                if (!builder.is_schema_locked()) {
×
1798
                    builder.lock_schema();
1799
                }
1800
                builder.reset(true);
×
1801
                builder.reserve(batch_size);
×
1802
            }
×
1803
        }
×
1804

1805
        if (builder.num_rows() > 0 &&
×
1806
            !state->cancelled.load(std::memory_order_acquire)) {
1807
            auto result = builder.finish();
×
1808
            auto result_bytes = dftracer::utils::python::byte_size(result);
×
1809
            state->bytes_in_queue.fetch_add(result_bytes,
1810
                                            std::memory_order_acq_rel);
1811
            co_await producer.send(std::move(result));
×
1812
        }
×
1813
    } catch (...) {
17!
1814
        state->set_error(std::current_exception());
×
1815
    }
×
1816
}
375!
1817

1818
#endif  // DFTRACER_UTILS_ENABLE_ARROW
1819

1820
#ifdef DFTRACER_UTILS_ENABLE_ARROW_IPC
1821

1822
struct WriteArrowStats {
39✔
1823
    std::unordered_map<std::string, PartitionWriteStats> partitions;
1824
    int64_t total_rows = 0;
39✔
1825
    int64_t total_uncompressed_bytes = 0;
39✔
1826
};
1827

1828
struct WriteArrowResult {
52✔
1829
    WriteArrowStats stats;
1830
    std::string error;
1831
    std::uint64_t chunks_scanned = 0;
39✔
1832
    std::uint64_t chunks_skipped = 0;
39✔
1833
};
1834

1835
CoroTask<WriteArrowResult> write_arrow_pipeline(
190!
1836
    std::string file_path, std::string index_path, std::size_t checkpoint_size,
1837
    std::vector<ViewDefinition> views, std::string output_path,
1838
    int64_t chunk_size_bytes, IpcCompression compression,
1839
    std::size_t event_batch_size) {
13!
1840
    namespace dft_internal =
1841
        dftracer::utils::utilities::composites::dft::internal;
1842
    WriteArrowResult result;
13✔
1843

1844
    try {
1845
        if (views.empty()) {
13✔
1846
            views.push_back(ViewDefinition().with_name("all"));
8!
1847
        }
8✔
1848

1849
        std::string resolved_index =
13✔
1850
            index_path.empty()
26!
1851
                ? dft_internal::determine_index_path(file_path, "")
13!
1852
                : index_path;
×
1853

1854
        auto meta_input = MetadataCollectorUtilityInput::from_file(file_path)
26!
1855
                              .with_checkpoint_size(checkpoint_size)
13!
1856
                              .with_index(resolved_index);
13!
1857
        auto metadata = co_await MetadataCollectorUtility{}.process(meta_input);
39!
1858
        if (!metadata.success) {
13!
1859
            result.error =
1860
                "Failed to collect metadata: " + metadata.error_message;
×
1861
            co_return result;
×
1862
        }
1863

1864
        for (const auto &view : views) {
99✔
1865
            std::string view_output = output_path;
30!
1866
            if (views.size() > 1 || view.name != "all") {
30✔
1867
                view_output = output_path + "/" + view.name;
42!
1868
            }
6✔
1869

1870
            PartitionWriter writer;
42!
1871
            int rc_open = co_await writer.open(view_output, chunk_size_bytes,
98!
1872
                                               compression);
42✔
1873
            if (rc_open != 0) {
42!
1874
                result.error =
1875
                    "Failed to open partition writer for view: " + view.name;
×
1876
                co_return result;
×
1877
            }
1878

1879
            ViewBuilderInput builder_input;
42✔
1880
            builder_input.with_view(view)
42✔
1881
                .with_file_path(file_path)
14!
1882
                .with_index_path(resolved_index)
14!
1883
                .with_uncompressed_size(metadata.uncompressed_size)
14!
1884
                .with_num_checkpoints(metadata.num_checkpoints);
14✔
1885

1886
            auto build_output =
42✔
1887
                co_await ViewBuilderUtility{}.process(builder_input);
56!
1888
            if (!build_output.success) {
58!
1889
                result.error = "ViewBuilder failed for view: " + view.name;
×
1890
                co_return result;
×
1891
            }
1892

1893
            result.chunks_skipped += build_output.skipped_checkpoints;
58✔
1894

1895
            if (!build_output.file_may_match) {
58✔
1896
                auto stats = co_await writer.close();
24!
1897
                result.stats.partitions[view.name] = std::move(stats);
6!
1898
                continue;
1899
            }
6✔
1900

1901
            RecordBatchBuilder builder;
40!
1902
            bool schema_locked = false;
40✔
1903

1904
            for (const auto &candidate : build_output.candidates) {
64!
1905
                ViewReaderInput reader_input;
40✔
1906
                reader_input.with_file_path(file_path)
40✔
1907
                    .with_index_path(resolved_index)
8!
1908
                    .with_checkpoint_size(checkpoint_size)
8!
1909
                    .with_byte_range(candidate.start_byte, candidate.end_byte)
8!
1910
                    .with_checkpoint_idx(candidate.checkpoint_idx)
8!
1911
                    .with_event_batch_size(event_batch_size)
8!
1912
                    .with_view(view);
8!
1913
                reader_input.query = view.query;
8✔
1914

1915
                ViewReaderUtility reader;
40!
1916
                auto gen = reader.process(reader_input);
40!
1917
                while (auto opt = co_await gen.next()) {
64!
1918
                    auto arrow_batch = opt->to_arrow(builder);
24!
1919
                    int rc_write = co_await writer.write_batch(arrow_batch);
32!
1920
                    if (rc_write != 0) {
8!
1921
                        result.error =
1922
                            "Failed to write batch for view: " + view.name;
×
1923
                        co_return result;
×
1924
                    }
1925
                    if (!schema_locked) {
8!
1926
                        builder.lock_schema();
8✔
1927
                        schema_locked = true;
8✔
1928
                    }
8✔
1929
                    builder.reset(true);
8!
1930
                }
32!
1931
                result.chunks_scanned++;
8✔
1932
            }
24✔
1933

1934
            auto stats = co_await writer.close();
32!
1935
            result.stats.partitions[view.name] = std::move(stats);
8!
1936
            result.stats.total_rows +=
8✔
1937
                result.stats.partitions[view.name].total_rows;
8!
1938
            result.stats.total_uncompressed_bytes +=
8✔
1939
                result.stats.partitions[view.name].total_uncompressed_bytes;
8!
1940
        }
86✔
1941
    } catch (const std::exception &e) {
69!
1942
        result.error = e.what();
×
1943
    }
×
1944
    co_return result;
13!
1945
}
405!
1946

1947
struct ViewChunkInfo {
1948
    std::uint64_t checkpoint_idx;
1949
    std::size_t start_byte;
1950
    std::size_t end_byte;
1951
};
1952

1953
struct GetViewChunksResult {
12✔
1954
    std::vector<ViewChunkInfo> chunks;
1955
    std::uint64_t total_checkpoints = 0;
9✔
1956
    std::uint64_t skipped_checkpoints = 0;
9✔
1957
    bool file_may_match = false;
9✔
1958
    std::string error;
1959
};
1960

1961
CoroTask<GetViewChunksResult> get_view_chunks_pipeline(
30!
1962
    std::string file_path, std::string index_path, std::size_t checkpoint_size,
1963
    ViewDefinition view) {
3!
1964
    namespace dft_internal =
1965
        dftracer::utils::utilities::composites::dft::internal;
1966
    GetViewChunksResult result;
3✔
1967

1968
    try {
1969
        std::string resolved_index =
3✔
1970
            index_path.empty()
6!
1971
                ? dft_internal::determine_index_path(file_path, "")
3!
1972
                : index_path;
×
1973

1974
        auto meta_input = MetadataCollectorUtilityInput::from_file(file_path)
3!
1975
                              .with_checkpoint_size(checkpoint_size)
3✔
1976
                              .with_index(resolved_index);
3!
1977
        auto metadata = co_await MetadataCollectorUtility{}.process(meta_input);
9!
1978
        if (!metadata.success) {
9!
1979
            result.error =
1980
                "Failed to collect metadata: " + metadata.error_message;
×
1981
            co_return result;
×
1982
        }
1983

1984
        ViewBuilderInput builder_input;
9✔
1985
        builder_input.with_view(view)
9✔
1986
            .with_file_path(file_path)
3!
1987
            .with_index_path(resolved_index)
3!
1988
            .with_uncompressed_size(metadata.uncompressed_size)
3!
1989
            .with_num_checkpoints(metadata.num_checkpoints);
3✔
1990

1991
        auto build_output =
9✔
1992
            co_await ViewBuilderUtility{}.process(builder_input);
12!
1993
        if (!build_output.success) {
3!
1994
            result.error = "ViewBuilder failed";
×
1995
            co_return result;
×
1996
        }
1997

1998
        result.file_may_match = build_output.file_may_match;
3✔
1999
        result.total_checkpoints = build_output.total_checkpoints;
3✔
2000
        result.skipped_checkpoints = build_output.skipped_checkpoints;
3✔
2001

2002
        for (const auto &candidate : build_output.candidates) {
7✔
2003
            result.chunks.push_back({candidate.checkpoint_idx,
12!
2004
                                     candidate.start_byte, candidate.end_byte});
8✔
2005
        }
4✔
2006
    } catch (const std::exception &e) {
9!
2007
        result.error = e.what();
×
2008
    }
×
2009
    co_return result;
3!
2010
}
39!
2011

2012
struct WriteViewChunkResult {
16✔
2013
    std::string output_file;
2014
    std::uint64_t events_matched = 0;
12✔
2015
    std::uint64_t events_scanned = 0;
12✔
2016
    int64_t rows_written = 0;
12✔
2017
    int64_t bytes_written = 0;
12✔
2018
    std::string error;
2019
};
2020

2021
CoroTask<WriteViewChunkResult> write_view_chunk_pipeline(
56!
2022
    std::string file_path, std::string index_path, std::size_t checkpoint_size,
2023
    ViewDefinition view, std::uint64_t checkpoint_idx, std::size_t start_byte,
2024
    std::size_t end_byte, std::string output_file, IpcCompression compression,
2025
    std::size_t event_batch_size) {
4!
2026
    namespace dft_internal =
2027
        dftracer::utils::utilities::composites::dft::internal;
2028
    WriteViewChunkResult result;
4✔
2029
    result.output_file = output_file;
4!
2030

2031
    try {
2032
        std::string resolved_index =
4✔
2033
            index_path.empty()
8!
2034
                ? dft_internal::determine_index_path(file_path, "")
4!
2035
                : index_path;
×
2036

2037
        dftracer::utils::utilities::common::arrow::IpcWriter writer;
4!
2038
        int rc_open = co_await writer.open(output_file, compression);
12!
2039
        if (rc_open != 0) {
20!
2040
            result.error = "Failed to open output file";
×
2041
            co_return result;
×
2042
        }
2043

2044
        ViewReaderInput reader_input;
20✔
2045
        reader_input.with_file_path(file_path)
20✔
2046
            .with_index_path(resolved_index)
4!
2047
            .with_checkpoint_size(checkpoint_size)
4!
2048
            .with_byte_range(start_byte, end_byte)
4!
2049
            .with_checkpoint_idx(checkpoint_idx)
4!
2050
            .with_event_batch_size(event_batch_size)
4!
2051
            .with_view(view);
4!
2052
        reader_input.query = view.query;
4✔
2053

2054
        RecordBatchBuilder builder;
20!
2055
        bool schema_locked = false;
20✔
2056

2057
        ViewReaderUtility reader;
20!
2058
        auto gen = reader.process(reader_input);
20!
2059
        while (auto opt = co_await gen.next()) {
32!
2060
            result.events_matched += opt->events_matched;
12✔
2061
            result.events_scanned += opt->events_scanned;
12✔
2062
            auto batch = opt->to_arrow(builder);
12!
2063
            if (batch.valid()) {
12!
2064
                result.rows_written += batch.num_rows();
12✔
2065
                int rc = co_await writer.write_batch(batch);
16!
2066
                if (rc != 0) {
4!
2067
                    result.error = "Failed to write batch";
×
2068
                    co_return result;
×
2069
                }
2070
                if (!schema_locked) {
4!
2071
                    builder.lock_schema();
4✔
2072
                    schema_locked = true;
4✔
2073
                }
4✔
2074
                builder.reset(true);
4!
2075
            }
4!
2076
        }
16!
2077

2078
        int rc = co_await writer.close();
16!
2079
        if (rc != 0) {
4!
2080
            result.error = "Failed to close output file";
×
2081
        }
2082
    } catch (const std::exception &e) {
20!
2083
        result.error = e.what();
×
2084
    }
×
2085
    co_return result;
4!
2086
}
116!
2087

2088
struct ChunkDescriptor {
2089
    std::uint64_t checkpoint_idx;
2090
    std::size_t start_byte;
2091
    std::size_t end_byte;
2092
    std::string output_file;
2093
};
2094

2095
struct WriteViewChunksResult {
3✔
2096
    std::vector<WriteViewChunkResult> results;
2097
    int64_t total_rows = 0;
3✔
2098
    int64_t total_events_matched = 0;
3✔
2099
};
2100

2101
CoroTask<WriteViewChunksResult> write_view_chunks_pipeline(
8!
2102
    std::string file_path, std::string index_path, std::size_t checkpoint_size,
2103
    ViewDefinition view, std::vector<ChunkDescriptor> chunks,
2104
    IpcCompression compression, std::size_t event_batch_size) {
1!
2105
    WriteViewChunksResult result;
3✔
2106

2107
    if (chunks.empty()) {
3!
2108
        co_return result;
1!
2109
    }
2110

2111
    std::vector<CoroTask<WriteViewChunkResult>> tasks;
3✔
2112
    tasks.reserve(chunks.size());
3!
2113

2114
    for (const auto &chunk : chunks) {
6✔
2115
        tasks.push_back(write_view_chunk_pipeline(
6!
2116
            file_path, index_path, checkpoint_size, view, chunk.checkpoint_idx,
3!
2117
            chunk.start_byte, chunk.end_byte, chunk.output_file, compression,
3!
2118
            event_batch_size));
3✔
2119
    }
3✔
2120

2121
    result.results = co_await when_all(std::move(tasks));
4!
2122

2123
    for (const auto &r : result.results) {
4✔
2124
        result.total_rows += r.rows_written;
3✔
2125
        result.total_events_matched += r.events_matched;
3✔
2126
    }
3✔
2127

2128
    co_return result;
1!
2129
}
7!
2130

2131
#endif  // DFTRACER_UTILS_ENABLE_ARROW_IPC
2132

2133
TraceReaderConfig build_config(TraceReaderObject *self) {
302✔
2134
    TraceReaderConfig cfg;
302✔
2135
    cfg.file_path = PyUnicode_AsUTF8(self->file_path);
302!
2136
    const char *idx = PyUnicode_AsUTF8(self->index_dir);
302!
2137
    if (idx) cfg.index_dir = idx;
302!
2138
    cfg.checkpoint_size = self->checkpoint_size;
302✔
2139
    cfg.auto_build_index = self->auto_build_index != 0;
302✔
2140
    return cfg;
302✔
2141
}
151!
2142

2143
static Runtime *get_runtime(TraceReaderObject *self) {
300✔
2144
    if (self->runtime_obj) {
300✔
2145
        return ((RuntimeObject *)self->runtime_obj)->runtime.get();
68✔
2146
    }
2147
    return get_default_runtime();
232✔
2148
}
150✔
2149

2150
static TraceReaderIteratorObject *make_memoryview_iterator(
142✔
2151
    std::shared_ptr<MemoryViewBatchIteratorState> state) {
2152
    TraceReaderIteratorObject *it =
71✔
2153
        (TraceReaderIteratorObject *)TraceReaderIteratorType.tp_alloc(
142✔
2154
            &TraceReaderIteratorType, 0);
2155
    if (!it) return NULL;
142✔
2156
    new (&it->batch_state)
142✔
2157
        std::shared_ptr<MemoryViewBatchIteratorState>(std::move(state));
142✔
2158
    it->current_batch = NULL;
142✔
2159
    it->batch_index = 0;
142✔
2160
    new (&it->json_dict_state) std::shared_ptr<JsonDictIteratorState>();
142✔
2161
    new (&it->json_dict_current_batch) std::shared_ptr<JsonDictBatch>();
142✔
2162
    it->json_dict_index = 0;
142✔
2163
#ifdef DFTRACER_UTILS_ENABLE_ARROW
2164
    new (&it->arrow_state) std::shared_ptr<ArrowIteratorState>();
142✔
2165
#endif
2166
    it->mode = IteratorMode::MEMORYVIEW;
142✔
2167
    return it;
142✔
2168
}
71✔
2169

2170
static TraceReaderIteratorObject *make_json_dict_iterator(
14✔
2171
    std::shared_ptr<JsonDictIteratorState> state) {
2172
    TraceReaderIteratorObject *it =
7✔
2173
        (TraceReaderIteratorObject *)TraceReaderIteratorType.tp_alloc(
14✔
2174
            &TraceReaderIteratorType, 0);
2175
    if (!it) return NULL;
14✔
2176
    new (&it->batch_state) std::shared_ptr<MemoryViewBatchIteratorState>();
14✔
2177
    it->current_batch = NULL;
14✔
2178
    it->batch_index = 0;
14✔
2179
    new (&it->json_dict_state)
14✔
2180
        std::shared_ptr<JsonDictIteratorState>(std::move(state));
14✔
2181
    new (&it->json_dict_current_batch) std::shared_ptr<JsonDictBatch>();
14✔
2182
    it->json_dict_index = 0;
14✔
2183
#ifdef DFTRACER_UTILS_ENABLE_ARROW
2184
    new (&it->arrow_state) std::shared_ptr<ArrowIteratorState>();
14✔
2185
#endif
2186
    it->mode = IteratorMode::JSON_DICT;
14✔
2187
    return it;
14✔
2188
}
7✔
2189

2190
#ifdef DFTRACER_UTILS_ENABLE_ARROW
2191
static TraceReaderIteratorObject *make_arrow_iterator(
54✔
2192
    std::shared_ptr<ArrowIteratorState> state) {
2193
    TraceReaderIteratorObject *it =
27✔
2194
        (TraceReaderIteratorObject *)TraceReaderIteratorType.tp_alloc(
54✔
2195
            &TraceReaderIteratorType, 0);
2196
    if (!it) return NULL;
54✔
2197
    new (&it->batch_state) std::shared_ptr<MemoryViewBatchIteratorState>();
54✔
2198
    it->current_batch = NULL;
54✔
2199
    it->batch_index = 0;
54✔
2200
    new (&it->json_dict_state) std::shared_ptr<JsonDictIteratorState>();
54✔
2201
    new (&it->json_dict_current_batch) std::shared_ptr<JsonDictBatch>();
54✔
2202
    it->json_dict_index = 0;
54✔
2203
    new (&it->arrow_state)
54✔
2204
        std::shared_ptr<ArrowIteratorState>(std::move(state));
54✔
2205
    it->mode = IteratorMode::ARROW;
54✔
2206
    return it;
54✔
2207
}
27✔
2208
#endif
2209

2210
}  // namespace
2211

2212
static void TraceReader_dealloc(TraceReaderObject *self) {
316✔
2213
    Py_XDECREF(self->file_path);
316✔
2214
    Py_XDECREF(self->index_dir);
316✔
2215
    Py_XDECREF(self->runtime_obj);
316✔
2216
    Py_TYPE(self)->tp_free((PyObject *)self);
316✔
2217
}
316✔
2218

2219
static PyObject *TraceReader_new(PyTypeObject *type, PyObject *args,
316✔
2220
                                 PyObject *kwds) {
2221
    TraceReaderObject *self = (TraceReaderObject *)type->tp_alloc(type, 0);
316✔
2222
    if (self) {
316✔
2223
        self->file_path = NULL;
316✔
2224
        self->index_dir = NULL;
316✔
2225
        self->checkpoint_size = 32 * 1024 * 1024;
316✔
2226
        self->auto_build_index = 0;
316✔
2227
        self->has_index = 0;
316✔
2228
        self->runtime_obj = NULL;
316✔
2229
    }
158✔
2230
    return (PyObject *)self;
316✔
2231
}
2232

2233
static int TraceReader_init(TraceReaderObject *self, PyObject *args,
316✔
2234
                            PyObject *kwds) {
2235
    static const char *kwlist[] = {
2236
        "path",    "index_dir", "checkpoint_size", "auto_build_index",
2237
        "runtime", NULL};
2238

2239
    const char *file_path;
2240
    const char *index_dir = "";
316✔
2241
    std::size_t checkpoint_size = 32 * 1024 * 1024;
316✔
2242
    int auto_build_index = 0;
316✔
2243
    PyObject *runtime_arg = NULL;
316✔
2244

2245
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|snpO", (char **)kwlist,
316!
2246
                                     &file_path, &index_dir, &checkpoint_size,
2247
                                     &auto_build_index, &runtime_arg)) {
2248
        return -1;
×
2249
    }
2250

2251
    if (runtime_arg && runtime_arg != Py_None) {
316✔
2252
        if (PyObject_TypeCheck(runtime_arg, &RuntimeType)) {
64!
2253
            // Direct C++ Runtime object
2254
            Py_INCREF(runtime_arg);
×
2255
            self->runtime_obj = runtime_arg;
×
2256
        } else {
2257
            // Python wrapper, extract _native attribute
2258
            PyObject *native = PyObject_GetAttrString(runtime_arg, "_native");
64!
2259
            if (native && PyObject_TypeCheck(native, &RuntimeType)) {
64!
2260
                self->runtime_obj = native;  // already incref'd by GetAttr
64✔
2261
            } else {
32✔
2262
                Py_XDECREF(native);
×
2263
                PyErr_SetString(PyExc_TypeError,
×
2264
                                "runtime must be a Runtime instance or None");
2265
                return -1;
×
2266
            }
2267
        }
2268
    }
32✔
2269

2270
    self->file_path = PyUnicode_FromString(file_path);
316!
2271
    if (!self->file_path) return -1;
316✔
2272

2273
    self->index_dir = PyUnicode_FromString(index_dir);
316!
2274
    if (!self->index_dir) {
316✔
2275
        Py_DECREF(self->file_path);
×
2276
        self->file_path = NULL;
×
2277
        return -1;
×
2278
    }
2279

2280
    self->checkpoint_size = checkpoint_size;
316✔
2281
    self->auto_build_index = auto_build_index;
316✔
2282

2283
    try {
2284
        TraceReaderConfig cfg;
316✔
2285
        cfg.file_path = file_path;
316!
2286
        cfg.index_dir = index_dir;
316!
2287
        cfg.checkpoint_size = checkpoint_size;
316✔
2288
        cfg.auto_build_index = auto_build_index != 0;
316✔
2289
        TraceReader probe(std::move(cfg));
316!
2290
        self->has_index = probe.has_index() ? 1 : 0;
316!
2291
    } catch (const std::exception &e) {
316!
2292
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2293
        Py_DECREF(self->file_path);
×
2294
        Py_DECREF(self->index_dir);
×
2295
        self->file_path = NULL;
×
2296
        self->index_dir = NULL;
×
2297
        return -1;
×
2298
    }
×
2299

2300
    return 0;
316✔
2301
}
158✔
2302

2303
static PyObject *TraceReader_iter_lines(TraceReaderObject *self, PyObject *args,
124✔
2304
                                        PyObject *kwds) {
2305
    static const char *kwlist[] = {"start_line",    "end_line",    "start_byte",
2306
                                   "end_byte",      "buffer_size", "query",
2307
                                   "memory_budget", NULL};
2308
    Py_ssize_t start_line = 0, end_line = 0;
124✔
2309
    Py_ssize_t start_byte = 0, end_byte = 0;
124✔
2310
    Py_ssize_t buffer_size = 4 * 1024 * 1024;
124✔
2311
    const char *query_str = NULL;
124✔
2312
    Py_ssize_t memory_budget = 0;
124✔
2313

2314
    if (!PyArg_ParseTupleAndKeywords(
124!
2315
            args, kwds, "|nnnnnzn", (char **)kwlist, &start_line, &end_line,
62✔
2316
            &start_byte, &end_byte, &buffer_size, &query_str, &memory_budget)) {
2317
        return NULL;
×
2318
    }
2319

2320
    if (start_line < 0 || end_line < 0 || start_byte < 0 || end_byte < 0 ||
124!
2321
        buffer_size <= 0) {
118!
2322
        PyErr_SetString(
6!
2323
            PyExc_ValueError,
3✔
2324
            "range arguments must be >= 0; buffer_size must be > 0");
2325
        return NULL;
6✔
2326
    }
2327

2328
    TraceReaderConfig cfg;
118✔
2329
    try {
2330
        cfg = build_config(self);
118!
2331
    } catch (const std::exception &e) {
59!
2332
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2333
        return NULL;
×
2334
    }
×
2335

2336
    ReadConfig rc;
118✔
2337
    rc.start_line = static_cast<std::size_t>(start_line);
118✔
2338
    rc.end_line = static_cast<std::size_t>(end_line);
118✔
2339
    rc.start_byte = static_cast<std::size_t>(start_byte);
118✔
2340
    rc.end_byte = static_cast<std::size_t>(end_byte);
118✔
2341
    rc.buffer_size = static_cast<std::size_t>(buffer_size);
118✔
2342
    if (query_str) rc.query = query_str;
118!
2343

2344
    auto state = std::make_shared<MemoryViewBatchIteratorState>();
118!
2345
    state->memory_budget_bytes = dftracer::utils::compute_memory_budget(
118!
2346
        static_cast<std::size_t>(memory_budget));
59✔
2347

2348
    Runtime *rt = get_runtime(self);
118!
2349
    std::size_t max_workers = rt->threads();
118!
2350
    constexpr std::size_t LINE_BATCH_SIZE = 1024;
118✔
2351
    std::size_t capacity = dftracer::utils::compute_channel_capacity(
118!
2352
        state->memory_budget_bytes, LINE_BATCH_SIZE * ESTIMATED_BYTES_PER_LINE,
118✔
2353
        max_workers);
59✔
2354
    state->channel =
118✔
2355
        dftracer::utils::coro::make_channel<MemoryViewBatchData>(capacity);
177!
2356
    auto *sp = state.get();
118✔
2357

2358
    try {
2359
        bool is_dir = fs::is_directory(cfg.file_path);
118!
2360
        if (is_dir) {
118✔
2361
            auto handle = rt->scope(
6!
2362
                "iter_lines_parallel",
3!
2363
                [sp, dir_path = cfg.file_path, index_dir = cfg.index_dir,
27!
2364
                 checkpoint_size = cfg.checkpoint_size,
6✔
2365
                 auto_build_index = cfg.auto_build_index, rc,
9!
2366
                 max_workers](CoroScope &scope) -> CoroTask<void> {
6!
2367
                    co_await produce_lines_parallel(
24!
2368
                        scope, sp, dir_path, index_dir, checkpoint_size,
9!
2369
                        auto_build_index, rc, LINE_BATCH_SIZE, max_workers);
9!
2370
                });
18!
2371
            state->task_future = handle.future;
6!
2372
        } else {
6✔
2373
            auto handle = rt->submit(
112!
2374
                produce_lines_batched(state, state->channel->producer(), cfg,
224!
2375
                                      rc, LINE_BATCH_SIZE),
56!
2376
                "iter_lines");
224!
2377
            state->task_future = handle.future;
112!
2378
        }
112✔
2379
    } catch (const std::exception &e) {
59!
2380
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2381
        return NULL;
×
2382
    }
×
2383

2384
    TraceReaderIteratorObject *it = make_memoryview_iterator(std::move(state));
118!
2385
    return (PyObject *)it;
118✔
2386
}
121✔
2387

2388
static PyObject *TraceReader_iter_raw(TraceReaderObject *self, PyObject *args,
26✔
2389
                                      PyObject *kwds) {
2390
    static const char *kwlist[] = {"start_line", "end_line",    "start_byte",
2391
                                   "end_byte",   "buffer_size", "line_aligned",
2392
                                   "multi_line", "query",       "memory_budget",
2393
                                   NULL};
2394
    Py_ssize_t start_line = 0, end_line = 0;
26✔
2395
    Py_ssize_t start_byte = 0, end_byte = 0;
26✔
2396
    Py_ssize_t buffer_size = 4 * 1024 * 1024;
26✔
2397
    int line_aligned = 1;
26✔
2398
    int multi_line = 1;
26✔
2399
    const char *query_str = NULL;
26✔
2400
    Py_ssize_t memory_budget = 0;
26✔
2401

2402
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|nnnnnppzn", (char **)kwlist,
26!
2403
                                     &start_line, &end_line, &start_byte,
2404
                                     &end_byte, &buffer_size, &line_aligned,
2405
                                     &multi_line, &query_str, &memory_budget)) {
2406
        return NULL;
×
2407
    }
2408

2409
    if (start_line < 0 || end_line < 0 || start_byte < 0 || end_byte < 0 ||
26!
2410
        buffer_size <= 0) {
24!
2411
        PyErr_SetString(
2!
2412
            PyExc_ValueError,
1✔
2413
            "range arguments must be >= 0; buffer_size must be > 0");
2414
        return NULL;
2✔
2415
    }
2416

2417
    TraceReaderConfig cfg;
24✔
2418
    try {
2419
        cfg = build_config(self);
24!
2420
    } catch (const std::exception &e) {
12!
2421
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2422
        return NULL;
×
2423
    }
×
2424

2425
    ReadConfig rc;
24✔
2426
    rc.start_line = static_cast<std::size_t>(start_line);
24✔
2427
    rc.end_line = static_cast<std::size_t>(end_line);
24✔
2428
    rc.start_byte = static_cast<std::size_t>(start_byte);
24✔
2429
    rc.end_byte = static_cast<std::size_t>(end_byte);
24✔
2430
    rc.buffer_size = static_cast<std::size_t>(buffer_size);
24✔
2431
    rc.line_aligned = line_aligned != 0;
24✔
2432
    rc.multi_line = multi_line != 0;
24✔
2433
    if (query_str) rc.query = query_str;
24!
2434

2435
    auto state = std::make_shared<MemoryViewBatchIteratorState>();
24!
2436
    state->memory_budget_bytes = dftracer::utils::compute_memory_budget(
24!
2437
        static_cast<std::size_t>(memory_budget));
12✔
2438

2439
    Runtime *rt = get_runtime(self);
24!
2440
    std::size_t max_workers = rt->threads();
24!
2441
    std::size_t capacity = dftracer::utils::compute_channel_capacity(
24!
2442
        state->memory_budget_bytes, ESTIMATED_BYTES_PER_RAW_CHUNK, max_workers);
24✔
2443
    state->channel =
24✔
2444
        dftracer::utils::coro::make_channel<MemoryViewBatchData>(capacity);
36!
2445
    auto *sp = state.get();
24✔
2446

2447
    try {
2448
        bool is_dir = fs::is_directory(cfg.file_path);
24!
2449
        if (is_dir) {
24✔
2450
            auto handle = rt->scope(
4!
2451
                "iter_raw_parallel",
2!
2452
                [sp, dir_path = cfg.file_path, index_dir = cfg.index_dir,
18!
2453
                 checkpoint_size = cfg.checkpoint_size,
4✔
2454
                 auto_build_index = cfg.auto_build_index, rc,
6!
2455
                 max_workers](CoroScope &scope) -> CoroTask<void> {
4!
2456
                    co_await produce_raw_parallel(
16!
2457
                        scope, sp, dir_path, index_dir, checkpoint_size,
6!
2458
                        auto_build_index, rc, max_workers);
6!
2459
                });
12!
2460
            state->task_future = handle.future;
4!
2461
        } else {
4✔
2462
            auto handle = rt->submit(
20!
2463
                produce_raw_batched(state, state->channel->producer(), cfg, rc),
30!
2464
                "iter_raw");
40!
2465
            state->task_future = handle.future;
20!
2466
        }
20✔
2467
    } catch (const std::exception &e) {
12!
2468
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2469
        return NULL;
×
2470
    }
×
2471

2472
    TraceReaderIteratorObject *it = make_memoryview_iterator(std::move(state));
24!
2473
    return (PyObject *)it;
24✔
2474
}
25✔
2475

2476
static PyObject *TraceReader_read_lines(TraceReaderObject *self, PyObject *args,
92✔
2477
                                        PyObject *kwds) {
2478
    PyObject *iter = TraceReader_iter_lines(self, args, kwds);
92✔
2479
    if (!iter) return NULL;
92✔
2480
    PyObject *list = PySequence_List(iter);
88✔
2481
    Py_DECREF(iter);
44✔
2482
    return list;
88✔
2483
}
46✔
2484

2485
static PyObject *TraceReader_iter_json(TraceReaderObject *self, PyObject *args,
14✔
2486
                                       PyObject *kwds) {
2487
    static const char *kwlist[] = {"start_line", "end_line",      "start_byte",
2488
                                   "end_byte",   "buffer_size",   "query",
2489
                                   "batch_size", "memory_budget", NULL};
2490
    Py_ssize_t start_line = 0, end_line = 0;
14✔
2491
    Py_ssize_t start_byte = 0, end_byte = 0;
14✔
2492
    Py_ssize_t buffer_size = 4 * 1024 * 1024;
14✔
2493
    const char *query_str = NULL;
14✔
2494
    Py_ssize_t batch_size = 1024;
14✔
2495
    Py_ssize_t memory_budget = 0;
14✔
2496

2497
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|nnnnnznn", (char **)kwlist,
14!
2498
                                     &start_line, &end_line, &start_byte,
2499
                                     &end_byte, &buffer_size, &query_str,
2500
                                     &batch_size, &memory_budget)) {
2501
        return NULL;
×
2502
    }
2503

2504
    if (start_line < 0 || end_line < 0 || start_byte < 0 || end_byte < 0 ||
14!
2505
        buffer_size <= 0 || batch_size <= 0) {
14!
2506
        PyErr_SetString(PyExc_ValueError,
×
2507
                        "range arguments must be >= 0; buffer_size and "
2508
                        "batch_size must be > 0");
2509
        return NULL;
×
2510
    }
2511

2512
    TraceReaderConfig cfg;
14✔
2513
    try {
2514
        cfg = build_config(self);
14!
2515
    } catch (const std::exception &e) {
7!
2516
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2517
        return NULL;
×
2518
    }
×
2519

2520
    ReadConfig rc;
14✔
2521
    rc.start_line = static_cast<std::size_t>(start_line);
14✔
2522
    rc.end_line = static_cast<std::size_t>(end_line);
14✔
2523
    rc.start_byte = static_cast<std::size_t>(start_byte);
14✔
2524
    rc.end_byte = static_cast<std::size_t>(end_byte);
14✔
2525
    rc.buffer_size = static_cast<std::size_t>(buffer_size);
14✔
2526
    if (query_str) rc.query = query_str;
14!
2527

2528
    auto state = std::make_shared<JsonDictIteratorState>();
14!
2529
    state->memory_budget_bytes = dftracer::utils::compute_memory_budget(
14!
2530
        static_cast<std::size_t>(memory_budget));
7✔
2531

2532
    Runtime *rt = get_runtime(self);
14!
2533
    std::size_t max_workers = rt->threads();
14!
2534
    auto bs = static_cast<std::size_t>(batch_size);
14✔
2535
    std::size_t capacity = dftracer::utils::compute_channel_capacity(
14!
2536
        state->memory_budget_bytes, bs * ESTIMATED_BYTES_PER_JSON_EVENT,
14✔
2537
        max_workers);
7✔
2538
    state->channel =
14✔
2539
        dftracer::utils::coro::make_channel<JsonDictBatch>(capacity);
21!
2540
    auto *sp = state.get();
14✔
2541

2542
    try {
2543
        bool is_dir = fs::is_directory(cfg.file_path);
14!
2544
        if (is_dir) {
14✔
2545
            auto handle = rt->scope(
12!
2546
                "iter_json_parallel",
6!
2547
                [sp, dir_path = cfg.file_path, index_dir = cfg.index_dir,
54!
2548
                 checkpoint_size = cfg.checkpoint_size,
12✔
2549
                 auto_build_index = cfg.auto_build_index, rc, bs,
18!
2550
                 max_workers](CoroScope &scope) -> CoroTask<void> {
12!
2551
                    co_await produce_json_dicts_parallel(
48!
2552
                        scope, sp, dir_path, index_dir, checkpoint_size,
18!
2553
                        auto_build_index, rc, bs, max_workers);
18!
2554
                });
36!
2555
            state->task_future = handle.future;
12!
2556
        } else {
12✔
2557
            auto handle =
2558
                rt->submit(produce_json_dicts(state, state->channel->producer(),
5!
2559
                                              cfg, rc, bs),
1!
2560
                           "iter_json");
4!
2561
            state->task_future = handle.future;
2!
2562
        }
2✔
2563
    } catch (const std::exception &e) {
7!
2564
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2565
        return NULL;
×
2566
    }
×
2567

2568
    TraceReaderIteratorObject *it = make_json_dict_iterator(std::move(state));
14!
2569
    return (PyObject *)it;
14✔
2570
}
14✔
2571

2572
static PyObject *TraceReader_read_json_py(TraceReaderObject *self,
2✔
2573
                                          PyObject *args, PyObject *kwds) {
2574
    PyObject *iter = TraceReader_iter_json(self, args, kwds);
2✔
2575
    if (!iter) return NULL;
2✔
2576
    PyObject *list = PySequence_List(iter);
2✔
2577
    Py_DECREF(iter);
1✔
2578
    return list;
2✔
2579
}
1✔
2580

2581
static PyObject *TraceReader_read_raw(TraceReaderObject *self, PyObject *args,
8✔
2582
                                      PyObject *kwds) {
2583
    PyObject *iter = TraceReader_iter_raw(self, args, kwds);
8✔
2584
    if (!iter) return NULL;
8✔
2585
    PyObject *list = PySequence_List(iter);
8✔
2586
    Py_DECREF(iter);
4✔
2587
    return list;
8✔
2588
}
4✔
2589

2590
#ifdef DFTRACER_UTILS_ENABLE_ARROW
2591

2592
static PyObject *TraceReader_iter_arrow(TraceReaderObject *self, PyObject *args,
54✔
2593
                                        PyObject *kwds) {
2594
    static const char *kwlist[] = {
2595
        "batch_size", "start_line",    "end_line", "start_byte",
2596
        "end_byte",   "buffer_size",   "query",    "flatten_objects",
2597
        "normalize",  "memory_budget", NULL};
2598
    Py_ssize_t batch_size = 10000;
54✔
2599
    Py_ssize_t start_line = 0, end_line = 0;
54✔
2600
    Py_ssize_t start_byte = 0, end_byte = 0;
54✔
2601
    Py_ssize_t buffer_size = 4 * 1024 * 1024;
54✔
2602
    const char *query_str = NULL;
54✔
2603
    int flatten_objects = 1;  // default: expand top-level objects
54✔
2604
    int normalize = 0;
54✔
2605
    Py_ssize_t memory_budget = 0;
54✔
2606

2607
    if (!PyArg_ParseTupleAndKeywords(
54!
2608
            args, kwds, "|nnnnnnzppn", (char **)kwlist, &batch_size,
27✔
2609
            &start_line, &end_line, &start_byte, &end_byte, &buffer_size,
2610
            &query_str, &flatten_objects, &normalize, &memory_budget)) {
2611
        return NULL;
×
2612
    }
2613

2614
    if (batch_size <= 0) {
54!
2615
        PyErr_SetString(PyExc_ValueError, "batch_size must be > 0");
×
2616
        return NULL;
×
2617
    }
2618
    if (start_line < 0 || end_line < 0 || start_byte < 0 || end_byte < 0 ||
54!
2619
        buffer_size <= 0) {
54!
2620
        PyErr_SetString(
×
2621
            PyExc_ValueError,
2622
            "range arguments must be >= 0; buffer_size must be > 0");
2623
        return NULL;
×
2624
    }
2625

2626
    TraceReaderConfig cfg;
54✔
2627
    try {
2628
        cfg = build_config(self);
54!
2629
    } catch (const std::exception &e) {
27!
2630
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2631
        return NULL;
×
2632
    }
×
2633

2634
    ReadConfig rc;
54✔
2635
    rc.start_line = static_cast<std::size_t>(start_line);
54✔
2636
    rc.end_line = static_cast<std::size_t>(end_line);
54✔
2637
    rc.start_byte = static_cast<std::size_t>(start_byte);
54✔
2638
    rc.end_byte = static_cast<std::size_t>(end_byte);
54✔
2639
    rc.buffer_size = static_cast<std::size_t>(buffer_size);
54✔
2640
    rc.flatten_objects = flatten_objects != 0;
54✔
2641
    if (query_str) rc.query = query_str;
54!
2642

2643
    auto state = std::make_shared<ArrowIteratorState>();
54!
2644
    state->memory_budget_bytes = dftracer::utils::compute_memory_budget(
54!
2645
        static_cast<std::size_t>(memory_budget));
27✔
2646

2647
    Runtime *rt = get_runtime(self);
54!
2648
    std::size_t max_workers = rt->threads();
54!
2649
    auto bs = static_cast<std::size_t>(batch_size);
54✔
2650
    std::size_t capacity = dftracer::utils::compute_channel_capacity(
54!
2651
        state->memory_budget_bytes, bs * ESTIMATED_BYTES_PER_ARROW_ROW,
54✔
2652
        max_workers);
27✔
2653
    state->channel =
54✔
2654
        dftracer::utils::coro::make_channel<ArrowIteratorState::BatchType>(
81!
2655
            capacity);
54✔
2656
    auto *sp = state.get();
54✔
2657

2658
    try {
2659
        bool is_dir = fs::is_directory(cfg.file_path);
54!
2660
        if (is_dir) {
54✔
2661
            auto handle = rt->scope(
10!
2662
                "iter_arrow_parallel",
5!
2663
                [sp, dir_path = cfg.file_path, index_dir = cfg.index_dir,
45!
2664
                 checkpoint_size = cfg.checkpoint_size,
10✔
2665
                 auto_build_index = cfg.auto_build_index, rc, bs,
15!
2666
                 norm = normalize != 0,
10✔
2667
                 max_workers](CoroScope &scope) -> CoroTask<void> {
10!
2668
                    co_await produce_arrow_batches_parallel(
40!
2669
                        scope, sp, dir_path, index_dir, checkpoint_size,
15!
2670
                        auto_build_index, rc, bs, norm, max_workers);
15!
2671
                });
30!
2672
            state->task_future = handle.future;
10!
2673
        } else if (normalize) {
54!
2674
            auto handle = rt->submit(
×
2675
                produce_arrow_batches(state, state->channel->producer(), cfg,
×
2676
                                      rc, static_cast<std::size_t>(batch_size),
×
2677
                                      flatten_objects != 0, normalize != 0),
2678
                "iter_arrow");
×
2679
            state->task_future = handle.future;
×
2680
        } else {
×
2681
            std::vector<std::string> files_vec{cfg.file_path};
88!
2682
            auto handle = rt->scope(
44!
2683
                "iter_arrow_parallel",
22!
2684
                [sp, files = std::move(files_vec), index_dir = cfg.index_dir,
198!
2685
                 checkpoint_size = cfg.checkpoint_size,
44✔
2686
                 auto_build_index = cfg.auto_build_index, rc, bs,
66!
2687
                 norm = normalize != 0,
44✔
2688
                 max_workers](CoroScope &scope) mutable -> CoroTask<void> {
44!
2689
                    co_await produce_arrow_batches_for_files(
176!
2690
                        scope, sp, std::move(files), index_dir, checkpoint_size,
66!
2691
                        auto_build_index, rc, bs, norm, max_workers);
66!
2692
                });
132!
2693
            state->task_future = handle.future;
44!
2694
        }
44✔
2695
    } catch (const std::exception &e) {
27!
2696
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2697
        return NULL;
×
2698
    }
×
2699

2700
    TraceReaderIteratorObject *it = make_arrow_iterator(std::move(state));
54!
2701
    return (PyObject *)it;
54✔
2702
}
54✔
2703

2704
// Build ArrowIteratorState + spawn the producer task. Same plumbing as
2705
// TraceReader_iter_arrow but returns the state so callers can wrap it as
2706
// either a per-batch iterator or an ArrowArrayStream.
2707
static std::shared_ptr<ArrowIteratorState> spawn_arrow_producer(
54✔
2708
    TraceReaderObject *self, PyObject *args, PyObject *kwds) {
2709
    static const char *kwlist[] = {
2710
        "batch_size", "start_line",    "end_line", "start_byte",
2711
        "end_byte",   "buffer_size",   "query",    "flatten_objects",
2712
        "normalize",  "memory_budget", NULL};
2713
    Py_ssize_t batch_size = 10000;
54✔
2714
    Py_ssize_t start_line = 0, end_line = 0;
54✔
2715
    Py_ssize_t start_byte = 0, end_byte = 0;
54✔
2716
    Py_ssize_t buffer_size = 4 * 1024 * 1024;
54✔
2717
    const char *query_str = NULL;
54✔
2718
    int flatten_objects = 1;  // default: expand top-level objects
54✔
2719
    int normalize = 0;
54✔
2720
    Py_ssize_t memory_budget = 0;
54✔
2721

2722
    if (!PyArg_ParseTupleAndKeywords(
54!
2723
            args, kwds, "|nnnnnnzppn", (char **)kwlist, &batch_size,
27✔
2724
            &start_line, &end_line, &start_byte, &end_byte, &buffer_size,
2725
            &query_str, &flatten_objects, &normalize, &memory_budget)) {
2726
        return nullptr;
×
2727
    }
2728

2729
    if (batch_size <= 0) {
54!
2730
        PyErr_SetString(PyExc_ValueError, "batch_size must be > 0");
×
2731
        return nullptr;
×
2732
    }
2733
    if (start_line < 0 || end_line < 0 || start_byte < 0 || end_byte < 0 ||
54!
2734
        buffer_size <= 0) {
54!
2735
        PyErr_SetString(
×
2736
            PyExc_ValueError,
2737
            "range arguments must be >= 0; buffer_size must be > 0");
2738
        return nullptr;
×
2739
    }
2740

2741
    TraceReaderConfig cfg;
54✔
2742
    try {
2743
        cfg = build_config(self);
54!
2744
    } catch (const std::exception &e) {
27!
2745
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2746
        return nullptr;
×
2747
    }
×
2748

2749
    ReadConfig rc;
54✔
2750
    rc.start_line = static_cast<std::size_t>(start_line);
54✔
2751
    rc.end_line = static_cast<std::size_t>(end_line);
54✔
2752
    rc.start_byte = static_cast<std::size_t>(start_byte);
54✔
2753
    rc.end_byte = static_cast<std::size_t>(end_byte);
54✔
2754
    rc.buffer_size = static_cast<std::size_t>(buffer_size);
54✔
2755
    rc.flatten_objects = flatten_objects != 0;
54✔
2756
    if (query_str) rc.query = query_str;
54!
2757

2758
    auto state = std::make_shared<ArrowIteratorState>();
54!
2759
    state->memory_budget_bytes = dftracer::utils::compute_memory_budget(
54!
2760
        static_cast<std::size_t>(memory_budget));
27✔
2761

2762
    Runtime *rt = get_runtime(self);
54!
2763
    std::size_t max_workers = rt->threads();
54!
2764
    auto bs = static_cast<std::size_t>(batch_size);
54✔
2765
    std::size_t capacity = dftracer::utils::compute_channel_capacity(
54!
2766
        state->memory_budget_bytes, bs * ESTIMATED_BYTES_PER_ARROW_ROW,
54✔
2767
        max_workers);
27✔
2768
    state->channel =
54✔
2769
        dftracer::utils::coro::make_channel<ArrowIteratorState::BatchType>(
81!
2770
            capacity);
54✔
2771
    auto *sp = state.get();
54✔
2772

2773
    try {
2774
        bool is_dir = fs::is_directory(cfg.file_path);
54!
2775
        if (is_dir) {
54✔
2776
            auto handle = rt->scope(
20!
2777
                "iter_arrow_parallel",
10!
2778
                [sp, dir_path = cfg.file_path, index_dir = cfg.index_dir,
90!
2779
                 checkpoint_size = cfg.checkpoint_size,
20✔
2780
                 auto_build_index = cfg.auto_build_index, rc, bs,
30!
2781
                 norm = normalize != 0,
20✔
2782
                 max_workers](CoroScope &scope) -> CoroTask<void> {
20!
2783
                    co_await produce_arrow_batches_parallel(
80!
2784
                        scope, sp, dir_path, index_dir, checkpoint_size,
30!
2785
                        auto_build_index, rc, bs, norm, max_workers);
30!
2786
                });
60!
2787
            state->task_future = handle.future;
20!
2788
        } else {
20✔
2789
            auto handle = rt->submit(
34!
2790
                produce_arrow_batches(state, state->channel->producer(), cfg,
68!
2791
                                      rc, static_cast<std::size_t>(batch_size),
17!
2792
                                      flatten_objects != 0, normalize != 0),
17✔
2793
                "iter_arrow");
68!
2794
            state->task_future = handle.future;
34!
2795
        }
34✔
2796
    } catch (const std::exception &e) {
27!
2797
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
2798
        return nullptr;
×
2799
    }
×
2800

2801
    return state;
54✔
2802
}
54✔
2803

2804
static PyObject *TraceReader_iter_arrow_stream(TraceReaderObject *self,
28✔
2805
                                               PyObject *args, PyObject *kwds) {
2806
    auto state = spawn_arrow_producer(self, args, kwds);
28!
2807
    if (!state) return NULL;
28!
2808
    return make_arrow_batch_stream(std::move(state));
28!
2809
}
28✔
2810

2811
static PyObject *TraceReader_read_arrow(TraceReaderObject *self, PyObject *args,
26✔
2812
                                        PyObject *kwds) {
2813
    auto state = spawn_arrow_producer(self, args, kwds);
26!
2814
    if (!state) return NULL;
26!
2815
    PyObject *stream = make_arrow_batch_stream(std::move(state));
26!
2816
    if (!stream) return NULL;
26✔
2817
    return dftracer::utils::python::wrap_arrow_stream_table(stream);
26!
2818
}
26✔
2819

2820
#endif  // DFTRACER_UTILS_ENABLE_ARROW
2821

2822
#ifdef DFTRACER_UTILS_ENABLE_ARROW_IPC
2823

2824
static int parse_str_list_trace(PyObject *obj, std::vector<std::string> &out,
2825
                                const char *param_name) {
2826
    if (!obj || obj == Py_None) return 0;
×
2827
    if (!PyList_Check(obj)) {
×
2828
        PyErr_Format(PyExc_TypeError, "%s must be a list of str", param_name);
2829
        return -1;
2830
    }
2831
    Py_ssize_t n = PyList_Size(obj);
2832
    for (Py_ssize_t i = 0; i < n; i++) {
×
2833
        const char *s = PyUnicode_AsUTF8(PyList_GetItem(obj, i));
×
2834
        if (!s) return -1;
×
2835
        out.emplace_back(s);
×
2836
    }
2837
    return 0;
2838
}
2839

2840
static PyObject *TraceReader_write_arrow(TraceReaderObject *self,
26✔
2841
                                         PyObject *args, PyObject *kwds) {
2842
    static const char *kwlist[] = {"path",        "views",      "chunk_size_mb",
2843
                                   "compression", "batch_size", NULL};
2844
    const char *path = NULL;
26✔
2845
    PyObject *views_obj = Py_None;
26✔
2846
    int chunk_size_mb = 32;
26✔
2847
    const char *compression_str = "zstd";
26✔
2848
    Py_ssize_t batch_size = 10000;
26✔
2849

2850
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "s|Oisn", (char **)kwlist,
26!
2851
                                     &path, &views_obj, &chunk_size_mb,
2852
                                     &compression_str, &batch_size)) {
2853
        return NULL;
×
2854
    }
2855

2856
    if (chunk_size_mb < 0) {
26✔
2857
        PyErr_SetString(PyExc_ValueError, "chunk_size_mb must be >= 0");
×
2858
        return NULL;
×
2859
    }
2860

2861
    std::vector<ViewDefinition> views;
26✔
2862
    if (views_obj && views_obj != Py_None) {
26!
2863
        if (!PyList_Check(views_obj)) {
10!
2864
            PyErr_SetString(PyExc_TypeError, "views must be a list or None");
×
2865
            return NULL;
×
2866
        }
2867
        Py_ssize_t n = PyList_Size(views_obj);
10!
2868
        for (Py_ssize_t i = 0; i < n; i++) {
22✔
2869
            PyObject *item = PyList_GetItem(views_obj, i);
12!
2870
            ViewDefinition vd;
12✔
2871

2872
            if (PyUnicode_Check(item)) {
12✔
2873
                const char *name = PyUnicode_AsUTF8(item);
2!
2874
                if (!name) return NULL;
2✔
2875
                std::string name_str(name);
2!
2876
                if (name_str == "io") {
2!
2877
                    vd = ViewDefinition::io_view();
2!
2878
                } else if (name_str == "compute") {
1!
2879
                    vd = ViewDefinition::compute_view();
×
2880
                } else if (name_str == "dlio") {
×
2881
                    vd = ViewDefinition::dlio_view();
×
2882
                } else {
2883
                    vd.with_name(name_str);
×
2884
                }
2885
            } else if (PyDict_Check(item)) {
12!
2886
                PyObject *name_obj = PyDict_GetItemString(item, "name");
10!
2887
                if (!name_obj || !PyUnicode_Check(name_obj)) {
10!
2888
                    PyErr_SetString(PyExc_ValueError,
×
2889
                                    "view dict must have 'name' string");
2890
                    return NULL;
×
2891
                }
2892
                vd.with_name(PyUnicode_AsUTF8(name_obj));
10!
2893

2894
                PyObject *query_obj = PyDict_GetItemString(item, "query");
10!
2895
                if (query_obj && query_obj != Py_None) {
10!
2896
                    if (!PyUnicode_Check(query_obj)) {
10!
2897
                        PyErr_SetString(PyExc_ValueError,
×
2898
                                        "view 'query' must be a string");
2899
                        return NULL;
×
2900
                    }
2901
                    vd.with_query(PyUnicode_AsUTF8(query_obj));
10!
2902
                }
5✔
2903

2904
                PyObject *meta_obj =
5✔
2905
                    PyDict_GetItemString(item, "include_metadata");
10!
2906
                if (meta_obj && meta_obj != Py_None) {
10!
2907
                    vd.with_include_metadata(PyObject_IsTrue(meta_obj));
2!
2908
                }
1✔
2909
            } else {
5✔
2910
                PyErr_SetString(PyExc_TypeError,
×
2911
                                "views list must contain strings or dicts");
2912
                return NULL;
×
2913
            }
2914
            views.push_back(std::move(vd));
12!
2915
        }
12✔
2916
    }
5✔
2917

2918
    IpcCompression compression = IpcCompression::ZSTD;
26✔
2919
    if (compression_str) {
26!
2920
        std::string comp_lower(compression_str);
26!
2921
        for (auto &c : comp_lower) c = std::tolower(c);
130!
2922
        if (comp_lower == "none") {
26✔
2923
            compression = IpcCompression::NONE;
2✔
2924
        } else if (comp_lower == "zstd") {
25✔
2925
#ifdef DFTRACER_UTILS_ENABLE_ZSTD
2926
            compression = IpcCompression::ZSTD;
24✔
2927
#else
2928
            PyErr_SetString(
2929
                PyExc_ValueError,
2930
                "ZSTD compression not available (built without ZSTD)");
2931
            return NULL;
2932
#endif
2933
        } else {
12✔
2934
            PyErr_Format(PyExc_ValueError,
×
2935
                         "Unknown compression: %s (use 'none' or 'zstd')",
2936
                         compression_str);
2937
            return NULL;
×
2938
        }
2939
    }
26✔
2940

2941
    int64_t chunk_size_bytes =
26✔
2942
        static_cast<int64_t>(chunk_size_mb) * 1024 * 1024;
26✔
2943

2944
    std::string file_path = PyUnicode_AsUTF8(self->file_path);
26!
2945
    std::string index_path;
26✔
2946
    const char *idx = PyUnicode_AsUTF8(self->index_dir);
26!
2947
    if (idx && idx[0] != '\0') {
26!
2948
        index_path = idx;
×
2949
    }
2950
    std::size_t checkpoint_size = self->checkpoint_size;
26✔
2951

2952
    std::string output_path(path);
26!
2953
    WriteArrowResult result;
26✔
2954
    std::string error_msg;
26✔
2955

2956
    Py_BEGIN_ALLOW_THREADS try {
26!
2957
        Runtime *rt = get_runtime(self);
26!
2958
        result =
13✔
2959
            rt->submit(write_arrow_pipeline(
78!
2960
                           file_path, index_path, checkpoint_size,
13!
2961
                           std::move(views), output_path, chunk_size_bytes,
26!
2962
                           compression, static_cast<std::size_t>(batch_size)),
13✔
2963
                       "write_arrow")
13!
2964
                .get();
26!
2965
    } catch (const std::exception &e) {
13!
2966
        error_msg = e.what();
×
2967
    }
×
2968
    Py_END_ALLOW_THREADS
26!
2969

2970
        if (!error_msg.empty()) {
26!
2971
        PyErr_SetString(PyExc_RuntimeError, error_msg.c_str());
×
2972
        return NULL;
×
2973
    }
2974

2975
    if (!result.error.empty()) {
26!
2976
        PyErr_SetString(PyExc_RuntimeError, result.error.c_str());
×
2977
        return NULL;
×
2978
    }
2979

2980
    // Build result dict
2981
    PyObject *dict = PyDict_New();
26!
2982
    if (!dict) return NULL;
26✔
2983

2984
    // Build files list per partition
2985
    PyObject *partitions_dict = PyDict_New();
26!
2986
    if (!partitions_dict) {
26!
2987
        Py_DECREF(dict);
×
2988
        return NULL;
×
2989
    }
2990

2991
    for (const auto &[partition_name, partition_stats] :
54!
2992
         result.stats.partitions) {
53✔
2993
        PyObject *partition_dict = PyDict_New();
28!
2994
        if (!partition_dict) {
28!
2995
            Py_DECREF(partitions_dict);
×
2996
            Py_DECREF(dict);
×
2997
            return NULL;
×
2998
        }
2999

3000
        PyObject *files_list = PyList_New(0);
28!
3001
        if (!files_list) {
28!
3002
            Py_DECREF(partition_dict);
×
3003
            Py_DECREF(partitions_dict);
×
3004
            Py_DECREF(dict);
×
3005
            return NULL;
×
3006
        }
3007

3008
        for (const auto &f : partition_stats.files) {
44✔
3009
            PyObject *file_str = PyUnicode_FromString(f.c_str());
16!
3010
            if (!file_str || PyList_Append(files_list, file_str) < 0) {
16!
3011
                Py_XDECREF(file_str);
×
3012
                Py_DECREF(files_list);
×
3013
                Py_DECREF(partition_dict);
×
3014
                Py_DECREF(partitions_dict);
×
3015
                Py_DECREF(dict);
×
3016
                return NULL;
×
3017
            }
3018
            Py_DECREF(file_str);
8!
3019
        }
3020

3021
        PyDict_SetItemString(partition_dict, "files", files_list);
28!
3022
        PyDict_SetItemString(partition_dict, "rows",
28!
3023
                             PyLong_FromLongLong(partition_stats.total_rows));
28!
3024
        PyDict_SetItemString(
28!
3025
            partition_dict, "bytes",
14✔
3026
            PyLong_FromLongLong(partition_stats.total_uncompressed_bytes));
28!
3027
        Py_DECREF(files_list);
14!
3028

3029
        PyObject *key = partition_name.empty()
42!
3030
                            ? PyUnicode_FromString("_default")
14!
3031
                            : PyUnicode_FromString(partition_name.c_str());
28!
3032
        PyDict_SetItem(partitions_dict, key, partition_dict);
28!
3033
        Py_DECREF(key);
14!
3034
        Py_DECREF(partition_dict);
14!
3035
    }
3036

3037
    PyDict_SetItemString(dict, "partitions", partitions_dict);
26!
3038
    PyDict_SetItemString(dict, "total_rows",
26!
3039
                         PyLong_FromLongLong(result.stats.total_rows));
26!
3040
    PyDict_SetItemString(
26!
3041
        dict, "total_bytes",
13✔
3042
        PyLong_FromLongLong(result.stats.total_uncompressed_bytes));
26!
3043
    PyDict_SetItemString(dict, "chunks_scanned",
26!
3044
                         PyLong_FromUnsignedLongLong(result.chunks_scanned));
26!
3045
    PyDict_SetItemString(dict, "chunks_skipped",
26!
3046
                         PyLong_FromUnsignedLongLong(result.chunks_skipped));
26!
3047
    Py_DECREF(partitions_dict);
13!
3048

3049
    return dict;
26✔
3050
}
26✔
3051

3052
static PyObject *TraceReader_get_view_chunks(TraceReaderObject *self,
6✔
3053
                                             PyObject *args, PyObject *kwds) {
3054
    static const char *kwlist[] = {"view", NULL};
3055
    PyObject *view_obj = Py_None;
6✔
3056

3057
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "|O", (char **)kwlist,
6!
3058
                                     &view_obj)) {
3059
        return NULL;
×
3060
    }
3061

3062
    ViewDefinition view;
6✔
3063
    if (view_obj && view_obj != Py_None) {
6!
3064
        if (PyUnicode_Check(view_obj)) {
2!
3065
            const char *name = PyUnicode_AsUTF8(view_obj);
×
3066
            if (!name) return NULL;
×
3067
            std::string name_str(name);
×
3068
            if (name_str == "io") {
×
3069
                view = ViewDefinition::io_view();
×
3070
            } else if (name_str == "compute") {
×
3071
                view = ViewDefinition::compute_view();
×
3072
            } else if (name_str == "dlio") {
×
3073
                view = ViewDefinition::dlio_view();
×
3074
            } else {
3075
                view.with_name(name_str);
×
3076
            }
3077
        } else if (PyDict_Check(view_obj)) {
2!
3078
            PyObject *name_obj = PyDict_GetItemString(view_obj, "name");
2!
3079
            if (name_obj && PyUnicode_Check(name_obj)) {
2!
3080
                view.with_name(PyUnicode_AsUTF8(name_obj));
2!
3081
            }
1✔
3082
            PyObject *query_obj = PyDict_GetItemString(view_obj, "query");
2!
3083
            if (query_obj && query_obj != Py_None &&
3!
3084
                PyUnicode_Check(query_obj)) {
2✔
3085
                view.with_query(PyUnicode_AsUTF8(query_obj));
2!
3086
            }
1✔
3087
        } else {
1✔
3088
            PyErr_SetString(PyExc_TypeError, "view must be a string or dict");
×
3089
            return NULL;
×
3090
        }
3091
    }
1✔
3092

3093
    std::string file_path = PyUnicode_AsUTF8(self->file_path);
6!
3094
    std::string index_path;
6✔
3095
    const char *idx = PyUnicode_AsUTF8(self->index_dir);
6!
3096
    if (idx && idx[0] != '\0') {
6!
3097
        index_path = idx;
×
3098
    }
3099
    std::size_t checkpoint_size = self->checkpoint_size;
6✔
3100

3101
    GetViewChunksResult result;
6✔
3102
    std::string error_msg;
6✔
3103

3104
    Py_BEGIN_ALLOW_THREADS try {
6!
3105
        Runtime *rt = get_runtime(self);
6!
3106
        result = rt->submit(get_view_chunks_pipeline(file_path, index_path,
18!
3107
                                                     checkpoint_size, view),
3!
3108
                            "get_view_chunks")
3!
3109
                     .get();
6!
3110
    } catch (const std::exception &e) {
3!
3111
        error_msg = e.what();
×
3112
    }
×
3113
    Py_END_ALLOW_THREADS
6!
3114

3115
        if (!error_msg.empty()) {
6!
3116
        PyErr_SetString(PyExc_RuntimeError, error_msg.c_str());
×
3117
        return NULL;
×
3118
    }
3119

3120
    if (!result.error.empty()) {
6!
3121
        PyErr_SetString(PyExc_RuntimeError, result.error.c_str());
×
3122
        return NULL;
×
3123
    }
3124

3125
    PyObject *dict = PyDict_New();
6!
3126
    if (!dict) return NULL;
6!
3127

3128
    PyObject *chunks_list = PyList_New(result.chunks.size());
6!
3129
    if (!chunks_list) {
6!
3130
        Py_DECREF(dict);
×
3131
        return NULL;
×
3132
    }
3133

3134
    for (std::size_t i = 0; i < result.chunks.size(); ++i) {
14✔
3135
        const auto &chunk = result.chunks[i];
8✔
3136
        PyObject *chunk_dict = PyDict_New();
8!
3137
        if (!chunk_dict) {
8!
3138
            Py_DECREF(chunks_list);
×
3139
            Py_DECREF(dict);
×
3140
            return NULL;
×
3141
        }
3142
        PyDict_SetItemString(chunk_dict, "checkpoint_idx",
8!
3143
                             PyLong_FromUnsignedLongLong(chunk.checkpoint_idx));
8!
3144
        PyDict_SetItemString(chunk_dict, "start_byte",
8!
3145
                             PyLong_FromSize_t(chunk.start_byte));
8!
3146
        PyDict_SetItemString(chunk_dict, "end_byte",
8!
3147
                             PyLong_FromSize_t(chunk.end_byte));
8!
3148
        PyList_SetItem(chunks_list, i, chunk_dict);
8!
3149
    }
4✔
3150

3151
    PyDict_SetItemString(dict, "chunks", chunks_list);
6!
3152
    PyDict_SetItemString(dict, "total_checkpoints",
6!
3153
                         PyLong_FromUnsignedLongLong(result.total_checkpoints));
6!
3154
    PyDict_SetItemString(
6!
3155
        dict, "skipped_checkpoints",
3✔
3156
        PyLong_FromUnsignedLongLong(result.skipped_checkpoints));
6!
3157
    PyDict_SetItemString(dict, "file_may_match",
6!
3158
                         PyBool_FromLong(result.file_may_match ? 1 : 0));
6✔
3159
    Py_DECREF(chunks_list);
3!
3160

3161
    return dict;
6✔
3162
}
6✔
3163

3164
static PyObject *TraceReader_write_view_chunk(TraceReaderObject *self,
2✔
3165
                                              PyObject *args, PyObject *kwds) {
3166
    static const char *kwlist[] = {
3167
        "output_file", "checkpoint_idx", "start_byte", "end_byte",
3168
        "view",        "compression",    "batch_size", NULL};
3169
    const char *output_file = NULL;
2✔
3170
    unsigned long long checkpoint_idx = 0;
2✔
3171
    Py_ssize_t start_byte = 0;
2✔
3172
    Py_ssize_t end_byte = 0;
2✔
3173
    PyObject *view_obj = Py_None;
2✔
3174
    const char *compression_str = "zstd";
2✔
3175
    Py_ssize_t batch_size = 10000;
2✔
3176

3177
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "sKnn|Osn", (char **)kwlist,
2!
3178
                                     &output_file, &checkpoint_idx, &start_byte,
3179
                                     &end_byte, &view_obj, &compression_str,
3180
                                     &batch_size)) {
3181
        return NULL;
×
3182
    }
3183

3184
    IpcCompression compression = IpcCompression::ZSTD;
2✔
3185
    if (compression_str) {
2✔
3186
        std::string comp_lower(compression_str);
2!
3187
        for (auto &c : comp_lower) c = std::tolower(c);
10!
3188
        if (comp_lower == "none") {
2!
3189
            compression = IpcCompression::NONE;
×
3190
        } else if (comp_lower == "zstd") {
2✔
3191
#ifdef DFTRACER_UTILS_ENABLE_ZSTD
3192
            compression = IpcCompression::ZSTD;
2✔
3193
#else
3194
            PyErr_SetString(PyExc_ValueError, "ZSTD compression not available");
3195
            return NULL;
3196
#endif
3197
        }
1✔
3198
    }
2✔
3199

3200
    ViewDefinition view;
2✔
3201
    if (view_obj && view_obj != Py_None) {
2!
3202
        if (PyUnicode_Check(view_obj)) {
×
3203
            const char *name = PyUnicode_AsUTF8(view_obj);
×
3204
            if (!name) return NULL;
×
3205
            std::string name_str(name);
×
3206
            if (name_str == "io") {
×
3207
                view = ViewDefinition::io_view();
×
3208
            } else if (name_str == "compute") {
×
3209
                view = ViewDefinition::compute_view();
×
3210
            } else if (name_str == "dlio") {
×
3211
                view = ViewDefinition::dlio_view();
×
3212
            } else {
3213
                view.with_name(name_str);
×
3214
            }
3215
        } else if (PyDict_Check(view_obj)) {
×
3216
            PyObject *name_obj = PyDict_GetItemString(view_obj, "name");
×
3217
            if (name_obj && PyUnicode_Check(name_obj)) {
×
3218
                view.with_name(PyUnicode_AsUTF8(name_obj));
×
3219
            }
3220
            PyObject *query_obj = PyDict_GetItemString(view_obj, "query");
×
3221
            if (query_obj && query_obj != Py_None &&
×
3222
                PyUnicode_Check(query_obj)) {
×
3223
                view.with_query(PyUnicode_AsUTF8(query_obj));
×
3224
            }
3225
        }
3226
    }
3227

3228
    std::string file_path = PyUnicode_AsUTF8(self->file_path);
2!
3229
    std::string index_path;
2✔
3230
    const char *idx = PyUnicode_AsUTF8(self->index_dir);
2!
3231
    if (idx && idx[0] != '\0') {
2!
3232
        index_path = idx;
×
3233
    }
3234
    std::size_t checkpoint_size = self->checkpoint_size;
2✔
3235

3236
    WriteViewChunkResult result;
2✔
3237
    std::string error_msg;
2✔
3238

3239
    Py_BEGIN_ALLOW_THREADS try {
2!
3240
        Runtime *rt = get_runtime(self);
2!
3241
        result =
1✔
3242
            rt->submit(write_view_chunk_pipeline(
7!
3243
                           file_path, index_path, checkpoint_size, view,
1!
3244
                           checkpoint_idx, static_cast<std::size_t>(start_byte),
1✔
3245
                           static_cast<std::size_t>(end_byte),
1✔
3246
                           std::string(output_file), compression,
3!
3247
                           static_cast<std::size_t>(batch_size)),
1✔
3248
                       "write_view_chunk")
1!
3249
                .get();
2!
3250
    } catch (const std::exception &e) {
1!
3251
        error_msg = e.what();
×
3252
    }
×
3253
    Py_END_ALLOW_THREADS
2!
3254

3255
        if (!error_msg.empty()) {
2!
3256
        PyErr_SetString(PyExc_RuntimeError, error_msg.c_str());
×
3257
        return NULL;
×
3258
    }
3259

3260
    if (!result.error.empty()) {
2!
3261
        PyErr_SetString(PyExc_RuntimeError, result.error.c_str());
×
3262
        return NULL;
×
3263
    }
3264

3265
    PyObject *dict = PyDict_New();
2!
3266
    if (!dict) return NULL;
2✔
3267

3268
    PyDict_SetItemString(dict, "output_file",
2!
3269
                         PyUnicode_FromString(result.output_file.c_str()));
1!
3270
    PyDict_SetItemString(dict, "events_matched",
2!
3271
                         PyLong_FromUnsignedLongLong(result.events_matched));
2!
3272
    PyDict_SetItemString(dict, "events_scanned",
2!
3273
                         PyLong_FromUnsignedLongLong(result.events_scanned));
2!
3274
    PyDict_SetItemString(dict, "rows_written",
2!
3275
                         PyLong_FromLongLong(result.rows_written));
2!
3276
    PyDict_SetItemString(dict, "bytes_written",
2!
3277
                         PyLong_FromLongLong(result.bytes_written));
2!
3278

3279
    return dict;
2✔
3280
}
2✔
3281

3282
static PyObject *TraceReader_write_view_chunks(TraceReaderObject *self,
2✔
3283
                                               PyObject *args, PyObject *kwds) {
3284
    static const char *kwlist[] = {"chunks",      "output_dir", "view",
3285
                                   "compression", "batch_size", NULL};
3286
    PyObject *chunks_list = NULL;
2✔
3287
    const char *output_dir = NULL;
2✔
3288
    PyObject *view_obj = Py_None;
2✔
3289
    const char *compression_str = "zstd";
2✔
3290
    Py_ssize_t batch_size = 10000;
2✔
3291

3292
    if (!PyArg_ParseTupleAndKeywords(args, kwds, "Os|Osn", (char **)kwlist,
2!
3293
                                     &chunks_list, &output_dir, &view_obj,
3294
                                     &compression_str, &batch_size)) {
3295
        return NULL;
×
3296
    }
3297

3298
    if (!PyList_Check(chunks_list)) {
2✔
3299
        PyErr_SetString(PyExc_TypeError, "chunks must be a list");
×
3300
        return NULL;
×
3301
    }
3302

3303
    IpcCompression compression = IpcCompression::ZSTD;
2✔
3304
    if (strcmp(compression_str, "none") == 0) {
2!
3305
        compression = IpcCompression::NONE;
×
3306
    } else if (strcmp(compression_str, "zstd") != 0) {
2!
3307
        PyErr_SetString(PyExc_ValueError,
×
3308
                        "compression must be 'zstd' or 'none'");
3309
        return NULL;
×
3310
    }
3311

3312
    ViewDefinition view;
2✔
3313
    if (view_obj && view_obj != Py_None) {
2!
3314
        if (PyUnicode_Check(view_obj)) {
×
3315
            const char *name = PyUnicode_AsUTF8(view_obj);
×
3316
            if (!name) return NULL;
×
3317
            std::string name_str(name);
×
3318
            if (name_str == "io") {
×
3319
                view = ViewDefinition::io_view();
×
3320
            } else if (name_str == "compute") {
×
3321
                view = ViewDefinition::compute_view();
×
3322
            } else if (name_str == "dlio") {
×
3323
                view = ViewDefinition::dlio_view();
×
3324
            } else {
3325
                view.with_name(name_str);
×
3326
            }
3327
        } else if (PyDict_Check(view_obj)) {
×
3328
            PyObject *name_obj = PyDict_GetItemString(view_obj, "name");
×
3329
            if (name_obj && PyUnicode_Check(name_obj)) {
×
3330
                view.with_name(PyUnicode_AsUTF8(name_obj));
×
3331
            }
3332
            PyObject *query_obj = PyDict_GetItemString(view_obj, "query");
×
3333
            if (query_obj && query_obj != Py_None &&
×
3334
                PyUnicode_Check(query_obj)) {
×
3335
                view.with_query(PyUnicode_AsUTF8(query_obj));
×
3336
            }
3337
        }
3338
    }
3339

3340
    std::vector<ChunkDescriptor> chunks;
2✔
3341
    Py_ssize_t num_chunks = PyList_Size(chunks_list);
2!
3342
    chunks.reserve(static_cast<std::size_t>(num_chunks));
2!
3343

3344
    for (Py_ssize_t i = 0; i < num_chunks; i++) {
8✔
3345
        PyObject *chunk_dict = PyList_GetItem(chunks_list, i);
6!
3346
        if (!PyDict_Check(chunk_dict)) {
6!
3347
            PyErr_SetString(PyExc_TypeError, "each chunk must be a dict");
×
3348
            return NULL;
×
3349
        }
3350

3351
        ChunkDescriptor desc;
6✔
3352

3353
        PyObject *cp_idx = PyDict_GetItemString(chunk_dict, "checkpoint_idx");
6!
3354
        PyObject *start = PyDict_GetItemString(chunk_dict, "start_byte");
6!
3355
        PyObject *end = PyDict_GetItemString(chunk_dict, "end_byte");
6!
3356

3357
        if (!cp_idx || !start || !end) {
6!
3358
            PyErr_SetString(
×
3359
                PyExc_KeyError,
3360
                "chunk must have checkpoint_idx, start_byte, end_byte");
3361
            return NULL;
×
3362
        }
3363

3364
        desc.checkpoint_idx =
6✔
3365
            static_cast<std::uint64_t>(PyLong_AsUnsignedLongLong(cp_idx));
6!
3366
        desc.start_byte =
6✔
3367
            static_cast<std::size_t>(PyLong_AsUnsignedLongLong(start));
6!
3368
        desc.end_byte =
6✔
3369
            static_cast<std::size_t>(PyLong_AsUnsignedLongLong(end));
6!
3370

3371
        char filename[64];
3372
        snprintf(filename, sizeof(filename), "chunk-%05llu.arrow",
9✔
3373
                 (unsigned long long)desc.checkpoint_idx);
6✔
3374
        desc.output_file = std::string(output_dir) + "/" + filename;
6!
3375

3376
        chunks.push_back(std::move(desc));
6!
3377
    }
6✔
3378

3379
    std::string file_path = PyUnicode_AsUTF8(self->file_path);
2!
3380
    std::string index_path;
2✔
3381
    const char *idx = PyUnicode_AsUTF8(self->index_dir);
2!
3382
    if (idx && idx[0] != '\0') {
2!
3383
        index_path = idx;
×
3384
    }
3385
    std::size_t checkpoint_size = self->checkpoint_size;
2✔
3386

3387
    WriteViewChunksResult result;
2✔
3388
    std::string error_msg;
2✔
3389

3390
    Py_BEGIN_ALLOW_THREADS try {
2!
3391
        Runtime *rt = get_runtime(self);
2!
3392
        result = rt->submit(write_view_chunks_pipeline(
7!
3393
                                file_path, index_path, checkpoint_size, view,
1!
3394
                                std::move(chunks), compression,
2✔
3395
                                static_cast<std::size_t>(batch_size)),
1✔
3396
                            "write_view_chunks")
1!
3397
                     .get();
2!
3398
    } catch (const std::exception &e) {
1!
3399
        error_msg = e.what();
×
3400
    }
×
3401
    Py_END_ALLOW_THREADS
2!
3402

3403
        if (!error_msg.empty()) {
2!
3404
        PyErr_SetString(PyExc_RuntimeError, error_msg.c_str());
×
3405
        return NULL;
×
3406
    }
3407

3408
    PyObject *dict = PyDict_New();
2!
3409
    if (!dict) return NULL;
2✔
3410

3411
    PyObject *results_list =
1✔
3412
        PyList_New(static_cast<Py_ssize_t>(result.results.size()));
2!
3413
    if (!results_list) {
2!
3414
        Py_DECREF(dict);
×
3415
        return NULL;
×
3416
    }
3417

3418
    for (std::size_t i = 0; i < result.results.size(); i++) {
8✔
3419
        const auto &r = result.results[i];
6✔
3420
        PyObject *item = PyDict_New();
6!
3421
        if (!item) {
6!
3422
            Py_DECREF(results_list);
×
3423
            Py_DECREF(dict);
×
3424
            return NULL;
×
3425
        }
3426
        PyDict_SetItemString(item, "output_file",
6!
3427
                             PyUnicode_FromString(r.output_file.c_str()));
3!
3428
        PyDict_SetItemString(item, "rows_written",
6!
3429
                             PyLong_FromLongLong(r.rows_written));
6!
3430
        PyDict_SetItemString(item, "events_matched",
6!
3431
                             PyLong_FromUnsignedLongLong(r.events_matched));
6!
3432
        if (!r.error.empty()) {
6!
3433
            PyDict_SetItemString(item, "error",
×
3434
                                 PyUnicode_FromString(r.error.c_str()));
×
3435
        }
3436
        PyList_SetItem(results_list, static_cast<Py_ssize_t>(i), item);
6!
3437
    }
3✔
3438

3439
    PyDict_SetItemString(dict, "results", results_list);
2!
3440
    Py_DECREF(results_list);
1!
3441
    PyDict_SetItemString(dict, "total_rows",
2!
3442
                         PyLong_FromLongLong(result.total_rows));
2!
3443
    PyDict_SetItemString(dict, "total_events_matched",
2!
3444
                         PyLong_FromLongLong(result.total_events_matched));
2!
3445

3446
    return dict;
2✔
3447
}
2✔
3448

3449
#endif  // DFTRACER_UTILS_ENABLE_ARROW_IPC
3450

3451
static PyObject *TraceReader_enter(TraceReaderObject *self,
116!
3452
                                   PyObject *Py_UNUSED(ignored)) {
3453
    Py_INCREF(self);
58✔
3454
    return (PyObject *)self;
116✔
3455
}
3456

3457
static PyObject *TraceReader_exit(TraceReaderObject *self, PyObject *args) {
114✔
3458
    Py_RETURN_NONE;
114✔
3459
}
3460

3461
static PyObject *TraceReader_get_file_path(TraceReaderObject *self,
14✔
3462
                                           void *closure) {
3463
    Py_INCREF(self->file_path);
14!
3464
    return self->file_path;
14✔
3465
}
3466

3467
static PyObject *TraceReader_get_index_dir(TraceReaderObject *self,
6✔
3468
                                           void *closure) {
3469
    Py_INCREF(self->index_dir);
6✔
3470
    return self->index_dir;
6✔
3471
}
3472

3473
static PyObject *TraceReader_get_has_index(TraceReaderObject *self,
12✔
3474
                                           void *closure) {
3475
    return PyBool_FromLong(self->has_index);
12✔
3476
}
3477

3478
static PyObject *TraceReader_get_num_lines_prop(TraceReaderObject *self,
8✔
3479
                                                void *closure) {
3480
    try {
3481
        TraceReaderConfig cfg = build_config(self);
8!
3482
        TraceReader reader(std::move(cfg));
8!
3483
        std::size_t n = reader.get_num_lines();
8!
3484
        if (n > 0) return PyLong_FromSize_t(n);
8!
3485
    } catch (...) {
8!
3486
    }
×
3487
    PyObject *empty_args = PyTuple_New(0);
8✔
3488
    if (!empty_args) return NULL;
8✔
3489
    PyObject *list = TraceReader_read_lines(self, empty_args, NULL);
8✔
3490
    Py_DECREF(empty_args);
4✔
3491
    if (!list) return NULL;
8✔
3492
    Py_ssize_t n = PyList_GET_SIZE(list);
8✔
3493
    Py_DECREF(list);
4✔
3494
    return PyLong_FromSsize_t(n);
8✔
3495
}
4✔
3496

3497
static PyObject *TraceReader_get_max_bytes(TraceReaderObject *self,
24✔
3498
                                           PyObject *Py_UNUSED(ignored)) {
3499
    try {
3500
        TraceReaderConfig cfg = build_config(self);
24!
3501
        TraceReader reader(std::move(cfg));
24!
3502
        return PyLong_FromSize_t(reader.get_max_bytes());
24!
3503
    } catch (const std::exception &e) {
24!
3504
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
3505
        return NULL;
×
3506
    }
×
3507
}
12✔
3508

3509
static PyObject *TraceReader_get_num_lines(TraceReaderObject *self,
6✔
3510
                                           PyObject *Py_UNUSED(ignored)) {
3511
    try {
3512
        TraceReaderConfig cfg = build_config(self);
6!
3513
        TraceReader reader(std::move(cfg));
6!
3514
        return PyLong_FromSize_t(reader.get_num_lines());
6!
3515
    } catch (const std::exception &e) {
6!
3516
        PyErr_SetString(PyExc_RuntimeError, e.what());
×
3517
        return NULL;
×
3518
    }
×
3519
}
3✔
3520

3521
static PyMethodDef TraceReader_methods[] = {
3522
    {"iter_lines", (PyCFunction)TraceReader_iter_lines,
3523
     METH_VARARGS | METH_KEYWORDS,
3524
     "Return an iterator over decoded lines.\n"
3525
     "\n"
3526
     "Args:\n"
3527
     "    start_line (int): First line (0 = beginning).\n"
3528
     "    end_line (int): Last line (0 = end of file).\n"
3529
     "    start_byte (int): First byte offset (0 = beginning).\n"
3530
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3531
     "    buffer_size (int): Internal read buffer size in bytes.\n"},
3532
    {"iter_raw", (PyCFunction)TraceReader_iter_raw,
3533
     METH_VARARGS | METH_KEYWORDS,
3534
     "Return an iterator over raw byte chunks.\n"
3535
     "\n"
3536
     "Args:\n"
3537
     "    start_line (int): First line (0 = beginning).\n"
3538
     "    end_line (int): Last line (0 = end of file).\n"
3539
     "    start_byte (int): First byte offset (0 = beginning).\n"
3540
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3541
     "    buffer_size (int): Internal read buffer size in bytes.\n"
3542
     "    line_aligned (bool): Align chunks to line boundaries.\n"
3543
     "    multi_line (bool): Allow multiple lines per chunk.\n"},
3544
    {"read_lines", (PyCFunction)TraceReader_read_lines,
3545
     METH_VARARGS | METH_KEYWORDS,
3546
     "Read all lines and return as list.\n"
3547
     "\n"
3548
     "Args:\n"
3549
     "    start_line (int): First line (0 = beginning).\n"
3550
     "    end_line (int): Last line (0 = end of file).\n"
3551
     "    start_byte (int): First byte offset (0 = beginning).\n"
3552
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3553
     "    buffer_size (int): Internal read buffer size in bytes.\n"},
3554
    {"iter_json", (PyCFunction)TraceReader_iter_json,
3555
     METH_VARARGS | METH_KEYWORDS,
3556
     "Return an iterator over parsed JSON events as Python dicts.\n"
3557
     "\n"
3558
     "Each event is parsed once in C++ (single-pass simdjson ondemand)\n"
3559
     "and yielded as a Python dict. No double-parsing overhead.\n"
3560
     "\n"
3561
     "Args:\n"
3562
     "    start_line (int): First line (0 = beginning).\n"
3563
     "    end_line (int): Last line (0 = end of file).\n"
3564
     "    start_byte (int): First byte offset (0 = beginning).\n"
3565
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3566
     "    buffer_size (int): Internal read buffer size in bytes.\n"
3567
     "    query (str): Optional query filter.\n"
3568
     "    batch_size (int): Events per internal batch (default 1024).\n"},
3569
    {"read_json", (PyCFunction)TraceReader_read_json_py,
3570
     METH_VARARGS | METH_KEYWORDS,
3571
     "Read all events as parsed Python dicts (list).\n"
3572
     "\n"
3573
     "Equivalent to list(iter_json(...)).\n"},
3574
    {"read_raw", (PyCFunction)TraceReader_read_raw,
3575
     METH_VARARGS | METH_KEYWORDS,
3576
     "Read all raw chunks and return as list.\n"
3577
     "\n"
3578
     "Args:\n"
3579
     "    start_line (int): First line (0 = beginning).\n"
3580
     "    end_line (int): Last line (0 = end of file).\n"
3581
     "    start_byte (int): First byte offset (0 = beginning).\n"
3582
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3583
     "    buffer_size (int): Internal read buffer size in bytes.\n"
3584
     "    line_aligned (bool): Align chunks to line boundaries.\n"
3585
     "    multi_line (bool): Allow multiple lines per chunk.\n"},
3586
#ifdef DFTRACER_UTILS_ENABLE_ARROW
3587
    {"iter_arrow", (PyCFunction)TraceReader_iter_arrow,
3588
     METH_VARARGS | METH_KEYWORDS,
3589
     "Return an iterator over Arrow record batches.\n"
3590
     "\n"
3591
     "Args:\n"
3592
     "    batch_size (int): Maximum rows per Arrow batch.\n"
3593
     "    start_line (int): First line (0 = beginning).\n"
3594
     "    end_line (int): Last line (0 = end of file).\n"
3595
     "    start_byte (int): First byte offset (0 = beginning).\n"
3596
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3597
     "    buffer_size (int): Internal read buffer size in bytes.\n"},
3598
    {"iter_arrow_stream", (PyCFunction)TraceReader_iter_arrow_stream,
3599
     METH_VARARGS | METH_KEYWORDS,
3600
     "Return an _ArrowBatchStream that exposes Arrow record batches\n"
3601
     "via the Arrow C Data Interface stream protocol\n"
3602
     "(__arrow_c_stream__). PyArrow can drain the producer channel\n"
3603
     "with a single call, without per-batch Python iteration.\n"},
3604
    {"read_arrow", (PyCFunction)TraceReader_read_arrow,
3605
     METH_VARARGS | METH_KEYWORDS,
3606
     "Read all events as a materialized ArrowTable.\n"
3607
     "\n"
3608
     "Args:\n"
3609
     "    batch_size (int): Maximum rows per Arrow batch.\n"
3610
     "    start_line (int): First line (0 = beginning).\n"
3611
     "    end_line (int): Last line (0 = end of file).\n"
3612
     "    start_byte (int): First byte offset (0 = beginning).\n"
3613
     "    end_byte (int): Last byte offset (0 = end of file).\n"
3614
     "    buffer_size (int): Internal read buffer size in bytes.\n"},
3615
#endif
3616
#ifdef DFTRACER_UTILS_ENABLE_ARROW_IPC
3617
    {"write_arrow", (PyCFunction)TraceReader_write_arrow,
3618
     METH_VARARGS | METH_KEYWORDS,
3619
     "Write trace data to partitioned Arrow IPC files.\n"
3620
     "\n"
3621
     "Args:\n"
3622
     "    path (str): Output directory path.\n"
3623
     "    partition_by (list[str] or None): Column names to partition by.\n"
3624
     "    num_buckets (int): Number of hash buckets (0 = no bucketing).\n"
3625
     "    chunk_size_mb (int): Max uncompressed MB per file (default 32).\n"
3626
     "    compression (str): 'zstd' or 'none' (default 'zstd').\n"
3627
     "    batch_size (int): Rows per internal batch (default 10000).\n"
3628
     "    normalize (bool): Use normalized schema (default False).\n"
3629
     "\n"
3630
     "Returns:\n"
3631
     "    dict: Statistics including partitions, total_rows, total_bytes.\n"},
3632
    {"get_view_chunks", (PyCFunction)TraceReader_get_view_chunks,
3633
     METH_VARARGS | METH_KEYWORDS,
3634
     "Get candidate chunks for a view after bloom filter pruning.\n"
3635
     "\n"
3636
     "Args:\n"
3637
     "    view (str or dict): View name ('io', 'compute', 'dlio') or\n"
3638
     "                        dict with 'name' and optional 'query'.\n"
3639
     "\n"
3640
     "Returns:\n"
3641
     "    dict: chunks list, total_checkpoints, skipped_checkpoints.\n"},
3642
    {"write_view_chunk", (PyCFunction)TraceReader_write_view_chunk,
3643
     METH_VARARGS | METH_KEYWORDS,
3644
     "Write a single chunk to an Arrow IPC file.\n"
3645
     "\n"
3646
     "Args:\n"
3647
     "    output_file (str): Path to output Arrow IPC file.\n"
3648
     "    checkpoint_idx (int): Checkpoint index.\n"
3649
     "    start_byte (int): Start byte offset.\n"
3650
     "    end_byte (int): End byte offset.\n"
3651
     "    view (str or dict): View definition.\n"
3652
     "    compression (str): 'zstd' or 'none' (default 'zstd').\n"
3653
     "    batch_size (int): Events per batch (default 10000).\n"
3654
     "\n"
3655
     "Returns:\n"
3656
     "    dict: output_file, events_matched, rows_written, bytes_written.\n"},
3657
    {"write_view_chunks", (PyCFunction)TraceReader_write_view_chunks,
3658
     METH_VARARGS | METH_KEYWORDS,
3659
     "Write multiple chunks to Arrow IPC files in parallel.\n"
3660
     "\n"
3661
     "All chunks are processed concurrently on the Runtime thread pool.\n"
3662
     "\n"
3663
     "Args:\n"
3664
     "    chunks (list): List of dicts with checkpoint_idx, start_byte, "
3665
     "end_byte.\n"
3666
     "    output_dir (str): Directory for output Arrow IPC files.\n"
3667
     "    view (str or dict): View definition.\n"
3668
     "    compression (str): 'zstd' or 'none' (default 'zstd').\n"
3669
     "    batch_size (int): Events per batch (default 10000).\n"
3670
     "\n"
3671
     "Returns:\n"
3672
     "    dict: results list, total_rows, total_events_matched.\n"},
3673
#endif
3674
    {"get_max_bytes", (PyCFunction)TraceReader_get_max_bytes, METH_NOARGS,
3675
     "Get the maximum byte position (0 if unknown for compressed\n"
3676
     "files without index)."},
3677
    {"get_num_lines", (PyCFunction)TraceReader_get_num_lines, METH_NOARGS,
3678
     "Get the total number of lines (0 if unknown for files without\n"
3679
     "index)."},
3680
    {"__enter__", (PyCFunction)TraceReader_enter, METH_NOARGS,
3681
     "Enter the runtime context for the with statement."},
3682
    {"__exit__", (PyCFunction)TraceReader_exit, METH_VARARGS,
3683
     "Exit the runtime context for the with statement.\n"
3684
     "\n"
3685
     "TraceReader does not own the shared RocksDB instance for an index path;\n"
3686
     "any shared DB lifetime remains manager-owned on the native side."},
3687
    {NULL}};
3688

3689
static PyGetSetDef TraceReader_getsetters[] = {
3690
    {"path", (getter)TraceReader_get_file_path, NULL,
3691
     "Path to the trace file or directory", NULL},
3692
    {"index_dir", (getter)TraceReader_get_index_dir, NULL,
3693
     "Directory for index files", NULL},
3694
    {"has_index", (getter)TraceReader_get_has_index, NULL,
3695
     "True if a checkpoint index was found", NULL},
3696
    {"num_lines", (getter)TraceReader_get_num_lines_prop, NULL,
3697
     "Total line count (reads all lines if needed)", NULL},
3698
    {NULL}};
3699

3700
PyTypeObject TraceReaderType = {
3701
    PyVarObject_HEAD_INIT(NULL, 0) "dftracer_utils_ext.TraceReader",
3702
    sizeof(TraceReaderObject),                /* tp_basicsize */
3703
    0,                                        /* tp_itemsize */
3704
    (destructor)TraceReader_dealloc,          /* tp_dealloc */
3705
    0,                                        /* tp_vectorcall_offset */
3706
    0,                                        /* tp_getattr */
3707
    0,                                        /* tp_setattr */
3708
    0,                                        /* tp_as_async */
3709
    0,                                        /* tp_repr */
3710
    0,                                        /* tp_as_number */
3711
    0,                                        /* tp_as_sequence */
3712
    0,                                        /* tp_as_mapping */
3713
    0,                                        /* tp_hash */
3714
    0,                                        /* tp_call */
3715
    0,                                        /* tp_str */
3716
    0,                                        /* tp_getattro */
3717
    0,                                        /* tp_setattro */
3718
    0,                                        /* tp_as_buffer */
3719
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
3720
    "TraceReader(file_path: str, index_dir: str = '',\n"
3721
    "            checkpoint_size: int = 33554432,\n"
3722
    "            auto_build_index: bool = False,\n"
3723
    "            runtime: Runtime | None = None)\n"
3724
    "--\n"
3725
    "\n"
3726
    "Smart trace file reader that auto-selects sequential or indexed\n"
3727
    "reading based on whether a ``.dftindex`` store exists.\n"
3728
    "\n"
3729
    "Args:\n"
3730
    "    file_path (str): Path to the trace file (.pfw.gz or plain "
3731
    "text).\n"
3732
    "    index_dir (str): Directory to search for ``.dftindex`` "
3733
    "stores.\n"
3734
    "        Empty string (default) searches next to the trace file.\n"
3735
    "    checkpoint_size (int): Checkpoint interval in bytes for index\n"
3736
    "        building (default 32 MB).\n"
3737
    "    auto_build_index (bool): If True, automatically build an "
3738
    "index\n"
3739
    "        when none exists.\n"
3740
    "    runtime (Runtime or None): Runtime instance for thread pool "
3741
    "control.\n"
3742
    "        If None, uses the default global Runtime.\n"
3743
    "\n"
3744
    "Raises:\n"
3745
    "    RuntimeError: If *file_path* does not exist or cannot be "
3746
    "opened.\n",                /* tp_doc */
3747
    0,                          /* tp_traverse */
3748
    0,                          /* tp_clear */
3749
    0,                          /* tp_richcompare */
3750
    0,                          /* tp_weaklistoffset */
3751
    0,                          /* tp_iter */
3752
    0,                          /* tp_iternext */
3753
    TraceReader_methods,        /* tp_methods */
3754
    0,                          /* tp_members */
3755
    TraceReader_getsetters,     /* tp_getset */
3756
    0,                          /* tp_base */
3757
    0,                          /* tp_dict */
3758
    0,                          /* tp_descr_get */
3759
    0,                          /* tp_descr_set */
3760
    0,                          /* tp_dictoffset */
3761
    (initproc)TraceReader_init, /* tp_init */
3762
    0,                          /* tp_alloc */
3763
    TraceReader_new,            /* tp_new */
3764
};
3765

3766
int init_trace_reader(PyObject *m) {
2✔
3767
    if (PyType_Ready(&TraceReaderType) < 0) return -1;
2✔
3768

3769
    Py_INCREF(&TraceReaderType);
1✔
3770
    if (PyModule_AddObject(m, "TraceReader", (PyObject *)&TraceReaderType) <
2✔
3771
        0) {
3772
        Py_DECREF(&TraceReaderType);
3773
        Py_DECREF(m);
3774
        return -1;
×
3775
    }
3776

3777
    return 0;
2✔
3778
}
1✔
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc