pxf: DatasetReader returns a complete row from a stream that stays open (#171)

CHANGELOG.md

@@ -21,6 +21,7 @@ number says nothing about any other port. See the Releases section of
 - **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.
+- **`DatasetReader` returns a complete row from a stream that stays open** (#171), as protowire-go does. On a pipe, a socket or a file still being appended to, `Next` could wait for 4 KiB more input, or for the stream to end, before returning a row that had already arrived. The read took a full 4 KiB (`std::istream::read`), and the scan held a token until four more bytes followed it. `Next` now reads what the stream has, blocking only when nothing is buffered. It takes a token as complete as soon as the bytes the lexer looks at past it have arrived. Input that ends reads as before.
 
 ## [2.0.0] — 2026-11-05
 

include/protowire/pxf/dataset_reader.h

@@ -143,7 +143,7 @@ class DatasetReader {
   Status RowError(int start) const;
   RowScan ScanRow(int* start, int* end);
   RowScan ScanHeader(int* end);
-  bool Complete(int tok_end) const;
+  bool Complete(const Token& tok, int tok_end) const;
   void Wait(int at);
   bool Resume();
   bool MayEndToken() const;

src/pxf/dataset_reader.cc

@@ -398,7 +398,7 @@ DatasetReader::RowScan DatasetReader::ScanRow(int* start, int* end) {
     const Token tok = lx.Next();
     const int tok_end = lx.CurrentPos().offset;
     scanned_ += static_cast<size_t>(tok_end - from);
-    if (tok.kind == TokenKind::kEOF || !Complete(tok_end)) {
+    if (tok.kind == TokenKind::kEOF || !Complete(tok, tok_end)) {
       if (!src_eof_) Wait(tok.pos.offset);
       *start = row_start_;
       return RowScan::kNeedMore;
@@ -443,39 +443,82 @@ DatasetReader::RowScan DatasetReader::ScanRow(int* start, int* end) {
   }
 }
 
+// Pull reads up to n bytes, taking what the stream already has rather than
+// waiting for all n: on a stream that stays open, a complete row must not
+// wait for bytes that may never come (#171). It blocks only when nothing is
+// buffered, for one byte or the end of the stream.
 Status DatasetReader::Pull(size_t n) {
   if (src_eof_) return Status::OK();
-  std::string buf(n, '\0');
-  src_->read(buf.data(), static_cast<std::streamsize>(n));
-  std::streamsize got = src_->gcount();
-  if (got > 0) {
-    pending_.append(buf.data(), static_cast<size_t>(got));
-    // A stream that never ends a row cannot grow the buffer without bound.
-    const int max_size = opts_.max_message_size > 0 ? opts_.max_message_size : kMaxMessageSize;
-    if (pending_.size() > static_cast<size_t>(max_size)) {
-      return Status::Error(
-          "dataset input holds " + std::to_string(pending_.size()) +
-          " bytes without a row boundary, MaxMessageSize=" + std::to_string(max_size));
-    }
-  }
-  if (src_->eof()) {
+  using Traits = std::istream::traits_type;
+  std::streambuf* buf = src_->rdbuf();
+  if (buf == nullptr) return Status::Error("pxf: DatasetReader: istream has no buffer");
+  if (buf->in_avail() <= 0 && Traits::eq_int_type(src_->peek(), Traits::eof())) {
+    if (src_->bad()) return Status::Error("pxf: DatasetReader: istream read error");
     src_eof_ = true;
     return Status::OK();
   }
+  const std::streamsize avail = buf->in_avail();
+  const size_t want = avail > 0 ? std::min(n, static_cast<size_t>(avail)) : 1;
+  const size_t old = pending_.size();
+  pending_.resize(old + want);
+  const std::streamsize got =
+      src_->readsome(pending_.data() + old, static_cast<std::streamsize>(want));
+  pending_.resize(old + static_cast<size_t>(std::max<std::streamsize>(got, 0)));
   if (src_->bad()) return Status::Error("pxf: DatasetReader: istream read error");
-  // `fail` without `eof` is unusual; treat as a read error.
-  if (src_->fail() && got == 0) return Status::Error("pxf: DatasetReader: istream read failed");
+  if (got <= 0) return Status::Error("pxf: DatasetReader: istream read failed");
+  // A stream that never ends a row cannot grow the buffer without bound.
+  const int max_size = opts_.max_message_size > 0 ? opts_.max_message_size : kMaxMessageSize;
+  if (pending_.size() > static_cast<size_t>(max_size)) {
+    return Status::Error(
+        "dataset input holds " + std::to_string(pending_.size()) +
+        " bytes without a row boundary, MaxMessageSize=" + std::to_string(max_size));
+  }
   return Status::OK();
 }
 
-// Complete reports whether a token ending at pending_[tok_end] is known
-// complete: the source has ended, or a few bytes follow it. The lexer looks
-// at most one byte past a token to end it (the `.` of `1.5`, the `p` of a
-// precision, the second byte of `µs`), so bytes that arrive later cannot
-// change it.
-bool DatasetReader::Complete(int tok_end) const {
-  constexpr int kLookahead = 4;
-  return src_eof_ || tok_end + kLookahead <= static_cast<int>(pending_.size());
+// Complete reports whether tok, ending at pending_[tok_end], is known
+// complete: bytes that arrive later cannot change it. The source has ended;
+// or it is a one-byte token (punctuation, a newline); or the lexer has the
+// bytes it looks at past the token to end it. A closed string, bytes literal
+// or block comment needs none, but for `""`. Otherwise it is the byte at tok_end,
+// and the one after when that byte is `.` (`1.5`), `p` (a precision), `/`
+// (a comment after a timestamp) or the first byte of `µ` (`µs`). An error
+// token, or one that starts with a byte past ASCII, may have been decided
+// from a few bytes further, so it waits for four. Being exact here is what
+// lets a complete row through on a stream that stays open (#171).
+bool DatasetReader::Complete(const Token& tok, int tok_end) const {
+  if (src_eof_) return true;
+  const int n = static_cast<int>(pending_.size());
+  switch (tok.kind) {
+    case TokenKind::kNewline:
+    case TokenKind::kLBrace:
+    case TokenKind::kRBrace:
+    case TokenKind::kLBracket:
+    case TokenKind::kRBracket:
+    case TokenKind::kLParen:
+    case TokenKind::kRParen:
+    case TokenKind::kEquals:
+    case TokenKind::kColon:
+    case TokenKind::kComma:
+      return true;
+    case TokenKind::kIllegal:
+      return tok_end + 4 <= n;
+    default:
+      break;
+  }
+  if (static_cast<unsigned char>(pending_[tok.pos.offset]) >= 0x80) return tok_end + 4 <= n;
+  if (tok_end >= n) {
+    // A closed string, bytes literal or block comment ends at its own last
+    // byte, whatever follows: all but `""`, which a third quote would make
+    // the opening of a triple-quoted string.
+    const std::string_view text =
+        std::string_view(pending_).substr(tok.pos.offset, tok_end - tok.pos.offset);
+    if (tok.kind == TokenKind::kString || tok.kind == TokenKind::kBytes) return text != "\"\"";
+    return tok.kind == TokenKind::kComment && text.substr(0, 2) == "/*";
+  }
+  const unsigned char next = static_cast<unsigned char>(pending_[tok_end]);
+  if (next == '.' || next == 'p' || next == '/' || next == 0xC2) return tok_end + 1 < n;
+  return true;
 }
 
 // Wait records the token at pending_[at], which the scan could not yet take
@@ -619,7 +662,7 @@ DatasetReader::RowScan DatasetReader::ScanHeader(int* end) {
       const int tok_end = lx.CurrentPos().offset;
       scanned_ += static_cast<size_t>(tok_end - from);
       if (tok.kind == TokenKind::kEOF && src_eof_) return RowScan::kNeedMore;
-      if (tok.kind == TokenKind::kEOF || !Complete(tok_end)) {
+      if (tok.kind == TokenKind::kEOF || !Complete(tok, tok_end)) {
         Wait(tok.pos.offset);
         return RowScan::kNeedMore;
       }

test/pxf_dataset_reader_test.cc

@@ -13,6 +13,7 @@
 #include <fstream>
 #include <memory>
 #include <sstream>
+#include <streambuf>
 #include <string>
 
 #include <google/protobuf/compiler/importer.h>
@@ -664,4 +665,97 @@ TEST(DatasetReaderScan, HeaderIsLinear) {
   }
 }
 
+// ---- A stream that stays open (#171) ----------------------------------------
+
+// OneByteBuf serves its bytes one per read, as a pipe fed a byte at a time
+// does. Asked for a byte past them, it records it: on a live stream that read
+// would block until more bytes came, or forever.
+class OneByteBuf : public std::streambuf {
+ public:
+  explicit OneByteBuf(std::string data) : data_(std::move(data)) {}
+  bool asked_past_end() const { return asked_past_end_; }
+
+ protected:
+  int_type underflow() override {
+    if (gptr() != nullptr && gptr() < egptr()) return traits_type::to_int_type(*gptr());
+    if (next_ >= data_.size()) {
+      asked_past_end_ = true;
+      return traits_type::eof();
+    }
+    char* at = &data_[next_++];
+    setg(at, at, at + 1);
+    return traits_type::to_int_type(*at);
+  }
+
+ private:
+  std::string data_;
+  size_t next_ = 0;
+  bool asked_past_end_ = false;
+};
+
+// A row whose `)` is the last byte read is returned at once: the reader must
+// not wait for bytes that may never come (a live, append-only stream). The
+// rows are protowire-go's (TestDatasetReader_CompleteRowIsNotHeldBack); each
+// arrives a byte per read, so every byte of it ends a read.
+TEST(DatasetReaderLive, CompleteRowIsNotHeldBack) {
+  const char* rows[] = {
+      "(1)",
+      "(\"\")",
+      R"(("a\"b"))",
+      "(abc)",
+      "(2026-10-23T00:00:00Z)",
+      "(5s)",
+      "(-1.5e3)",
+      "(b\"QQ==\")",
+      "({ k = [1, 2] })",
+      "(\"\"\"\nx\n\"\"\")",
+      "(1 /* c */)",
+      "(1h.5m)",
+      "(1\xc2\xb5s)",
+      "(-inf)",
+      R"(("\x41\U0001F600"))",
+      "(aaaaaaaaaaaa)",
+  };
+  for (const char* row : rows) {
+    OneByteBuf buf(std::string("@dataset T (a)\n") + row);
+    std::istream src(&buf);
+    auto tr = DatasetReader::Create(&src);
+    ASSERT_TRUE(tr.ok()) << row << ": " << tr.status().ToString();
+    DatasetRow got;
+    ASSERT_TRUE((*tr)->Next(&got).ok()) << row << ": " << (*tr)->StickyStatus().ToString();
+    EXPECT_FALSE((*tr)->Done()) << row;
+    EXPECT_FALSE(buf.asked_past_end()) << row << ": the complete row waited for more input";
+  }
+}
+
+// The same scans stay linear when every read returns one byte, as protowire-go
+// holds them with iotest.OneByteReader.
+TEST(DatasetReaderScan, IsLinearOneByteAtATime) {
+  const std::pair<const char*, std::string> shapes[] = {
+      {"short tokens",
+       [] {
+         std::string s;
+         for (int i = 0; i < (128 << 10); ++i) s += "1 ";
+         return s;
+       }()},
+      {"unterminated string", "\"" + std::string(256 << 10, 'a')},
+      {"line comment", "1 # " + std::string(256 << 10, 'x')},
+      {"one long number", std::string(256 << 10, '7')},
+  };
+  for (const auto& [name, body] : shapes) {
+    const std::string in = "@dataset T (a)\n(" + body;
+    OneByteBuf buf(in);
+    std::istream src(&buf);
+    auto tr = DatasetReader::Create(&src);
+    ASSERT_TRUE(tr.ok()) << name << ": " << tr.status().ToString();
+    size_t& scanned = protowire::pxf::DatasetReaderTestPeer::Scanned(**tr);
+    scanned = 0;
+    DatasetRow row;
+    EXPECT_FALSE((*tr)->Next(&row).ok()) << name << ": an unclosed row was accepted";
+    // Short tokens in one-byte reads come to about 1.5x.
+    EXPECT_GT(scanned, 0u) << name;
+    EXPECT_LE(scanned, 2 * in.size()) << name;
+  }
+}
+
 }  // namespace