ytsaurus_job/
worker.rs

1//! One worker API selected by [`ytsaurus_format::DataFormat`].
2//!
3//! [`WorkerReader`] and [`WorkerWriter`] choose YSON or Skiff at the process
4//! boundary. They intentionally preserve each format's row representation:
5//! a YSON row is byte-exact and can be forwarded without decoding, while a
6//! Skiff row is a validated dynamic [`ytsaurus_skiff::Value`]. Treating those
7//! as one invented row type would discard the most useful property of both.
8
9use std::io::{Read, Write};
10
11use ytsaurus_format::DataFormat;
12use ytsaurus_skiff::Value;
13
14use crate::{
15    Event, JobError, JobReader, JobWriter, Result, SkiffJobReader, SkiffJobWriter, SkiffRow,
16    TableId,
17};
18
19/// A reader selected by the operation's input [`DataFormat`].
20#[derive(Debug)]
21pub enum WorkerReader<R> {
22    /// A YSON job stream.
23    Yson(JobReader<R>),
24    /// A schema-described Skiff job stream.
25    Skiff(SkiffJobReader<R>),
26}
27
28/// An input event returned by [`WorkerReader::next_event`].
29#[derive(Debug)]
30pub enum WorkerEvent<'input> {
31    /// A YSON data row or control event.
32    Yson(Event<'input>),
33    /// A Skiff data row, including its current control values.
34    Skiff(SkiffRow),
35}
36
37impl WorkerReader<std::io::BufReader<std::io::Stdin>> {
38    /// Reads stdin using the operation's selected input format.
39    ///
40    /// Buffered, because this constructor may end up on the Skiff side, which
41    /// reads field by field — see `skiff::STDIN_BUFFER_BYTES`. The YSON side
42    /// asks for a whole buffer at a time and is handed the descriptor's bytes
43    /// without a second copy.
44    ///
45    /// # Errors
46    ///
47    /// Returns an error if the format is unknown to this runtime or a Skiff
48    /// schema is not suitable for job input.
49    pub fn from_stdin(format: DataFormat) -> Result<Self> {
50        Self::new(
51            std::io::BufReader::with_capacity(crate::skiff::STDIN_BUFFER_BYTES, std::io::stdin()),
52            format,
53        )
54    }
55}
56
57impl<R: Read> WorkerReader<R> {
58    /// Creates a worker input reader selected by `format`.
59    ///
60    /// # Errors
61    ///
62    /// Returns an error if the format is unknown to this runtime or a Skiff
63    /// schema is not suitable for job input.
64    pub fn new(input: R, format: DataFormat) -> Result<Self> {
65        match format {
66            DataFormat::Yson(format) => Ok(Self::Yson(JobReader::with_format(input, format))),
67            DataFormat::Skiff(format) => Ok(Self::Skiff(SkiffJobReader::new(input, format)?)),
68            _ => Err(JobError::UnsupportedDataFormat),
69        }
70    }
71
72    /// Returns the next row or control event, or `None` at clean end of input.
73    ///
74    /// A returned YSON event borrows the reader's buffer; process it before
75    /// calling this method again. Skiff rows are owned.
76    pub fn next_event(&mut self) -> Result<Option<WorkerEvent<'_>>> {
77        match self {
78            Self::Yson(reader) => reader
79                .next_event()
80                .map(|event| event.map(WorkerEvent::Yson)),
81            Self::Skiff(reader) => reader.next_row().map(|row| row.map(WorkerEvent::Skiff)),
82        }
83    }
84}
85
86/// A row accepted by [`WorkerWriter::write`].
87pub enum WorkerRow<'row> {
88    /// One complete, already-encoded YSON value. The writer adds its `;`
89    /// record separator, just as [`JobWriter::write_raw`] does.
90    YsonRaw(&'row [u8]),
91    /// One value matching the selected Skiff table schema.
92    Skiff(&'row Value),
93}
94
95/// A writer selected by the operation's output [`DataFormat`].
96///
97/// Existing [`JobWriter`] and [`SkiffJobWriter`] APIs remain available for
98/// callers that want serde-based YSON output or direct format-specific access.
99pub enum WorkerWriter {
100    /// A YSON output stream.
101    Yson(JobWriter),
102    /// Schema-described Skiff output streams.
103    Skiff(SkiffJobWriter),
104}
105
106impl std::fmt::Debug for WorkerWriter {
107    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
108        match self {
109            Self::Yson(writer) => formatter
110                .debug_tuple("WorkerWriter::Yson")
111                .field(writer)
112                .finish(),
113            Self::Skiff(writer) => formatter
114                .debug_tuple("WorkerWriter::Skiff")
115                .field(writer)
116                .finish(),
117        }
118    }
119}
120
121impl WorkerWriter {
122    /// Opens YTsaurus output descriptors for `format`.
123    ///
124    /// For Skiff, `table_count` must equal the number of table schemas in the
125    /// format. One physical descriptor is opened per output table.
126    ///
127    /// # Errors
128    ///
129    /// Returns an error if the format is unknown, its Skiff schemas are not
130    /// valid job-output schemas, or its schema count differs from `table_count`.
131    #[cfg(unix)]
132    pub fn descriptors(format: DataFormat, table_count: usize) -> Result<Self> {
133        match format {
134            DataFormat::Yson(format) => {
135                JobWriter::descriptors_with_format(table_count, format).map(Self::Yson)
136            }
137            DataFormat::Skiff(format) => {
138                let schemas = format.table_schemas().len();
139                if schemas != table_count {
140                    return Err(JobError::SkiffOutputSchemaCount {
141                        sinks: table_count,
142                        schemas,
143                    });
144                }
145                SkiffJobWriter::descriptors(format).map(Self::Skiff)
146            }
147            _ => Err(JobError::UnsupportedDataFormat),
148        }
149    }
150
151    /// Builds a writer over arbitrary sinks, primarily for offline tests.
152    ///
153    /// For Skiff, the number of supplied sinks must equal the format's table
154    /// schemas. For YSON it determines the number of output tables.
155    ///
156    /// # Errors
157    ///
158    /// Returns an error if the format is unknown or its Skiff output schemas
159    /// are unsuitable.
160    pub fn from_writers(tables: Vec<Box<dyn Write>>, format: DataFormat) -> Result<Self> {
161        match format {
162            DataFormat::Yson(format) => Ok(Self::Yson(JobWriter::from_writers(tables, format))),
163            DataFormat::Skiff(format) => {
164                SkiffJobWriter::from_writers(tables, format).map(Self::Skiff)
165            }
166            _ => Err(JobError::UnsupportedDataFormat),
167        }
168    }
169
170    /// Number of output tables this writer can address.
171    #[must_use]
172    pub fn table_count(&self) -> usize {
173        match self {
174            Self::Yson(writer) => writer.table_count(),
175            Self::Skiff(writer) => writer.table_count(),
176        }
177    }
178
179    /// Writes one row to an output table.
180    ///
181    /// `WorkerRow` must use the representation selected by this writer's
182    /// [`DataFormat`].
183    ///
184    /// # Errors
185    ///
186    /// Returns an error for a format mismatch, invalid row, unknown table, or
187    /// failed output write.
188    pub fn write(&mut self, table: impl Into<TableId>, row: WorkerRow<'_>) -> Result<()> {
189        match (self, row) {
190            (Self::Yson(writer), WorkerRow::YsonRaw(row)) => writer.write_raw(table, row),
191            (Self::Skiff(writer), WorkerRow::Skiff(row)) => writer.write(table, row),
192            (Self::Yson(_), WorkerRow::Skiff(_)) => Err(JobError::WorkerRowFormatMismatch {
193                writer: "YSON",
194                row: "Skiff",
195            }),
196            (Self::Skiff(_), WorkerRow::YsonRaw(_)) => Err(JobError::WorkerRowFormatMismatch {
197                writer: "Skiff",
198                row: "YSON",
199            }),
200        }
201    }
202
203    /// Flushes every output table.
204    pub fn flush(&mut self) -> Result<()> {
205        match self {
206            Self::Yson(writer) => writer.flush(),
207            Self::Skiff(writer) => writer.flush(),
208        }
209    }
210
211    /// Flushes every output table and marks the writer complete.
212    pub fn finish(&mut self) -> Result<()> {
213        match self {
214            Self::Yson(writer) => writer.finish(),
215            Self::Skiff(writer) => writer.finish(),
216        }
217    }
218}
219
220#[cfg(test)]
221mod tests {
222    use std::io::Cursor;
223
224    use ytsaurus_format::SkiffFormat;
225    use ytsaurus_skiff::{Encoder, Schema, SchemaRef, Value, WireType};
226
227    use super::*;
228
229    fn skiff_format() -> SkiffFormat {
230        SkiffFormat::new(vec![SchemaRef::Inline(Schema::tuple([Schema::named(
231            "value",
232            WireType::String32,
233        )]))])
234        .unwrap()
235    }
236
237    #[test]
238    fn reader_selects_yson_and_skiff_from_the_same_enum() {
239        let mut yson =
240            WorkerReader::new(Cursor::new(b"{value=one};"), DataFormat::text_yson()).unwrap();
241        assert!(matches!(
242            yson.next_event().unwrap(),
243            Some(WorkerEvent::Yson(_))
244        ));
245
246        let schema = skiff_format().table_schema(0).unwrap().clone();
247        let mut encoder = Encoder::new(Vec::new(), schema).unwrap();
248        encoder
249            .write(&Value::Tuple(vec![Value::Bytes(b"one".to_vec())]))
250            .unwrap();
251        let stream = encoder.into_inner().unwrap();
252        let mut skiff =
253            WorkerReader::new(Cursor::new(stream), DataFormat::skiff(skiff_format())).unwrap();
254        assert!(matches!(
255            skiff.next_event().unwrap(),
256            Some(WorkerEvent::Skiff(_))
257        ));
258    }
259
260    #[test]
261    fn writer_rejects_a_row_from_the_other_format() {
262        let mut writer =
263            WorkerWriter::from_writers(vec![Box::new(Vec::new())], DataFormat::binary_yson())
264                .unwrap();
265        let error = writer
266            .write(0, WorkerRow::Skiff(&Value::Tuple(Vec::new())))
267            .unwrap_err();
268        assert_eq!(error.kind(), "worker_row_format_mismatch");
269    }
270}
ytsaurus_job/worker.rs

ytsaurus_job/
worker.rs
