mirror of
https://git.pumpin.online/funkpower/catena.git
synced 2026-10-23 03:07:13 +00:00
workspace, licence, and the three record types
This commit is contained in:
@@ -0,0 +1,18 @@
|
||||
[package]
|
||||
name = "corpus-core"
|
||||
description = "Record types and stream plumbing shared by the catena tools."
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
schemars.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
|
||||
[lints.rust]
|
||||
unsafe_code = "forbid"
|
||||
missing_docs = "warn"
|
||||
@@ -0,0 +1,135 @@
|
||||
//! A retrievable passage, with enough provenance to cite it.
|
||||
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// One passage of a document or a recording.
|
||||
///
|
||||
/// Documents fill `page` and `chapter`; media fill `start_time` and
|
||||
/// `end_time`. The shape is the same either way, so one index serves a book
|
||||
/// and a lecture recording and can cite both.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct Chunk {
|
||||
/// Record version.
|
||||
#[serde(default = "crate::v")]
|
||||
pub v: u32,
|
||||
/// Position in the stream this chunk was cut from, from zero.
|
||||
pub idx: u32,
|
||||
/// The passage.
|
||||
pub text: String,
|
||||
/// Token count of `text`, by the tokeniser the chunker used.
|
||||
pub tokens: u32,
|
||||
/// Page the passage starts on, for paginated sources.
|
||||
pub page: Option<u32>,
|
||||
/// Nearest enclosing heading, as written.
|
||||
pub chapter: Option<String>,
|
||||
/// What part of the work the passage belongs to.
|
||||
pub role: Role,
|
||||
/// Seconds from the start of the recording, for timed sources.
|
||||
pub start_time: Option<f64>,
|
||||
/// Seconds from the start of the recording, for timed sources.
|
||||
pub end_time: Option<f64>,
|
||||
/// Set tags, sorted and deduplicated; see [`flags`].
|
||||
#[serde(default)]
|
||||
pub flags: Vec<String>,
|
||||
}
|
||||
|
||||
impl Chunk {
|
||||
/// A body chunk with no provenance yet.
|
||||
pub fn new(idx: u32, text: impl Into<String>, tokens: u32) -> Self {
|
||||
Self {
|
||||
v: crate::V,
|
||||
idx,
|
||||
text: text.into(),
|
||||
tokens,
|
||||
page: None,
|
||||
chapter: None,
|
||||
role: Role::Body,
|
||||
start_time: None,
|
||||
end_time: None,
|
||||
flags: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Set a flag. Flags stay sorted and unique so the same input always
|
||||
/// serialises to the same bytes.
|
||||
pub fn flag(&mut self, flag: impl Into<String>) {
|
||||
let flag = flag.into();
|
||||
if let Err(at) = self.flags.binary_search(&flag) {
|
||||
self.flags.insert(at, flag);
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether a flag is set.
|
||||
pub fn has_flag(&self, flag: &str) -> bool {
|
||||
self.flags
|
||||
.binary_search_by(|f| f.as_str().cmp(flag))
|
||||
.is_ok()
|
||||
}
|
||||
}
|
||||
|
||||
/// The flag names the tools in this workspace set.
|
||||
pub mod flags {
|
||||
/// The passage's footnote markers were resolved against a citation list.
|
||||
pub const FOOTNOTE_CITATIONS: &str = "footnote_citations";
|
||||
}
|
||||
|
||||
/// What part of a work a passage belongs to.
|
||||
///
|
||||
/// Retrieval wants body text; an index or a bibliography matches lexically
|
||||
/// far too well for what it is worth, which is why the distinction is carried
|
||||
/// on the record rather than rediscovered at query time.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum Role {
|
||||
/// Running text.
|
||||
#[default]
|
||||
Body,
|
||||
/// A works-cited or references list.
|
||||
Bibliography,
|
||||
/// A back-of-book index.
|
||||
Index,
|
||||
/// A table of contents.
|
||||
Toc,
|
||||
/// Title page, copyright page, dedication, and the rest of the front.
|
||||
Frontmatter,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn role_spellings_are_the_wire_spellings() {
|
||||
for (role, wire) in [
|
||||
(Role::Body, "\"body\""),
|
||||
(Role::Bibliography, "\"bibliography\""),
|
||||
(Role::Index, "\"index\""),
|
||||
(Role::Toc, "\"toc\""),
|
||||
(Role::Frontmatter, "\"frontmatter\""),
|
||||
] {
|
||||
assert_eq!(serde_json::to_string(&role).unwrap(), wire);
|
||||
assert_eq!(serde_json::from_str::<Role>(wire).unwrap(), role);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_field_is_emitted_even_when_empty() {
|
||||
let chunk = Chunk::new(0, "hello", 1);
|
||||
assert_eq!(
|
||||
serde_json::to_string(&chunk).unwrap(),
|
||||
r#"{"v":1,"idx":0,"text":"hello","tokens":1,"page":null,"chapter":null,"role":"body","start_time":null,"end_time":null,"flags":[]}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flags_stay_sorted_and_unique() {
|
||||
let mut chunk = Chunk::new(0, "x", 1);
|
||||
chunk.flag("zebra");
|
||||
chunk.flag(flags::FOOTNOTE_CITATIONS);
|
||||
chunk.flag("zebra");
|
||||
assert_eq!(chunk.flags, ["footnote_citations", "zebra"]);
|
||||
assert!(chunk.has_flag("zebra"));
|
||||
assert!(!chunk.has_flag("aardvark"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,144 @@
|
||||
//! What a file is.
|
||||
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// A source file, as identified by inspection rather than by its name.
|
||||
///
|
||||
/// Fields a given tool cannot know are `None`: a format sniffer fills `kind`
|
||||
/// and leaves `text_layer` empty, a triage tool does the reverse.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct Doc {
|
||||
/// Record version.
|
||||
#[serde(default = "crate::v")]
|
||||
pub v: u32,
|
||||
/// Path the file was read from, as given.
|
||||
pub path: String,
|
||||
/// What the bytes say the file is.
|
||||
pub kind: Kind,
|
||||
/// Extension the name claimed, lowercased, without the dot.
|
||||
pub declared_ext: Option<String>,
|
||||
/// Size on disk.
|
||||
pub bytes: u64,
|
||||
/// Page count, for paginated formats.
|
||||
pub pages: Option<u32>,
|
||||
/// ISO 639-2/B language code, when one was detected.
|
||||
pub language: Option<String>,
|
||||
/// Verdict on the embedded text layer.
|
||||
pub text_layer: Option<TextLayer>,
|
||||
/// Whether the file has to be OCRed to be searchable.
|
||||
pub needs_ocr: Option<bool>,
|
||||
/// Human-readable observations, in the order they were made.
|
||||
pub notes: Vec<String>,
|
||||
}
|
||||
|
||||
impl Doc {
|
||||
/// A doc with only what inspection of the bytes always yields.
|
||||
pub fn new(path: impl Into<String>, kind: Kind, bytes: u64) -> Self {
|
||||
Self {
|
||||
v: crate::V,
|
||||
path: path.into(),
|
||||
kind,
|
||||
declared_ext: None,
|
||||
bytes,
|
||||
pages: None,
|
||||
language: None,
|
||||
text_layer: None,
|
||||
needs_ocr: None,
|
||||
notes: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Append an observation.
|
||||
pub fn note(&mut self, note: impl Into<String>) {
|
||||
self.notes.push(note.into());
|
||||
}
|
||||
}
|
||||
|
||||
/// The format a file actually is.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum Kind {
|
||||
/// Portable Document Format.
|
||||
Pdf,
|
||||
/// EPUB, any version.
|
||||
Epub,
|
||||
/// DjVu.
|
||||
Djvu,
|
||||
/// Amazon KF8.
|
||||
Azw3,
|
||||
/// Amazon MOBI.
|
||||
Mobi,
|
||||
/// Plain text.
|
||||
Txt,
|
||||
/// HTML or XHTML.
|
||||
Html,
|
||||
/// An OLE2 compound document — a `.doc`, whatever the name says.
|
||||
Ole2,
|
||||
/// A zip archive whose only real content is a single PDF.
|
||||
ZipWrappedPdf,
|
||||
/// Recognised as none of the above.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
/// The state of a PDF's embedded text layer, and therefore whether OCR is
|
||||
/// worth its cost.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum TextLayer {
|
||||
/// Extractable and readable as it stands.
|
||||
Healthy,
|
||||
/// Present but mojibake — extraction yields characters, not words.
|
||||
Garbled,
|
||||
/// Present but too thin to be the whole page; a scan with a caption layer.
|
||||
Sparse,
|
||||
/// Present and permuted by a non-standard font encoding, recoverable
|
||||
/// without OCR once the permutation is undone.
|
||||
ReflectedFont,
|
||||
/// No text at all; a scan.
|
||||
Absent,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn kind_spellings_are_the_wire_spellings() {
|
||||
let cases = [
|
||||
(Kind::Pdf, "\"pdf\""),
|
||||
(Kind::Ole2, "\"ole2\""),
|
||||
(Kind::Azw3, "\"azw3\""),
|
||||
(Kind::ZipWrappedPdf, "\"zip-wrapped-pdf\""),
|
||||
(Kind::Unknown, "\"unknown\""),
|
||||
];
|
||||
for (kind, wire) in cases {
|
||||
assert_eq!(serde_json::to_string(&kind).unwrap(), wire);
|
||||
assert_eq!(serde_json::from_str::<Kind>(wire).unwrap(), kind);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn text_layer_spellings_are_the_wire_spellings() {
|
||||
assert_eq!(
|
||||
serde_json::to_string(&TextLayer::ReflectedFont).unwrap(),
|
||||
"\"reflected-font\""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_field_is_emitted_even_when_empty() {
|
||||
let doc = Doc::new("/tmp/a.pdf", Kind::Pdf, 12);
|
||||
assert_eq!(
|
||||
serde_json::to_string(&doc).unwrap(),
|
||||
r#"{"v":1,"path":"/tmp/a.pdf","kind":"pdf","declared_ext":null,"bytes":12,"pages":null,"language":null,"text_layer":null,"needs_ocr":null,"notes":[]}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_record_written_without_a_version_reads_as_version_one() {
|
||||
let doc: Doc =
|
||||
serde_json::from_str(r#"{"path":"a","kind":"txt","bytes":1,"notes":[]}"#).unwrap();
|
||||
assert_eq!(doc.v, 1);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,69 @@
|
||||
//! What a tool's exit code means.
|
||||
//!
|
||||
//! A predicate answers through its status, so it is usable in a shell `if`
|
||||
//! without parsing anything. That only works if "no" and "broke" are
|
||||
//! different numbers, which is why the failure codes below start where
|
||||
//! `sysexits.h` puts them and leave 0, 1 and 2 to the answer.
|
||||
|
||||
/// Success, or a predicate answering yes.
|
||||
pub const OK: i32 = 0;
|
||||
/// A predicate answering no.
|
||||
pub const NO: i32 = 1;
|
||||
/// A predicate that could not tell — a damaged file, a format it does not
|
||||
/// know. Distinct from no, because a caller may want to retry or escalate.
|
||||
pub const INDETERMINATE: i32 = 2;
|
||||
/// The command line was wrong.
|
||||
pub const USAGE: i32 = 64;
|
||||
/// The input was not what the tool can work on.
|
||||
pub const DATA: i32 = 65;
|
||||
/// An external program the tool needs is not on `PATH`.
|
||||
pub const UNAVAILABLE: i32 = 69;
|
||||
/// Reading or writing failed.
|
||||
pub const IO: i32 = 74;
|
||||
|
||||
/// The answer a predicate tool gives.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Verdict {
|
||||
/// Yes.
|
||||
Yes,
|
||||
/// No.
|
||||
No,
|
||||
/// Could not tell.
|
||||
Unknown,
|
||||
}
|
||||
|
||||
impl Verdict {
|
||||
/// The process exit code for this answer.
|
||||
pub const fn code(self) -> i32 {
|
||||
match self {
|
||||
Verdict::Yes => OK,
|
||||
Verdict::No => NO,
|
||||
Verdict::Unknown => INDETERMINATE,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<bool> for Verdict {
|
||||
fn from(yes: bool) -> Self {
|
||||
if yes { Verdict::Yes } else { Verdict::No }
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn the_three_answers_are_zero_one_two() {
|
||||
assert_eq!(Verdict::Yes.code(), 0);
|
||||
assert_eq!(Verdict::No.code(), 1);
|
||||
assert_eq!(Verdict::Unknown.code(), 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn failure_codes_never_collide_with_an_answer() {
|
||||
for code in [USAGE, DATA, UNAVAILABLE, IO] {
|
||||
assert!(code > INDETERMINATE);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
//! One scored candidate, on its way into a fusion.
|
||||
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// A single result from a single retrieval channel.
|
||||
///
|
||||
/// `rank` is the candidate's position within its own channel, from one, and
|
||||
/// is what rank fusion consumes; `score` is the channel's raw number, carried
|
||||
/// along because it is useful to a human reading a trace and useless for
|
||||
/// comparing across channels.
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)]
|
||||
pub struct Hit {
|
||||
/// Record version.
|
||||
#[serde(default = "crate::v")]
|
||||
pub v: u32,
|
||||
/// Identifier of the thing retrieved, opaque to fusion.
|
||||
pub id: String,
|
||||
/// Which channel produced this hit.
|
||||
pub channel: Channel,
|
||||
/// Position within the channel's own ranking, from one.
|
||||
pub rank: u32,
|
||||
/// The channel's raw score, on whatever scale that channel uses.
|
||||
pub score: f64,
|
||||
/// Anything the producer wants to carry through fusion untouched.
|
||||
#[serde(default)]
|
||||
pub payload: Map<String, Value>,
|
||||
}
|
||||
|
||||
impl Hit {
|
||||
/// A hit with an empty payload.
|
||||
pub fn new(id: impl Into<String>, channel: Channel, rank: u32, score: f64) -> Self {
|
||||
Self {
|
||||
v: crate::V,
|
||||
id: id.into(),
|
||||
channel,
|
||||
rank,
|
||||
score,
|
||||
payload: Map::new(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A retrieval channel.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, JsonSchema)]
|
||||
#[serde(rename_all = "kebab-case")]
|
||||
pub enum Channel {
|
||||
/// Nearest neighbours in an embedding space.
|
||||
#[serde(rename = "vec")]
|
||||
Vector,
|
||||
/// Full-text search.
|
||||
Fts,
|
||||
/// A match on metadata — title, author, tag.
|
||||
Meta,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn channel_spellings_are_the_wire_spellings() {
|
||||
for (channel, wire) in [
|
||||
(Channel::Vector, "\"vec\""),
|
||||
(Channel::Fts, "\"fts\""),
|
||||
(Channel::Meta, "\"meta\""),
|
||||
] {
|
||||
assert_eq!(serde_json::to_string(&channel).unwrap(), wire);
|
||||
assert_eq!(serde_json::from_str::<Channel>(wire).unwrap(), channel);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_field_is_emitted_even_when_empty() {
|
||||
let hit = Hit::new("abc", Channel::Fts, 1, 0.5);
|
||||
assert_eq!(
|
||||
serde_json::to_string(&hit).unwrap(),
|
||||
r#"{"v":1,"id":"abc","channel":"fts","rank":1,"score":0.5,"payload":{}}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn payload_keys_serialise_in_a_fixed_order() {
|
||||
let mut hit = Hit::new("abc", Channel::Vector, 2, 0.25);
|
||||
hit.payload.insert("zulu".into(), Value::from(1));
|
||||
hit.payload.insert("alpha".into(), Value::from(2));
|
||||
assert_eq!(
|
||||
serde_json::to_string(&hit).unwrap(),
|
||||
r#"{"v":1,"id":"abc","channel":"vec","rank":2,"score":0.25,"payload":{"alpha":2,"zulu":1}}"#
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,29 @@
|
||||
//! Record types and stream plumbing shared by every catena tool.
|
||||
//!
|
||||
//! Three records travel between the tools: a [`Doc`] is what a file is, a
|
||||
//! [`Chunk`] is a retrievable passage with enough provenance to cite it, and a
|
||||
//! [`Hit`] is one scored candidate on its way into a fusion. Each carries a
|
||||
//! version field so a reader can refuse a stream it does not understand.
|
||||
//!
|
||||
//! Every list-shaped stream is NDJSON — see [`ndjson`] — and every predicate
|
||||
//! answers through an exit code — see [`exit`].
|
||||
|
||||
pub mod chunk;
|
||||
pub mod doc;
|
||||
pub mod exit;
|
||||
pub mod hit;
|
||||
pub mod ndjson;
|
||||
|
||||
pub use chunk::{Chunk, Role};
|
||||
pub use doc::{Doc, Kind, TextLayer};
|
||||
pub use hit::{Channel, Hit};
|
||||
pub use ndjson::{Error, read, write_line};
|
||||
|
||||
/// The record version these types serialise as.
|
||||
pub const V: u32 = 1;
|
||||
|
||||
/// Serde default for the `v` field, so a record written by an older tool that
|
||||
/// omitted it still reads back as version 1 rather than 0.
|
||||
pub(crate) const fn v() -> u32 {
|
||||
V
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
//! Prints the JSON Schema for the record types, so a consumer in another
|
||||
//! language can generate its own bindings instead of reading Rust.
|
||||
|
||||
use std::process::ExitCode;
|
||||
|
||||
use corpus_core::{Chunk, Doc, Hit, exit};
|
||||
|
||||
const USAGE: &str = "\
|
||||
usage: corpus-core schema [doc|chunk|hit|all]
|
||||
|
||||
Prints the JSON Schema for one record type, or an object of all three.
|
||||
|
||||
cargo run -p corpus-core -- schema chunk > chunk.schema.json
|
||||
";
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let mut args = std::env::args().skip(1);
|
||||
match args.next().as_deref() {
|
||||
Some("schema") => {
|
||||
let which = args.next().unwrap_or_else(|| "all".into());
|
||||
match schema(&which) {
|
||||
Some(json) => {
|
||||
println!("{json}");
|
||||
ExitCode::from(exit::OK as u8)
|
||||
}
|
||||
None => {
|
||||
eprintln!("corpus-core: unknown record type {which:?}\n\n{USAGE}");
|
||||
ExitCode::from(exit::USAGE as u8)
|
||||
}
|
||||
}
|
||||
}
|
||||
Some("--help" | "-h") => {
|
||||
println!("{USAGE}");
|
||||
ExitCode::from(exit::OK as u8)
|
||||
}
|
||||
_ => {
|
||||
eprintln!("{USAGE}");
|
||||
ExitCode::from(exit::USAGE as u8)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn schema(which: &str) -> Option<String> {
|
||||
let json = match which {
|
||||
"doc" => serde_json::to_string_pretty(&schemars::schema_for!(Doc)),
|
||||
"chunk" => serde_json::to_string_pretty(&schemars::schema_for!(Chunk)),
|
||||
"hit" => serde_json::to_string_pretty(&schemars::schema_for!(Hit)),
|
||||
"all" => serde_json::to_string_pretty(&serde_json::json!({
|
||||
"doc": schemars::schema_for!(Doc),
|
||||
"chunk": schemars::schema_for!(Chunk),
|
||||
"hit": schemars::schema_for!(Hit),
|
||||
})),
|
||||
_ => return None,
|
||||
};
|
||||
Some(json.expect("a schema is always serialisable"))
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
//! Newline-delimited JSON, read lazily and written flushed.
|
||||
//!
|
||||
//! Every list-shaped stream between these tools is NDJSON: one record per
|
||||
//! line, produced and consumed one at a time, so a slow stage never has to
|
||||
//! hold the whole list and never blocks the pipe behind it.
|
||||
|
||||
use std::io::{BufRead, Write};
|
||||
|
||||
use serde::{Serialize, de::DeserializeOwned};
|
||||
|
||||
/// Something went wrong reading or writing a stream.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
/// A line was not the record the reader expected.
|
||||
#[error("line {line}: {source}")]
|
||||
Parse {
|
||||
/// One-based line number within the stream, counting blank lines.
|
||||
line: u64,
|
||||
/// What serde made of it.
|
||||
source: serde_json::Error,
|
||||
},
|
||||
/// A record could not be turned into JSON at all.
|
||||
#[error("encoding a record: {0}")]
|
||||
Encode(#[source] serde_json::Error),
|
||||
/// The stream itself failed.
|
||||
#[error("{0}")]
|
||||
Io(#[from] std::io::Error),
|
||||
}
|
||||
|
||||
/// Read a stream of records, one per line.
|
||||
///
|
||||
/// Blank lines are skipped, so a stream that was concatenated or hand-edited
|
||||
/// still reads. The iterator stops at the first error rather than skipping
|
||||
/// it: a malformed line means the producer and the consumer disagree about
|
||||
/// the contract, and continuing would bury that.
|
||||
pub fn read<T, R>(reader: R) -> impl Iterator<Item = Result<T, Error>>
|
||||
where
|
||||
T: DeserializeOwned,
|
||||
R: BufRead,
|
||||
{
|
||||
let mut line_no = 0u64;
|
||||
reader.lines().filter_map(move |line| {
|
||||
line_no += 1;
|
||||
match line {
|
||||
Err(e) => Some(Err(Error::Io(e))),
|
||||
Ok(line) if line.trim().is_empty() => None,
|
||||
Ok(line) => Some(serde_json::from_str(&line).map_err(|source| Error::Parse {
|
||||
line: line_no,
|
||||
source,
|
||||
})),
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// Write one record and flush it.
|
||||
///
|
||||
/// The flush is the point: a consumer downstream should see a record as soon
|
||||
/// as it exists, not when the producer's buffer happens to fill.
|
||||
pub fn write_line<T, W>(writer: &mut W, record: &T) -> Result<(), Error>
|
||||
where
|
||||
T: Serialize,
|
||||
W: Write,
|
||||
{
|
||||
serde_json::to_writer(&mut *writer, record).map_err(Error::Encode)?;
|
||||
writer.write_all(b"\n")?;
|
||||
writer.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::{Channel, Hit};
|
||||
|
||||
#[test]
|
||||
fn reads_a_record_per_line_and_skips_blanks() {
|
||||
let input = "{\"id\":\"a\",\"channel\":\"fts\",\"rank\":1,\"score\":1.0}\n\n\
|
||||
{\"id\":\"b\",\"channel\":\"vec\",\"rank\":2,\"score\":0.5}\n";
|
||||
let hits: Vec<Hit> = read(input.as_bytes()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(hits.len(), 2);
|
||||
assert_eq!(hits[1].id, "b");
|
||||
assert_eq!(hits[1].channel, Channel::Vector);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_bad_line_reports_its_line_number_and_stops_the_stream() {
|
||||
let input = "{\"id\":\"a\",\"channel\":\"fts\",\"rank\":1,\"score\":1.0}\nnot json\n";
|
||||
let mut stream = read::<Hit, _>(input.as_bytes());
|
||||
assert!(stream.next().unwrap().is_ok());
|
||||
match stream.next().unwrap() {
|
||||
Err(Error::Parse { line, .. }) => assert_eq!(line, 2),
|
||||
other => panic!("expected a parse error, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn writes_one_compact_line_per_record() {
|
||||
let mut out = Vec::new();
|
||||
write_line(&mut out, &Hit::new("a", Channel::Meta, 1, 0.25)).unwrap();
|
||||
write_line(&mut out, &Hit::new("b", Channel::Meta, 2, 0.5)).unwrap();
|
||||
let text = String::from_utf8(out).unwrap();
|
||||
assert_eq!(text.lines().count(), 2);
|
||||
assert!(text.ends_with('\n'));
|
||||
assert!(!text.contains("\n "));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn what_is_written_reads_back_identical() {
|
||||
let hit = Hit::new("a", Channel::Vector, 3, 0.125);
|
||||
let mut out = Vec::new();
|
||||
write_line(&mut out, &hit).unwrap();
|
||||
let back: Vec<Hit> = read(out.as_slice()).collect::<Result<_, _>>().unwrap();
|
||||
assert_eq!(back, vec![hit]);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user