diff --git a/CHANGELOG.md b/CHANGELOG.md index a9ae9a0..9bb7ba9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ number says nothing about any other port. See the Releases section of - **`DatasetReader::Scan` binds a row under the default `max_message_size`** (#164), as protowire-go does. The body a row is bound as repeats the column names, so it is wider than the row, and a row whose bytes fit the caller's limit could still be refused when bound. The row's bytes are still capped as they are read. The other limits apply as before. - **A broken `@dataset` header is refused with the error `Parse` gives** (#164), as protowire-go does. `Create` said `no @dataset directive in stream` for a header cut short, and gave its own text for an unterminated string in it. The reader now finds the header's and each row's end with the lexer `Parse` uses rather than a byte scanner of its own. Streaming is about 15% faster on a 300,000-row table. - **A timestamp cell right before a row's `)` parses** (#169), as protowire-go does. `(2024-01-02T03:04:05Z)` was refused with `invalid timestamp: 2024-01-02T03:04:05Z)`, because the lexer read the `)` into the timestamp. A space before the `)`, or a later cell, hid it. +- **`DatasetReader` reads a long token in time linear in its length** (#168), as protowire-go does. A string, bytes literal, comment or number spread over many reads was lexed again from its start on every 4 KiB read: an unclosed 1 MiB string cost 136 MB of lexing, and 2 MiB cost 540 MB. The scan now lexes a token still open again only once the bytes read since could end it, which is Go's `mayEndToken` rule, and the header scan resumes where it stopped too. The same 1 MiB string now costs 1.05 MB. ## [2.0.0] — 2026-11-05 diff --git a/include/protowire/pxf/dataset_reader.h b/include/protowire/pxf/dataset_reader.h index 47f3df5..5b312fb 100644 --- a/include/protowire/pxf/dataset_reader.h +++ b/include/protowire/pxf/dataset_reader.h @@ -39,6 +39,8 @@ namespace protowire::pxf { +struct BraceMatch; // src/pxf/parser_internal.h + // Default cap on the @dataset header (leading directives plus the // `@dataset TYPE ( cols )` declaration). Real headers are tiny — a few // hundred bytes at most. The cap exists to fail-fast on misuse: a @@ -121,7 +123,12 @@ class DatasetReader { // which will desync the next reader. std::unique_ptr Tail(); + ~DatasetReader(); + DatasetReader(const DatasetReader&) = delete; + DatasetReader& operator=(const DatasetReader&) = delete; + private: + friend class DatasetReaderTestPeer; // reads scanned_ DatasetReader() = default; // What ScanRow or ScanHeader found; see dataset_reader.cc. @@ -135,8 +142,11 @@ class DatasetReader { std::optional TryWholeRow(); Status RowError(int start) const; RowScan ScanRow(int* start, int* end); - RowScan ScanHeader(int* end) const; + RowScan ScanHeader(int* end); bool Complete(int tok_end) const; + void Wait(int at); + bool Resume(); + bool MayEndToken() const; std::istream* src_ = nullptr; Position PosAt(size_t i) const; @@ -149,6 +159,22 @@ class DatasetReader { // or -1 between rows. Both index pending_. int scan_pos_ = 0; int row_start_ = -1; + // While a token runs past the buffer's end, the scan does not lex it again + // until bytes arrive that could end it (MayEndToken): wait_at_ is the + // token's start, wait_from_ the buffer length already checked. + bool waiting_ = false; + int wait_at_ = 0; + int wait_from_ = 0; + // Where the header scan stands: past the @dataset token, how far into an + // @proto directive (1 after @proto, 2 after its type name), and in its + // body, whose braces are matched as bytes. + bool hdr_in_dataset_ = false; + int hdr_proto_ = 0; + bool hdr_in_body_ = false; + std::unique_ptr hdr_body_; + // The bytes the header and row scans have lexed, for the tests that hold + // them linear in the input. + size_t scanned_ = 0; bool src_eof_ = false; bool finished_ = false; Status err_; diff --git a/src/pxf/dataset_reader.cc b/src/pxf/dataset_reader.cc index 345307c..252a5b9 100644 --- a/src/pxf/dataset_reader.cc +++ b/src/pxf/dataset_reader.cc @@ -323,8 +323,12 @@ void DatasetReader::Consume(size_t n) { const int k = static_cast(n); scan_pos_ = std::max(scan_pos_ - k, 0); if (row_start_ >= 0) row_start_ = row_start_ >= k ? row_start_ - k : -1; + // Bytes are consumed only between tokens, when none is waiting. + waiting_ = false; } +DatasetReader::~DatasetReader() = default; + Status DatasetReader::Fail(Status st) { err_ = std::move(st); return err_; @@ -373,19 +377,29 @@ Status DatasetReader::RowError(int start) const { } // ScanRow finds the next row with the lexer Parse uses, resuming at -// scan_pos_, so a complete token is not lexed again on the next read. +// scan_pos_, so a complete token is not lexed again on the next read, and a +// token still open is lexed again only once bytes arrive that could end it +// (Resume): a row costs time linear in its size however many reads it +// arrives in (#168). // Between rows it skips newlines and comments: `(` starts a row, a directive // ends the sequence, and anything else is not a row. Inside a row it stops // at the `)` that closes it, or at a `(` (no value holds one) or a token the // lexer rejects, which make the row bad. Returns kNeedMore with *start the // open row's `(`, or -1 between rows, when the buffer ends first. DatasetReader::RowScan DatasetReader::ScanRow(int* start, int* end) { + if (!Resume()) { + *start = row_start_; + return RowScan::kNeedMore; + } Lexer lx(std::string_view(pending_).substr(scan_pos_), Position{1, 1, scan_pos_}); lx.SetMaxBytesLiteral(MaxBytesLiteral()); for (;;) { + const int from = lx.CurrentPos().offset; const Token tok = lx.Next(); const int tok_end = lx.CurrentPos().offset; + scanned_ += static_cast(tok_end - from); if (tok.kind == TokenKind::kEOF || !Complete(tok_end)) { + if (!src_eof_) Wait(tok.pos.offset); *start = row_start_; return RowScan::kNeedMore; } @@ -464,54 +478,179 @@ bool DatasetReader::Complete(int tok_end) const { return src_eof_ || tok_end + kLookahead <= static_cast(pending_.size()); } +// Wait records the token at pending_[at], which the scan could not yet take +// as complete. The scan resumes at its start: the spaces before it are +// decided, and lexing them again would cost as much as the token. +void DatasetReader::Wait(int at) { + waiting_ = true; + wait_at_ = at; + wait_from_ = static_cast(pending_.size()); + scan_pos_ = at; +} + +// Resume reports whether a scan may lex on: no token is waiting, the source +// has ended, or the bytes read since the waiting token was last lexed could +// end it. +bool DatasetReader::Resume() { + if (waiting_ && !src_eof_ && !MayEndToken()) { + wait_from_ = static_cast(pending_.size()); + return false; + } + waiting_ = false; + return true; +} + +namespace { + +// kShortWait is the length under which a waiting token is lexed again +// whatever bytes arrive. What may still change a token this short is a byte +// or two of the lexer's lookahead (`""` opening `"""`, `b` opening `b"`, `/` +// opening a comment, `-in` reading `-inf`, a multi-byte rune cut in half, no +// token at all yet), and lexing a few bytes again costs nothing. +constexpr int kShortWait = 8; + +bool IsDigitByte(char c) { + return c >= '0' && c <= '9'; +} + +// As the lexer's IsIdentPart. +bool IsIdentByte(char c) { + return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_' || IsDigitByte(c) || c == '.'; +} + +// The bytes that end a timestamp, as the lexer's LexTimestamp has them. +bool EndsRun(char c) { + return c == ' ' || c == '\n' || c == '\t' || c == '\r' || c == ',' || c == ']' || c == '}' || + c == ')' || c == '#'; +} + +} // namespace + +// MayEndToken reports whether the bytes read since the waiting token was +// last lexed could end it: a closing quote for a string (one not escaped by +// an odd run of backslashes) or a newline, which a simple string cannot +// hold; `"""` for a triple-quoted one; a quote or newline for a bytes +// literal, or enough characters that the lexer refuses it +// (max_bytes_literal_length); a newline for a line comment; `*/` for a block +// comment; a byte that ends a run, or a comment opener, for a timestamp; and +// for any other token a byte that cannot continue an identifier, number or +// duration. Each of those does end the token, so a token is lexed again only +// once it is short or decided, and the scan stays linear. It decides only +// when the lexer looks again, never what the token is. Ported from +// protowire-go's mayEndToken. +bool DatasetReader::MayEndToken() const { + const std::string_view buf = pending_; + const int at = wait_at_; + const int from = wait_from_; + if (from - at < kShortWait) return true; + const std::string_view tok = buf.substr(at); + const std::string_view fresh = buf.substr(from); + if (tok.substr(0, 3) == "\"\"\"") { + return buf.substr(std::max(from - 2, at + 3)).find("\"\"\"") != std::string_view::npos; + } + if (tok[0] == '"') { + for (size_t i = 0; i < fresh.size(); ++i) { + if (fresh[i] == '\n') return true; // a simple string cannot hold one + if (fresh[i] != '"') continue; + int backslashes = 0; + for (int j = from + static_cast(i) - 1; j > at && buf[j] == '\\'; --j) ++backslashes; + if (backslashes % 2 == 0) return true; + } + return false; + } + if (tok.substr(0, 2) == "b\"") { + // The lexer checks the cap before each character it reads past the + // opening quote, the last one included. + const int chars = static_cast(buf.size()) - 1 - (at + 2); + const int limit = MaxBytesLiteral(); + return fresh.find_first_of("\"\n") != std::string_view::npos || + (limit > 0 && chars / 4 * 3 > limit); + } + if (tok[0] == '#' || tok.substr(0, 2) == "//") return fresh.find('\n') != std::string_view::npos; + if (tok.substr(0, 2) == "/*") { + return buf.substr(std::max(from - 1, at + 2)).find("*/") != std::string_view::npos; + } + if (IsDigitByte(tok[0]) && IsDigitByte(tok[1]) && IsDigitByte(tok[2]) && IsDigitByte(tok[3]) && + tok[4] == '-') { + // A timestamp runs to a byte that ends a run, or to `//` or `/*`, whose + // '/' may be the last byte checked before. + for (size_t i = std::max(from - 1, at); i < buf.size(); ++i) { + if (EndsRun(buf[i])) return true; + if (buf[i] == '/' && i + 1 < buf.size() && (buf[i + 1] == '/' || buf[i + 1] == '*')) { + return true; + } + } + return false; + } + for (const char c : fresh) { + if (!(IsIdentByte(c) || c == '+' || c == '-' || static_cast(c) >= 0x80)) { + return true; + } + } + return false; +} + // ScanHeader finds the end of the first `@dataset TYPE ( cols )` header with // the lexer Parse uses: the offset just past the `)` that closes the column // list, the first `)` after an @dataset token. An @proto body is raw // protobuf source, which Parse skips by matching its braces, so the scan // does too. Returns kNeedMore when the buffer ends first, and kBad at a -// token the lexer rejects or an @proto body that can never close. -DatasetReader::RowScan DatasetReader::ScanHeader(int* end) const { - Lexer lx(pending_); - lx.SetMaxBytesLiteral(MaxBytesLiteral()); - bool in_dataset = false; - int proto = 0; // 1 after @proto, 2 after its type name +// token the lexer rejects or an @proto body that can never close. Like +// ScanRow it resumes where the last read left it, so a header costs time +// linear in its size however many reads it arrives in (#168). +DatasetReader::RowScan DatasetReader::ScanHeader(int* end) { + if (!hdr_in_body_ && !Resume()) return RowScan::kNeedMore; for (;;) { - const Token tok = lx.Next(); - const int tok_end = lx.CurrentPos().offset; - if (tok.kind == TokenKind::kEOF || !Complete(tok_end)) return RowScan::kNeedMore; - switch (tok.kind) { - case TokenKind::kIllegal: - return RowScan::kBad; - case TokenKind::kNewline: - case TokenKind::kComment: + if (hdr_in_body_) { + BraceMatch::State st = BraceMatch::State::kOpen; + const int i = hdr_body_->Scan(pending_, scan_pos_, &st); + scanned_ += static_cast(i - scan_pos_); + scan_pos_ = i; + if (st == BraceMatch::State::kOpen) return RowScan::kNeedMore; + if (st == BraceMatch::State::kFailed) return RowScan::kBad; + hdr_in_body_ = false; + ++scan_pos_; // the closing `}` + } + Lexer lx(std::string_view(pending_).substr(scan_pos_), Position{1, 1, scan_pos_}); + lx.SetMaxBytesLiteral(MaxBytesLiteral()); + for (;;) { + const int from = lx.CurrentPos().offset; + const Token tok = lx.Next(); + const int tok_end = lx.CurrentPos().offset; + scanned_ += static_cast(tok_end - from); + if (tok.kind == TokenKind::kEOF && src_eof_) return RowScan::kNeedMore; + if (tok.kind == TokenKind::kEOF || !Complete(tok_end)) { + Wait(tok.pos.offset); + return RowScan::kNeedMore; + } + if (tok.kind == TokenKind::kIllegal) return RowScan::kBad; + scan_pos_ = tok_end; + if (tok.kind == TokenKind::kNewline || tok.kind == TokenKind::kComment) { // Parse skips these between @proto, its name and `{`. continue; - default: - break; - } - if (in_dataset) { - if (tok.kind == TokenKind::kRParen) { - *end = tok_end; - return RowScan::kRow; } - continue; - } - if (tok.kind == TokenKind::kAtDataset) { - in_dataset = true; - proto = 0; - } else if (tok.kind == TokenKind::kAtProto) { - proto = 1; - } else if (tok.kind == TokenKind::kIdent && proto == 1) { - proto = 2; - } else if (tok.kind == TokenKind::kLBrace && proto > 0) { - bool failed = false; - const int close = FindMatchingBrace(pending_, tok.pos.offset, &failed); - if (failed) return RowScan::kBad; - if (close < 0) return RowScan::kNeedMore; - lx.RepositionTo(close + 1); - proto = 0; - } else { - proto = 0; + if (hdr_in_dataset_) { + if (tok.kind == TokenKind::kRParen) { + *end = tok_end; + return RowScan::kRow; + } + continue; + } + if (tok.kind == TokenKind::kAtDataset) { + hdr_in_dataset_ = true; + hdr_proto_ = 0; + } else if (tok.kind == TokenKind::kAtProto) { + hdr_proto_ = 1; + } else if (tok.kind == TokenKind::kIdent && hdr_proto_ == 1) { + hdr_proto_ = 2; + } else if (tok.kind == TokenKind::kLBrace && hdr_proto_ > 0) { + hdr_proto_ = 0; + hdr_in_body_ = true; + hdr_body_ = std::make_unique(); + break; // the body is matched as bytes, not lexed + } else { + hdr_proto_ = 0; + } } } } diff --git a/src/pxf/parser.cc b/src/pxf/parser.cc index 5236305..2b97efa 100644 --- a/src/pxf/parser.cc +++ b/src/pxf/parser.cc @@ -1065,13 +1065,17 @@ StatusOr> Parser::ParseRowCell() { } // namespace int FindMatchingBrace(std::string_view input, int open_offset, bool* failed) { - enum class Mode : uint8_t { kCode, kString, kTripleString, kBytes, kLineComment, kBlockComment }; - Mode mode = Mode::kCode; - int depth = 1; - bool esc = false; - int run = 0; // quotes just seen in a triple string, or 1 after a '*' + BraceMatch m; + BraceMatch::State st = BraceMatch::State::kOpen; + const int i = m.Scan(input, open_offset + 1, &st); + if (st == BraceMatch::State::kClosed) return i; + if (st == BraceMatch::State::kFailed && failed != nullptr) *failed = true; + return -1; +} + +int BraceMatch::Scan(std::string_view input, int i, State* state) { const int n = static_cast(input.size()); - for (int i = open_offset + 1; i < n; ++i) { + for (; i < n; ++i) { const char ch = input[i]; switch (mode) { case Mode::kString: @@ -1082,8 +1086,8 @@ int FindMatchingBrace(std::string_view input, int open_offset, bool* failed) { } else if (ch == '"') { mode = Mode::kCode; } else if (ch == '\n') { - if (failed != nullptr) *failed = true; - return -1; + *state = State::kFailed; + return i; } continue; case Mode::kTripleString: @@ -1098,8 +1102,8 @@ int FindMatchingBrace(std::string_view input, int open_offset, bool* failed) { if (ch == '"') { mode = Mode::kCode; } else if (ch == '\n') { - if (failed != nullptr) *failed = true; - return -1; + *state = State::kFailed; + return i; } continue; case Mode::kLineComment: @@ -1121,10 +1125,17 @@ int FindMatchingBrace(std::string_view input, int open_offset, bool* failed) { ++depth; break; case '}': - if (--depth == 0) return i; + if (--depth == 0) { + *state = State::kClosed; + return i; + } break; case '"': - if (i + 2 >= n) return -1; + // Whether it opens a triple-quoted string needs the next two bytes. + if (i + 2 >= n) { + *state = State::kOpen; + return i; + } if (input[i + 1] == '"' && input[i + 2] == '"') { mode = Mode::kTripleString; i += 2; @@ -1137,7 +1148,10 @@ int FindMatchingBrace(std::string_view input, int open_offset, bool* failed) { break; case 'b': case '/': { - if (i + 1 >= n) return -1; + if (i + 1 >= n) { + *state = State::kOpen; + return i; + } const char next = input[i + 1]; if (ch == 'b' && next == '"') { mode = Mode::kBytes; @@ -1155,7 +1169,8 @@ int FindMatchingBrace(std::string_view input, int open_offset, bool* failed) { break; } } - return -1; + *state = State::kOpen; + return i; } StatusOr ParseDatasetAt(Lexer* lex, Token* current, int max_depth) { diff --git a/src/pxf/parser_internal.h b/src/pxf/parser_internal.h index 055df7a..cc0fc3e 100644 --- a/src/pxf/parser_internal.h +++ b/src/pxf/parser_internal.h @@ -6,6 +6,7 @@ #pragma once +#include #include #include @@ -27,6 +28,22 @@ StatusOr ParseDatasetAt(Lexer* lex, Token* current, int max_de // a newline, *failed (if given) is set; otherwise the input ends first. int FindMatchingBrace(std::string_view input, int open_offset, bool* failed = nullptr); +// BraceMatch matches an @proto body's braces as FindMatchingBrace does, and +// can resume: Scan reads input from i and returns where it stopped, with +// *state kClosed at the matching `}`, kFailed where the body can never +// close, or kOpen where the input ended first, to scan again from there +// once more input has arrived. Mirrors protowire-go's braceMatch. +struct BraceMatch { + enum class State : uint8_t { kOpen, kClosed, kFailed }; + enum class Mode : uint8_t { kCode, kString, kTripleString, kBytes, kLineComment, kBlockComment }; + int Scan(std::string_view input, int i, State* state); + + int depth = 1; // the opening `{` is read + Mode mode = Mode::kCode; + bool esc = false; // in a string, after a backslash + int run = 0; // quotes just seen in a triple string, or 1 after a '*' +}; + // The error for anything after a @dataset header that is neither a row nor // a directive (draft §3.4.4). inline constexpr const char* kDatasetBesideEntries = diff --git a/test/pxf_dataset_reader_test.cc b/test/pxf_dataset_reader_test.cc index d926df6..7ebd5f8 100644 --- a/test/pxf_dataset_reader_test.cc +++ b/test/pxf_dataset_reader_test.cc @@ -19,6 +19,14 @@ #include #include +namespace protowire::pxf { +// Reads the bytes a reader's scans have lexed (#168). +class DatasetReaderTestPeer { + public: + static size_t& Scanned(DatasetReader& r) { return r.scanned_; } +}; +} // namespace protowire::pxf + namespace { namespace pb = google::protobuf; @@ -579,4 +587,81 @@ TEST(DatasetTimestampCell, EndsAtTheRowsParen) { EXPECT_NE(std::get_if>(&*row.cells[0]), nullptr); } +// ---- The scans are linear in their input (#168) --------------------------- + +// The row scan lexes each byte of a row about once, however many reads the +// row arrives in: a token still open at a read's end is lexed again only once +// bytes arrive that could end it. These count the bytes the scan lexes rather +// than timing it, so a slow machine cannot fail them. The shapes are +// protowire-go's (dataset_scan_linear_test.go). Without the waiting rule, an +// unclosed 1 MiB row in 4 KiB reads is lexed about 128 MB. +TEST(DatasetReaderScan, IsLinear) { + struct Shape { + const char* name; + std::string (*body)(int n); + }; + auto rep = [](const std::string& s, int count) { + std::string out; + out.reserve(s.size() * count); + for (int i = 0; i < count; ++i) out += s; + return out; + }; + static decltype(rep)* r = &rep; + const Shape shapes[] = { + {"short tokens", [](int n) { return (*r)("1 ", n / 2); }}, + {"unterminated string", [](int n) { return "\"" + std::string(n, 'a'); }}, + {"string of escaped quotes", [](int n) { return "\"" + (*r)("a\\\"", n / 3); }}, + {"unterminated triple string", [](int n) { return "\"\"\"" + (*r)("a\"\"b", n / 4); }}, + {"one long number", [](int n) { return std::string(n, '7'); }}, + {"line comment", [](int n) { return "1 # " + std::string(n, 'x'); }}, + {"unterminated block comment", [](int n) { return "1 /* " + (*r)("/ * x", n / 5); }}, + // A timestamp runs to a byte that ends a run; '=' and '/' are not one, + // and must not make the scan lex it again. + {"one long timestamp", [](int n) { return "2026-10-23T" + (*r)("0=/", n / 3); }}, + {"one long identifier", [](int n) { return "a" + (*r)("b.", n / 2); }}, + {"one long bytes literal", [](int n) { return "b\"" + std::string(n, 'Q'); }}, + }; + for (const Shape& sh : shapes) { + for (const int size : {1 << 20, 2 << 20}) { + const std::string in = "@dataset T (a)\n(" + sh.body(size); + std::istringstream src(in); + auto tr = DatasetReader::Create(&src); + ASSERT_TRUE(tr.ok()) << sh.name << ": " << tr.status().ToString(); + size_t& scanned = protowire::pxf::DatasetReaderTestPeer::Scanned(**tr); + scanned = 0; // the row, not the header + DatasetRow row; + EXPECT_FALSE((*tr)->Next(&row).ok()) << sh.name << ": an unclosed row was accepted"; + // A token cut by a read is lexed again once. + EXPECT_GT(scanned, 0u) << sh.name; + EXPECT_LE(scanned, 2 * in.size()) << sh.name << ", " << size << " bytes"; + } + } +} + +// The header scan resumes too: a header near the 64 KiB budget, read 4 KiB +// at a time, was lexed from its start on every read. +TEST(DatasetReaderScan, HeaderIsLinear) { + constexpr int n = 60 << 10; + auto rep = [](const std::string& s, int count) { + std::string out; + for (int i = 0; i < count; ++i) out += s; + return out; + }; + const std::pair shapes[] = { + {"short tokens", "@x { " + rep("a = 1 ", n / 6) + "}\n"}, + {"one long string", "@x { k = \"" + std::string(n, 'a') + "\" }\n"}, + {"@proto body", "@proto { " + rep("string a = 1; ", n / 14) + "}\n"}, + }; + for (const auto& [name, header] : shapes) { + const std::string in = header + "@dataset T (a)\n(1)\n"; + std::istringstream src(in); + auto tr = DatasetReader::Create(&src); + ASSERT_TRUE(tr.ok()) << name << ": " << tr.status().ToString(); + EXPECT_EQ((*tr)->Columns(), std::vector{"a"}) << name; + const size_t scanned = protowire::pxf::DatasetReaderTestPeer::Scanned(**tr); + EXPECT_GT(scanned, 0u) << name; + EXPECT_LE(scanned, 2 * in.size()) << name; + } +} + } // namespace