1use ytsaurus_format::DataFormat;
12use ytsaurus_skiff::Format as SkiffFormat;
13use ytsaurus_yson::YsonValue;
14
15use crate::yson_build::{boolean, insert, int, list, map, string, with_attributes};
16
17fn named_file(path: impl Into<String>, name: impl AsRef<str>) -> YsonValue {
19 with_attributes(string(path.into()), [("file_name", string(name.as_ref()))])
20}
21
22fn skiff_table_mismatch(
30 what: &str,
31 format: &DataFormat,
32 tables: usize,
33 kind: &str,
34) -> Option<String> {
35 let schemas = format.as_skiff()?.table_schemas().len();
36 if schemas == tables {
37 return None;
38 }
39 Some(format!(
40 "{what} declares {}, but this operation has {}",
41 plural(schemas, "Skiff table schema"),
42 plural(tables, kind)
43 ))
44}
45
46fn plural(count: usize, noun: &str) -> String {
47 if count == 1 {
48 format!("{count} {noun}")
49 } else {
50 format!("{count} {noun}s")
51 }
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
60pub enum OperationType {
61 Map,
63 MapReduce,
65 Reduce,
67 Sort,
69 Vanilla,
71 Merge,
73 Erase,
75 RemoteCopy,
77 JoinReduce,
99}
100
101impl OperationType {
102 #[must_use]
104 pub fn as_str(self) -> &'static str {
105 match self {
106 OperationType::Map => "map",
107 OperationType::MapReduce => "map_reduce",
108 OperationType::Reduce => "reduce",
109 OperationType::Sort => "sort",
110 OperationType::Vanilla => "vanilla",
111 OperationType::Merge => "merge",
112 OperationType::Erase => "erase",
113 OperationType::RemoteCopy => "remote_copy",
114 OperationType::JoinReduce => "join_reduce",
115 }
116 }
117}
118
119#[derive(Debug, Clone)]
121struct UserJob {
122 command: String,
123 files: Vec<YsonValue>,
128 memory_limit: Option<i64>,
129 environment: Vec<(String, String)>,
130 input_format: DataFormat,
131 output_format: DataFormat,
132}
133
134impl UserJob {
135 fn new(command: impl Into<String>) -> Self {
136 Self {
137 command: command.into(),
138 files: Vec::new(),
139 memory_limit: None,
140 environment: Vec::new(),
141 input_format: DataFormat::binary_yson(),
142 output_format: DataFormat::binary_yson(),
143 }
144 }
145
146 fn with_formats(&mut self, input: DataFormat, output: DataFormat) {
147 self.input_format = input;
148 self.output_format = output;
149 }
150
151 fn to_yson(&self) -> YsonValue {
152 let mut job = map([
153 ("command", string(&self.command)),
154 ("input_format", self.input_format.to_yson()),
157 ("output_format", self.output_format.to_yson()),
158 ]);
159
160 if !self.files.is_empty() {
161 insert(&mut job, "file_paths", list(self.files.iter().cloned()));
162 }
163 if let Some(limit) = self.memory_limit {
164 insert(&mut job, "memory_limit", int(limit));
165 }
166 if !self.environment.is_empty() {
167 insert(
168 &mut job,
169 "environment",
170 map(self
171 .environment
172 .iter()
173 .map(|(k, v)| (k.as_str(), string(v)))),
174 );
175 }
176 job
177 }
178}
179
180#[derive(Debug, Clone)]
190pub struct MapSpec {
191 mapper: UserJob,
192 inputs: Vec<String>,
193 outputs: Vec<String>,
194 job_count: Option<i64>,
195 input_table_index: bool,
196 extra: Vec<(String, YsonValue)>,
197}
198
199impl MapSpec {
200 #[must_use]
202 pub fn new<I, O>(command: impl Into<String>, inputs: I, outputs: O) -> Self
203 where
204 I: IntoIterator,
205 I::Item: Into<String>,
206 O: IntoIterator,
207 O::Item: Into<String>,
208 {
209 Self {
210 mapper: UserJob::new(command),
211 inputs: inputs.into_iter().map(Into::into).collect(),
212 outputs: outputs.into_iter().map(Into::into).collect(),
213 job_count: None,
214 input_table_index: false,
215 extra: Vec::new(),
216 }
217 }
218
219 #[must_use]
221 pub fn with_local_file(mut self, path: impl Into<String>) -> Self {
222 self.mapper.files.push(string(path.into()));
223 self
224 }
225
226 #[must_use]
231 pub fn with_local_file_named(mut self, path: impl Into<String>, name: impl AsRef<str>) -> Self {
232 self.mapper.files.push(named_file(path, name));
233 self
234 }
235
236 #[must_use]
238 pub fn with_memory_limit(mut self, bytes: i64) -> Self {
239 self.mapper.memory_limit = Some(bytes);
240 self
241 }
242
243 #[must_use]
249 pub fn with_formats(mut self, input: DataFormat, output: DataFormat) -> Self {
250 self.mapper.with_formats(input, output);
251 self
252 }
253
254 #[must_use]
258 pub fn with_skiff_formats(self, input: SkiffFormat, output: SkiffFormat) -> Self {
259 self.with_formats(DataFormat::skiff(input), DataFormat::skiff(output))
260 }
261
262 #[must_use]
264 pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
265 self.mapper.environment.push((key.into(), value.into()));
266 self
267 }
268
269 #[must_use]
273 pub fn with_input_table_index(mut self) -> Self {
274 self.input_table_index = true;
275 self
276 }
277
278 #[must_use]
280 pub fn with_job_count(mut self, count: i64) -> Self {
281 self.job_count = Some(count);
282 self
283 }
284
285 #[must_use]
295 pub fn skiff_table_mismatch(&self) -> Option<String> {
296 skiff_table_mismatch(
297 "the mapper's input_format",
298 &self.mapper.input_format,
299 self.inputs.len(),
300 "input table",
301 )
302 .or_else(|| {
303 skiff_table_mismatch(
304 "the mapper's output_format",
305 &self.mapper.output_format,
306 self.outputs.len(),
307 "output table",
308 )
309 })
310 }
311
312 #[must_use]
314 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
315 self.extra.push((key.into(), value));
316 self
317 }
318
319 #[must_use]
321 pub fn to_yson(&self) -> YsonValue {
322 let mut mapper = self.mapper.to_yson();
323 if self.input_table_index {
324 insert(&mut mapper, "enable_input_table_index", boolean(true));
325 }
326
327 let mut spec = map([
328 ("mapper", mapper),
329 ("input_table_paths", list(self.inputs.iter().map(string))),
330 ("output_table_paths", list(self.outputs.iter().map(string))),
331 ]);
332
333 if let Some(count) = self.job_count {
334 insert(&mut spec, "job_count", int(count));
335 }
336 for (key, value) in &self.extra {
337 insert(&mut spec, key, value.clone());
338 }
339 spec
340 }
341}
342
343#[derive(Debug, Clone)]
348pub struct MapReduceSpec {
349 mapper: Option<UserJob>,
350 mapper_formats: Option<(DataFormat, DataFormat)>,
356 reducer: UserJob,
357 files: Vec<YsonValue>,
362 memory_limit: Option<i64>,
363 inputs: Vec<String>,
364 outputs: Vec<String>,
365 reduce_by: Vec<String>,
366 sort_by: Vec<String>,
367 key_switch: bool,
368 extra: Vec<(String, YsonValue)>,
369}
370
371impl MapReduceSpec {
372 #[must_use]
374 pub fn new<I, O, K>(reducer: impl Into<String>, inputs: I, outputs: O, reduce_by: K) -> Self
375 where
376 I: IntoIterator,
377 I::Item: Into<String>,
378 O: IntoIterator,
379 O::Item: Into<String>,
380 K: IntoIterator,
381 K::Item: Into<String>,
382 {
383 Self {
384 mapper: None,
385 mapper_formats: None,
386 reducer: UserJob::new(reducer),
387 files: Vec::new(),
388 memory_limit: None,
389 inputs: inputs.into_iter().map(Into::into).collect(),
390 outputs: outputs.into_iter().map(Into::into).collect(),
391 reduce_by: reduce_by.into_iter().map(Into::into).collect(),
392 sort_by: Vec::new(),
393 key_switch: true,
396 extra: Vec::new(),
397 }
398 }
399
400 #[must_use]
402 pub fn with_mapper(mut self, command: impl Into<String>) -> Self {
403 self.mapper = Some(UserJob::new(command));
404 self
405 }
406
407 #[must_use]
412 pub fn with_mapper_formats(mut self, input: DataFormat, output: DataFormat) -> Self {
413 self.mapper_formats = Some((input, output));
414 self
415 }
416
417 #[must_use]
421 pub fn with_mapper_skiff_formats(self, input: SkiffFormat, output: SkiffFormat) -> Self {
422 self.with_mapper_formats(DataFormat::skiff(input), DataFormat::skiff(output))
423 }
424
425 #[must_use]
430 pub fn with_reducer_formats(mut self, input: DataFormat, output: DataFormat) -> Self {
431 self.reducer.with_formats(input, output);
432 self
433 }
434
435 #[must_use]
439 pub fn with_reducer_skiff_formats(self, input: SkiffFormat, output: SkiffFormat) -> Self {
440 self.with_reducer_formats(DataFormat::skiff(input), DataFormat::skiff(output))
441 }
442
443 #[must_use]
450 pub fn with_local_file(self, path: impl Into<String>) -> Self {
451 self.attach(string(path.into()))
452 }
453
454 #[must_use]
458 pub fn with_local_file_named(self, path: impl Into<String>, name: impl AsRef<str>) -> Self {
459 self.attach(named_file(path, name))
460 }
461
462 fn attach(mut self, file: YsonValue) -> Self {
463 self.files.push(file);
464 self
465 }
466
467 #[must_use]
472 pub fn with_memory_limit(mut self, bytes: i64) -> Self {
473 self.memory_limit = Some(bytes);
474 self
475 }
476
477 fn phase(&self, job: &UserJob, formats: Option<&(DataFormat, DataFormat)>) -> YsonValue {
479 let mut job = job.clone();
480 if let Some((input, output)) = formats {
481 job.with_formats(input.clone(), output.clone());
482 }
483 job.files.extend(self.files.iter().cloned());
484 if job.memory_limit.is_none() {
485 job.memory_limit = self.memory_limit;
486 }
487 job.to_yson()
488 }
489
490 #[must_use]
492 pub fn with_sort_by<K>(mut self, columns: K) -> Self
493 where
494 K: IntoIterator,
495 K::Item: Into<String>,
496 {
497 self.sort_by = columns.into_iter().map(Into::into).collect();
498 self
499 }
500
501 #[must_use]
506 pub fn without_key_switch(mut self) -> Self {
507 self.key_switch = false;
508 self
509 }
510
511 #[must_use]
525 pub fn skiff_table_mismatch(&self) -> Option<String> {
526 let split_outputs = self
527 .extra
528 .iter()
529 .any(|(key, _)| key == "mapper_output_table_count");
530
531 self.mapper
532 .as_ref()
533 .and(self.mapper_formats.as_ref())
534 .and_then(|(input, _)| {
535 skiff_table_mismatch(
536 "the mapper's input_format",
537 input,
538 self.inputs.len(),
539 "input table",
540 )
541 })
542 .or_else(|| {
543 if split_outputs {
544 return None;
545 }
546 skiff_table_mismatch(
547 "the reducer's output_format",
548 &self.reducer.output_format,
549 self.outputs.len(),
550 "output table",
551 )
552 })
553 }
554
555 #[must_use]
557 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
558 self.extra.push((key.into(), value));
559 self
560 }
561
562 #[must_use]
564 pub fn to_yson(&self) -> YsonValue {
565 let mut spec = map([
566 ("reducer", self.phase(&self.reducer, None)),
567 ("input_table_paths", list(self.inputs.iter().map(string))),
568 ("output_table_paths", list(self.outputs.iter().map(string))),
569 ("reduce_by", list(self.reduce_by.iter().map(string))),
570 ]);
571
572 if let Some(mapper) = &self.mapper {
573 insert(
574 &mut spec,
575 "mapper",
576 self.phase(mapper, self.mapper_formats.as_ref()),
577 );
578 }
579
580 let sort_by = if self.sort_by.is_empty() {
581 &self.reduce_by
582 } else {
583 &self.sort_by
584 };
585 insert(&mut spec, "sort_by", list(sort_by.iter().map(string)));
586
587 if self.key_switch {
588 insert(
593 &mut spec,
594 "reduce_job_io",
595 map([(
596 "control_attributes",
597 map([("enable_key_switch", boolean(true))]),
598 )]),
599 );
600 }
601
602 for (key, value) in &self.extra {
603 insert(&mut spec, key, value.clone());
604 }
605 spec
606 }
607}
608
609#[derive(Debug, Clone)]
623pub struct ReduceSpec {
624 reducer: UserJob,
625 inputs: Vec<String>,
626 outputs: Vec<String>,
627 reduce_by: Vec<String>,
628 sort_by: Vec<String>,
629 job_count: Option<i64>,
630 key_switch: bool,
631 input_table_index: bool,
632 extra: Vec<(String, YsonValue)>,
633}
634
635impl ReduceSpec {
636 #[must_use]
638 pub fn new<I, O, K>(command: impl Into<String>, inputs: I, outputs: O, reduce_by: K) -> Self
639 where
640 I: IntoIterator,
641 I::Item: Into<String>,
642 O: IntoIterator,
643 O::Item: Into<String>,
644 K: IntoIterator,
645 K::Item: Into<String>,
646 {
647 Self {
648 reducer: UserJob::new(command),
649 inputs: inputs.into_iter().map(Into::into).collect(),
650 outputs: outputs.into_iter().map(Into::into).collect(),
651 reduce_by: reduce_by.into_iter().map(Into::into).collect(),
652 sort_by: Vec::new(),
653 job_count: None,
654 key_switch: true,
657 input_table_index: false,
658 extra: Vec::new(),
659 }
660 }
661
662 #[must_use]
664 pub fn with_local_file(mut self, path: impl Into<String>) -> Self {
665 self.reducer.files.push(string(path.into()));
666 self
667 }
668
669 #[must_use]
673 pub fn with_local_file_named(mut self, path: impl Into<String>, name: impl AsRef<str>) -> Self {
674 self.reducer.files.push(named_file(path, name));
675 self
676 }
677
678 #[must_use]
680 pub fn with_memory_limit(mut self, bytes: i64) -> Self {
681 self.reducer.memory_limit = Some(bytes);
682 self
683 }
684
685 #[must_use]
699 pub fn with_formats(mut self, input: DataFormat, output: DataFormat) -> Self {
700 self.reducer.with_formats(input, output);
701 self
702 }
703
704 #[must_use]
708 pub fn with_skiff_formats(self, input: SkiffFormat, output: SkiffFormat) -> Self {
709 self.with_formats(DataFormat::skiff(input), DataFormat::skiff(output))
710 }
711
712 #[must_use]
714 pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
715 self.reducer.environment.push((key.into(), value.into()));
716 self
717 }
718
719 #[must_use]
726 pub fn with_sort_by<K>(mut self, columns: K) -> Self
727 where
728 K: IntoIterator,
729 K::Item: Into<String>,
730 {
731 self.sort_by = columns.into_iter().map(Into::into).collect();
732 self
733 }
734
735 #[must_use]
737 pub fn with_job_count(mut self, count: i64) -> Self {
738 self.job_count = Some(count);
739 self
740 }
741
742 #[must_use]
747 pub fn with_input_table_index(mut self) -> Self {
748 self.input_table_index = true;
749 self
750 }
751
752 #[must_use]
754 pub fn without_key_switch(mut self) -> Self {
755 self.key_switch = false;
756 self
757 }
758
759 #[must_use]
768 pub fn skiff_table_mismatch(&self) -> Option<String> {
769 skiff_table_mismatch(
770 "the reducer's input_format",
771 &self.reducer.input_format,
772 self.inputs.len(),
773 "input table",
774 )
775 .or_else(|| {
776 skiff_table_mismatch(
777 "the reducer's output_format",
778 &self.reducer.output_format,
779 self.outputs.len(),
780 "output table",
781 )
782 })
783 }
784
785 #[must_use]
787 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
788 self.extra.push((key.into(), value));
789 self
790 }
791
792 #[must_use]
794 pub fn to_yson(&self) -> YsonValue {
795 let mut reducer = self.reducer.to_yson();
796 if self.input_table_index {
797 insert(&mut reducer, "enable_input_table_index", boolean(true));
798 }
799
800 let mut spec = map([
801 ("reducer", reducer),
802 ("input_table_paths", list(self.inputs.iter().map(string))),
803 ("output_table_paths", list(self.outputs.iter().map(string))),
804 ("reduce_by", list(self.reduce_by.iter().map(string))),
805 ]);
806
807 if !self.sort_by.is_empty() {
808 insert(&mut spec, "sort_by", list(self.sort_by.iter().map(string)));
809 }
810 if let Some(count) = self.job_count {
811 insert(&mut spec, "job_count", int(count));
812 }
813
814 if self.key_switch {
815 insert(
820 &mut spec,
821 "job_io",
822 map([(
823 "control_attributes",
824 map([("enable_key_switch", boolean(true))]),
825 )]),
826 );
827 }
828
829 for (key, value) in &self.extra {
830 insert(&mut spec, key, value.clone());
831 }
832 spec
833 }
834}
835
836#[derive(Debug, Clone)]
847pub struct SortSpec {
848 inputs: Vec<String>,
849 output: String,
850 sort_by: Vec<String>,
851 extra: Vec<(String, YsonValue)>,
852}
853
854impl SortSpec {
855 #[must_use]
859 pub fn new<I, K>(inputs: I, output: impl Into<String>, sort_by: K) -> Self
860 where
861 I: IntoIterator,
862 I::Item: Into<String>,
863 K: IntoIterator,
864 K::Item: Into<String>,
865 {
866 Self {
867 inputs: inputs.into_iter().map(Into::into).collect(),
868 output: output.into(),
869 sort_by: sort_by.into_iter().map(Into::into).collect(),
870 extra: Vec::new(),
871 }
872 }
873
874 #[must_use]
877 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
878 self.extra.push((key.into(), value));
879 self
880 }
881
882 #[must_use]
884 pub fn to_yson(&self) -> YsonValue {
885 let mut spec = map([
886 ("input_table_paths", list(self.inputs.iter().map(string))),
887 ("output_table_path", string(&self.output)),
890 ("sort_by", list(self.sort_by.iter().map(string))),
891 ]);
892
893 for (key, value) in &self.extra {
894 insert(&mut spec, key, value.clone());
895 }
896 spec
897 }
898}
899
900#[derive(Debug, Clone, Copy, PartialEq, Eq)]
905pub enum MergeMode {
906 Unordered,
908 Ordered,
910 Sorted,
917}
918
919impl MergeMode {
920 #[must_use]
922 pub fn as_str(self) -> &'static str {
923 match self {
924 MergeMode::Unordered => "unordered",
925 MergeMode::Ordered => "ordered",
926 MergeMode::Sorted => "sorted",
927 }
928 }
929}
930
931#[derive(Debug, Clone)]
945pub struct MergeSpec {
946 inputs: Vec<String>,
947 output: String,
948 mode: MergeMode,
949 merge_by: Vec<String>,
950 combine_chunks: Option<bool>,
951 force_transform: Option<bool>,
952 job_count: Option<i64>,
953 extra: Vec<(String, YsonValue)>,
954}
955
956impl MergeSpec {
957 #[must_use]
959 pub fn new<I>(inputs: I, output: impl Into<String>) -> Self
960 where
961 I: IntoIterator,
962 I::Item: Into<String>,
963 {
964 Self {
965 inputs: inputs.into_iter().map(Into::into).collect(),
966 output: output.into(),
967 mode: MergeMode::Unordered,
968 merge_by: Vec::new(),
969 combine_chunks: None,
970 force_transform: None,
971 job_count: None,
972 extra: Vec::new(),
973 }
974 }
975
976 #[must_use]
978 pub fn with_mode(mut self, mode: MergeMode) -> Self {
979 self.mode = mode;
980 self
981 }
982
983 #[must_use]
990 pub fn with_merge_by<K>(mut self, columns: K) -> Self
991 where
992 K: IntoIterator,
993 K::Item: Into<String>,
994 {
995 self.merge_by = columns.into_iter().map(Into::into).collect();
996 self
997 }
998
999 #[must_use]
1004 pub fn with_combine_chunks(mut self, combine: bool) -> Self {
1005 self.combine_chunks = Some(combine);
1006 self
1007 }
1008
1009 #[must_use]
1015 pub fn with_force_transform(mut self, force: bool) -> Self {
1016 self.force_transform = Some(force);
1017 self
1018 }
1019
1020 #[must_use]
1025 pub fn with_job_count(mut self, count: i64) -> Self {
1026 self.job_count = Some(count);
1027 self
1028 }
1029
1030 #[must_use]
1033 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
1034 self.extra.push((key.into(), value));
1035 self
1036 }
1037
1038 #[must_use]
1040 pub fn to_yson(&self) -> YsonValue {
1041 let mut spec = map([
1042 ("input_table_paths", list(self.inputs.iter().map(string))),
1043 ("output_table_path", string(&self.output)),
1045 ("mode", string(self.mode.as_str())),
1046 ]);
1047
1048 if !self.merge_by.is_empty() {
1049 insert(
1050 &mut spec,
1051 "merge_by",
1052 list(self.merge_by.iter().map(string)),
1053 );
1054 }
1055 if let Some(combine) = self.combine_chunks {
1056 insert(&mut spec, "combine_chunks", boolean(combine));
1057 }
1058 if let Some(force) = self.force_transform {
1059 insert(&mut spec, "force_transform", boolean(force));
1060 }
1061 if let Some(count) = self.job_count {
1062 insert(&mut spec, "job_count", int(count));
1063 }
1064
1065 for (key, value) in &self.extra {
1066 insert(&mut spec, key, value.clone());
1067 }
1068 spec
1069 }
1070}
1071
1072#[derive(Debug, Clone)]
1085pub struct EraseSpec {
1086 table: String,
1087 combine_chunks: Option<bool>,
1088 extra: Vec<(String, YsonValue)>,
1089}
1090
1091impl EraseSpec {
1092 #[must_use]
1101 pub fn new(table: impl Into<String>) -> Self {
1102 Self {
1103 table: table.into(),
1104 combine_chunks: None,
1105 extra: Vec::new(),
1106 }
1107 }
1108
1109 #[must_use]
1111 pub fn with_combine_chunks(mut self, combine: bool) -> Self {
1112 self.combine_chunks = Some(combine);
1113 self
1114 }
1115
1116 #[must_use]
1118 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
1119 self.extra.push((key.into(), value));
1120 self
1121 }
1122
1123 #[must_use]
1125 pub fn to_yson(&self) -> YsonValue {
1126 let mut spec = map([("table_path", string(&self.table))]);
1129
1130 if let Some(combine) = self.combine_chunks {
1131 insert(&mut spec, "combine_chunks", boolean(combine));
1132 }
1133 for (key, value) in &self.extra {
1134 insert(&mut spec, key, value.clone());
1135 }
1136 spec
1137 }
1138}
1139
1140#[derive(Debug, Clone)]
1153pub struct RemoteCopySpec {
1154 cluster_name: String,
1155 inputs: Vec<String>,
1156 output: String,
1157 network_name: Option<String>,
1158 copy_attributes: Option<bool>,
1159 attribute_keys: Vec<String>,
1160 extra: Vec<(String, YsonValue)>,
1161}
1162
1163impl RemoteCopySpec {
1164 #[must_use]
1166 pub fn new<I>(cluster_name: impl Into<String>, inputs: I, output: impl Into<String>) -> Self
1167 where
1168 I: IntoIterator,
1169 I::Item: Into<String>,
1170 {
1171 Self {
1172 cluster_name: cluster_name.into(),
1173 inputs: inputs.into_iter().map(Into::into).collect(),
1174 output: output.into(),
1175 network_name: None,
1176 copy_attributes: None,
1177 attribute_keys: Vec::new(),
1178 extra: Vec::new(),
1179 }
1180 }
1181
1182 #[must_use]
1187 pub fn with_network_name(mut self, network: impl Into<String>) -> Self {
1188 self.network_name = Some(network.into());
1189 self
1190 }
1191
1192 #[must_use]
1197 pub fn with_copy_attributes(mut self, copy: bool) -> Self {
1198 self.copy_attributes = Some(copy);
1199 self
1200 }
1201
1202 #[must_use]
1206 pub fn with_attribute_keys<K>(mut self, keys: K) -> Self
1207 where
1208 K: IntoIterator,
1209 K::Item: Into<String>,
1210 {
1211 self.attribute_keys = keys.into_iter().map(Into::into).collect();
1212 self
1213 }
1214
1215 #[must_use]
1218 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
1219 self.extra.push((key.into(), value));
1220 self
1221 }
1222
1223 #[must_use]
1225 pub fn to_yson(&self) -> YsonValue {
1226 let mut spec = map([
1227 ("cluster_name", string(&self.cluster_name)),
1228 ("input_table_paths", list(self.inputs.iter().map(string))),
1229 ("output_table_path", string(&self.output)),
1230 ]);
1231
1232 if let Some(network) = &self.network_name {
1233 insert(&mut spec, "network_name", string(network));
1234 }
1235 if let Some(copy) = self.copy_attributes {
1236 insert(&mut spec, "copy_attributes", boolean(copy));
1237 }
1238 if !self.attribute_keys.is_empty() {
1239 insert(
1240 &mut spec,
1241 "attribute_keys",
1242 list(self.attribute_keys.iter().map(string)),
1243 );
1244 }
1245
1246 for (key, value) in &self.extra {
1247 insert(&mut spec, key, value.clone());
1248 }
1249 spec
1250 }
1251}
1252
1253#[derive(Debug, Clone)]
1259pub struct VanillaTask {
1260 name: String,
1261 job: UserJob,
1262 job_count: i64,
1263 outputs: Vec<String>,
1264 extra: Vec<(String, YsonValue)>,
1265}
1266
1267impl VanillaTask {
1268 #[must_use]
1273 pub fn new(name: impl Into<String>, command: impl Into<String>, job_count: i64) -> Self {
1274 Self {
1275 name: name.into(),
1276 job: UserJob::new(command),
1277 job_count,
1278 outputs: Vec::new(),
1279 extra: Vec::new(),
1280 }
1281 }
1282
1283 #[must_use]
1285 pub fn with_local_file(mut self, path: impl Into<String>) -> Self {
1286 self.job.files.push(string(path.into()));
1287 self
1288 }
1289
1290 #[must_use]
1294 pub fn with_local_file_named(mut self, path: impl Into<String>, name: impl AsRef<str>) -> Self {
1295 self.job.files.push(named_file(path, name));
1296 self
1297 }
1298
1299 #[must_use]
1304 pub fn with_outputs<O>(mut self, paths: O) -> Self
1305 where
1306 O: IntoIterator,
1307 O::Item: Into<String>,
1308 {
1309 self.outputs = paths.into_iter().map(Into::into).collect();
1310 self
1311 }
1312
1313 #[must_use]
1315 pub fn with_memory_limit(mut self, bytes: i64) -> Self {
1316 self.job.memory_limit = Some(bytes);
1317 self
1318 }
1319
1320 #[must_use]
1328 pub fn with_output_format(mut self, output: DataFormat) -> Self {
1329 self.job.output_format = output;
1330 self
1331 }
1332
1333 #[must_use]
1337 pub fn with_skiff_output_format(self, output: SkiffFormat) -> Self {
1338 self.with_output_format(DataFormat::skiff(output))
1339 }
1340
1341 #[must_use]
1343 pub fn with_env(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
1344 self.job.environment.push((key.into(), value.into()));
1345 self
1346 }
1347
1348 #[must_use]
1351 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
1352 self.extra.push((key.into(), value));
1353 self
1354 }
1355
1356 fn to_yson(&self) -> YsonValue {
1357 let mut task = self.job.to_yson();
1358 insert(&mut task, "job_count", int(self.job_count));
1359 insert(
1362 &mut task,
1363 "output_table_paths",
1364 list(self.outputs.iter().map(string)),
1365 );
1366
1367 for (key, value) in &self.extra {
1368 insert(&mut task, key, value.clone());
1369 }
1370 task
1371 }
1372}
1373
1374#[derive(Debug, Clone)]
1391pub struct VanillaSpec {
1392 tasks: Vec<VanillaTask>,
1393 extra: Vec<(String, YsonValue)>,
1394}
1395
1396impl VanillaSpec {
1397 #[must_use]
1399 pub fn new(task: VanillaTask) -> Self {
1400 Self {
1401 tasks: vec![task],
1402 extra: Vec::new(),
1403 }
1404 }
1405
1406 #[must_use]
1410 pub fn with_task(mut self, task: VanillaTask) -> Self {
1411 self.tasks.push(task);
1412 self
1413 }
1414
1415 #[must_use]
1424 pub fn duplicate_task(&self) -> Option<&str> {
1425 let mut seen = std::collections::HashSet::new();
1426 self.tasks
1427 .iter()
1428 .find(|task| !seen.insert(task.name.as_str()))
1429 .map(|task| task.name.as_str())
1430 }
1431
1432 #[must_use]
1438 pub fn skiff_table_mismatch(&self) -> Option<String> {
1439 self.tasks.iter().find_map(|task| {
1440 skiff_table_mismatch(
1441 &format!("task {:?}'s output_format", task.name),
1442 &task.job.output_format,
1443 task.outputs.len(),
1444 "output table",
1445 )
1446 })
1447 }
1448
1449 #[must_use]
1451 pub fn with_raw(mut self, key: impl Into<String>, value: YsonValue) -> Self {
1452 self.extra.push((key.into(), value));
1453 self
1454 }
1455
1456 #[must_use]
1458 pub fn to_yson(&self) -> YsonValue {
1459 let mut spec = map([(
1460 "tasks",
1461 map(self
1462 .tasks
1463 .iter()
1464 .map(|task| (task.name.as_str(), task.to_yson()))),
1465 )]);
1466
1467 for (key, value) in &self.extra {
1468 insert(&mut spec, key, value.clone());
1469 }
1470 spec
1471 }
1472}
1473
1474#[cfg(test)]
1475mod tests {
1476 use super::*;
1477 use ytsaurus_skiff::{Schema, SchemaRef, WireType};
1478 use ytsaurus_yson::{YsonFormat, to_string};
1479
1480 fn render(v: &YsonValue) -> String {
1481 to_string(v, YsonFormat::Text).expect("encodes")
1482 }
1483
1484 fn skiff_format(column: &str) -> SkiffFormat {
1485 SkiffFormat::new(vec![SchemaRef::Inline(Schema::tuple([Schema::named(
1486 column,
1487 WireType::Uint64,
1488 )]))])
1489 .expect("a named tuple is a table schema")
1490 }
1491
1492 fn skiff_tables(columns: &[&str]) -> SkiffFormat {
1493 SkiffFormat::new(
1494 columns
1495 .iter()
1496 .map(|column| {
1497 SchemaRef::Inline(Schema::tuple([Schema::named(*column, WireType::Uint64)]))
1498 })
1499 .collect(),
1500 )
1501 .expect("named tuples are table schemas")
1502 }
1503
1504 #[test]
1505 fn a_map_skiff_format_needs_one_schema_per_table() {
1506 let two_in_one_out = MapSpec::new("./w", ["//a", "//b"], ["//out"]);
1507
1508 let short_input = two_in_one_out
1509 .clone()
1510 .with_skiff_formats(skiff_tables(&["source"]), skiff_tables(&["result"]));
1511 let reason = short_input
1512 .skiff_table_mismatch()
1513 .expect("one schema cannot describe two input tables");
1514 assert!(reason.contains("input_format"), "{reason}");
1515 assert!(reason.contains("1 Skiff table schema,"), "{reason}");
1516 assert!(reason.contains("2 input tables"), "{reason}");
1517
1518 let long_output = two_in_one_out
1519 .clone()
1520 .with_skiff_formats(skiff_tables(&["a", "b"]), skiff_tables(&["x", "y"]));
1521 let reason = long_output
1522 .skiff_table_mismatch()
1523 .expect("two schemas cannot describe one output table");
1524 assert!(reason.contains("output_format"), "{reason}");
1525
1526 assert!(
1527 two_in_one_out
1528 .clone()
1529 .with_skiff_formats(skiff_tables(&["a", "b"]), skiff_tables(&["x"]))
1530 .skiff_table_mismatch()
1531 .is_none()
1532 );
1533 assert!(two_in_one_out.skiff_table_mismatch().is_none());
1535 }
1536
1537 #[test]
1538 fn a_reduce_skiff_format_needs_one_schema_per_table() {
1539 let spec = ReduceSpec::new("./w", ["//a", "//b"], ["//out"], ["key"]);
1540
1541 let reason = spec
1542 .clone()
1543 .with_skiff_formats(skiff_tables(&["source"]), skiff_tables(&["result"]))
1544 .skiff_table_mismatch()
1545 .expect("a reduce input format describes every input table");
1546 assert!(reason.contains("input_format"), "{reason}");
1547
1548 assert!(
1549 spec.with_skiff_formats(skiff_tables(&["a", "b"]), skiff_tables(&["x"]))
1550 .skiff_table_mismatch()
1551 .is_none()
1552 );
1553 }
1554
1555 #[test]
1556 fn map_reduce_checks_the_counts_it_knows_and_leaves_the_shuffle_alone() {
1557 let spec = MapReduceSpec::new("./r", ["//a", "//b"], ["//out"], ["key"])
1558 .with_mapper("./m")
1559 .with_mapper_skiff_formats(skiff_tables(&["one"]), skiff_tables(&["shuffle"]));
1560 let reason = spec
1561 .skiff_table_mismatch()
1562 .expect("the mapper still reads the operation's input tables");
1563 assert!(reason.contains("input_format"), "{reason}");
1564
1565 let shuffle = MapReduceSpec::new("./r", ["//a"], ["//out"], ["key"])
1568 .with_mapper("./m")
1569 .with_mapper_skiff_formats(skiff_tables(&["one"]), skiff_tables(&["x", "y"]))
1570 .with_reducer_skiff_formats(skiff_tables(&["x", "y"]), skiff_tables(&["out"]));
1571 assert!(shuffle.skiff_table_mismatch().is_none());
1572 }
1573
1574 #[test]
1575 fn a_vanilla_task_skiff_output_needs_one_schema_per_output() {
1576 let spec = VanillaSpec::new(
1577 VanillaTask::new("worker", "./w", 1)
1578 .with_outputs(["//one"])
1579 .with_skiff_output_format(skiff_tables(&["a", "b"])),
1580 );
1581 let reason = spec
1582 .skiff_table_mismatch()
1583 .expect("two schemas cannot describe one output table");
1584 assert!(reason.contains(r#"task "worker""#), "{reason}");
1585
1586 assert!(
1587 VanillaSpec::new(
1588 VanillaTask::new("worker", "./w", 1)
1589 .with_outputs(["//one"])
1590 .with_skiff_output_format(skiff_tables(&["a"])),
1591 )
1592 .skiff_table_mismatch()
1593 .is_none()
1594 );
1595 }
1596
1597 #[test]
1598 fn a_map_spec_carries_what_the_operation_needs() {
1599 let spec = MapSpec::new("./cat", ["//tmp/in"], ["//tmp/out"])
1600 .with_local_file("//tmp/cat")
1601 .with_memory_limit(1024);
1602 let out = render(&spec.to_yson());
1603
1604 assert!(out.contains(r#"command="./cat""#), "{out}");
1605 assert!(out.contains(r#"file_paths=["//tmp/cat"]"#), "{out}");
1606 assert!(out.contains("memory_limit=1024"), "{out}");
1607 assert!(out.contains(r#"input_table_paths=["//tmp/in"]"#), "{out}");
1608 assert!(out.contains(r#"output_table_paths=["//tmp/out"]"#), "{out}");
1609 assert!(out.contains("input_format=<format=binary>yson"), "{out}");
1610 }
1611
1612 #[test]
1613 fn multiple_outputs_are_preserved_in_order() {
1614 let spec = MapSpec::new("./cat", ["//tmp/a", "//tmp/b"], ["//tmp/x", "//tmp/y"]);
1615 let out = render(&spec.to_yson());
1616 assert!(
1617 out.contains(r#"input_table_paths=["//tmp/a";"//tmp/b"]"#),
1618 "{out}"
1619 );
1620 assert!(
1621 out.contains(r#"output_table_paths=["//tmp/x";"//tmp/y"]"#),
1622 "{out}"
1623 );
1624 }
1625
1626 #[test]
1627 fn map_can_select_schema_checked_skiff_for_both_directions() {
1628 let out = render(
1629 &MapSpec::new("./worker", ["//in"], ["//out"])
1630 .with_skiff_formats(skiff_format("source"), skiff_format("result"))
1631 .to_yson(),
1632 );
1633
1634 assert!(out.contains("input_format=<table_skiff_schemas="), "{out}");
1635 assert!(out.contains("output_format=<table_skiff_schemas="), "{out}");
1636 assert!(out.contains("name=source"), "{out}");
1637 assert!(out.contains("name=result"), "{out}");
1638 assert!(!out.contains("format=binary"), "{out}");
1639 }
1640
1641 #[test]
1642 fn map_can_select_yson_and_skiff_through_the_shared_format_enum() {
1643 let out = render(
1644 &MapSpec::new("./worker", ["//in"], ["//out"])
1645 .with_formats(
1646 DataFormat::text_yson(),
1647 DataFormat::skiff(skiff_format("result")),
1648 )
1649 .to_yson(),
1650 );
1651
1652 assert!(out.contains("input_format=<format=text>yson"), "{out}");
1653 assert!(out.contains("output_format=<table_skiff_schemas="), "{out}");
1654 assert!(out.contains("name=result"), "{out}");
1655 }
1656
1657 #[test]
1658 fn table_index_is_off_unless_asked_for() {
1659 let plain = render(&MapSpec::new("./c", ["//i"], ["//o"]).to_yson());
1660 assert!(!plain.contains("enable_input_table_index"), "{plain}");
1661
1662 let asked = render(
1663 &MapSpec::new("./c", ["//i"], ["//o"])
1664 .with_input_table_index()
1665 .to_yson(),
1666 );
1667 assert!(asked.contains("enable_input_table_index=%true"), "{asked}");
1668 }
1669
1670 #[test]
1673 fn map_reduce_puts_key_switch_under_reduce_job_io() {
1674 let spec = MapReduceSpec::new("./wc reduce", ["//in"], ["//out"], ["word"])
1675 .with_mapper("./wc map");
1676 let out = render(&spec.to_yson());
1677
1678 assert!(
1679 out.contains("reduce_job_io={control_attributes={enable_key_switch=%true}}"),
1680 "{out}"
1681 );
1682 assert!(
1685 !out.contains(";job_io=") && !out.contains("{job_io="),
1686 "must not use the plain job_io section: {out}"
1687 );
1688 }
1689
1690 #[test]
1691 fn map_reduce_can_select_skiff_per_job_phase() {
1692 let out = render(
1693 &MapReduceSpec::new("./worker reduce", ["//in"], ["//out"], ["key"])
1694 .with_mapper("./worker map")
1695 .with_mapper_skiff_formats(skiff_format("map_input"), skiff_format("map_output"))
1696 .with_reducer_skiff_formats(
1697 skiff_format("reduce_input"),
1698 skiff_format("reduce_output"),
1699 )
1700 .to_yson(),
1701 );
1702
1703 for column in ["map_input", "map_output", "reduce_input", "reduce_output"] {
1704 assert!(out.contains(&format!("name={column}")), "{out}");
1705 }
1706 assert_eq!(
1707 out.matches("input_format=<table_skiff_schemas=").count(),
1708 2,
1709 "{out}"
1710 );
1711 assert_eq!(
1712 out.matches("output_format=<table_skiff_schemas=").count(),
1713 2,
1714 "{out}"
1715 );
1716 }
1717
1718 #[test]
1721 fn a_mapper_added_last_still_gets_its_formats() {
1722 let before = render(
1723 &MapReduceSpec::new("./r", ["//in"], ["//out"], ["k"])
1724 .with_mapper_skiff_formats(skiff_format("map_input"), skiff_format("map_output"))
1725 .with_mapper("./m")
1726 .to_yson(),
1727 );
1728 let after = render(
1729 &MapReduceSpec::new("./r", ["//in"], ["//out"], ["k"])
1730 .with_mapper("./m")
1731 .with_mapper_skiff_formats(skiff_format("map_input"), skiff_format("map_output"))
1732 .to_yson(),
1733 );
1734
1735 assert_eq!(before, after);
1736 assert!(before.contains("name=map_input"), "{before}");
1737 assert!(before.contains("name=map_output"), "{before}");
1738 }
1739
1740 #[test]
1741 fn key_switch_can_be_turned_off() {
1742 let out = render(
1743 &MapReduceSpec::new("./r", ["//in"], ["//out"], ["k"])
1744 .without_key_switch()
1745 .to_yson(),
1746 );
1747 assert!(!out.contains("enable_key_switch"), "{out}");
1748 }
1749
1750 #[test]
1751 fn sort_by_defaults_to_reduce_by() {
1752 let out = render(&MapReduceSpec::new("./r", ["//in"], ["//out"], ["k"]).to_yson());
1753 assert!(out.contains("sort_by=[k]"), "{out}");
1754
1755 let out = render(
1756 &MapReduceSpec::new("./r", ["//in"], ["//out"], ["k"])
1757 .with_sort_by(["k", "ts"])
1758 .to_yson(),
1759 );
1760 assert!(out.contains("sort_by=[k;ts]"), "{out}");
1761 }
1762
1763 #[test]
1764 fn one_file_reaches_both_phases() {
1765 let out = render(
1766 &MapReduceSpec::new("./w reduce", ["//in"], ["//out"], ["k"])
1767 .with_mapper("./w map")
1768 .with_local_file("//tmp/w")
1769 .to_yson(),
1770 );
1771 assert_eq!(
1772 out.matches(r#"file_paths=["//tmp/w"]"#).count(),
1773 2,
1774 "the binary must be attached to both phases: {out}"
1775 );
1776 }
1777
1778 #[test]
1783 fn a_mapper_added_last_still_gets_the_files_and_the_limit() {
1784 let out = render(
1785 &MapReduceSpec::new("./w reduce", ["//in"], ["//out"], ["k"])
1786 .with_local_file("//tmp/w")
1787 .with_memory_limit(512 * 1024 * 1024)
1788 .with_mapper("./w map")
1789 .to_yson(),
1790 );
1791 assert_eq!(
1792 out.matches(r#"file_paths=["//tmp/w"]"#).count(),
1793 2,
1794 "the binary must reach both phases whatever the call order: {out}"
1795 );
1796 assert_eq!(
1797 out.matches("memory_limit=536870912").count(),
1798 2,
1799 "the limit must reach both phases whatever the call order: {out}"
1800 );
1801 }
1802
1803 #[test]
1806 fn a_named_file_carries_its_sandbox_name() {
1807 let cached = "//tmp/yt_wrapper/file_storage/new_cache/da/2c76e46b90e8b9d5ec25397e14c043da";
1808 let out = render(
1809 &MapSpec::new("./cat", ["//i"], ["//o"])
1810 .with_local_file_named(cached, "cat")
1811 .to_yson(),
1812 );
1813
1814 assert!(out.contains("file_name=cat"), "{out}");
1815 assert!(out.contains(cached), "{out}");
1816 }
1817
1818 #[test]
1819 fn a_named_file_reaches_both_map_reduce_phases() {
1820 let out = render(
1821 &MapReduceSpec::new("./w reduce", ["//in"], ["//out"], ["k"])
1822 .with_mapper("./w map")
1823 .with_local_file_named("//tmp/cache/ab/cd", "w")
1824 .to_yson(),
1825 );
1826 assert_eq!(
1827 out.matches("file_name=w").count(),
1828 2,
1829 "the binary must be attached to both phases: {out}"
1830 );
1831 }
1832
1833 #[test]
1834 fn a_plain_file_gets_no_attributes() {
1835 let out = render(
1836 &MapSpec::new("./cat", ["//i"], ["//o"])
1837 .with_local_file("//tmp/cat")
1838 .to_yson(),
1839 );
1840 assert!(out.contains(r#"file_paths=["//tmp/cat"]"#), "{out}");
1841 }
1842
1843 #[test]
1844 fn raw_fields_land_in_the_spec() {
1845 let out = render(
1846 &MapSpec::new("./c", ["//i"], ["//o"])
1847 .with_raw("max_failed_job_count", int(3))
1848 .to_yson(),
1849 );
1850 assert!(out.contains("max_failed_job_count=3"), "{out}");
1851 }
1852
1853 #[test]
1857 fn operation_type_wire_names() {
1858 assert_eq!(OperationType::Map.as_str(), "map");
1859 assert_eq!(OperationType::MapReduce.as_str(), "map_reduce");
1860 assert_eq!(OperationType::Reduce.as_str(), "reduce");
1861 assert_eq!(OperationType::Sort.as_str(), "sort");
1862 assert_eq!(OperationType::Vanilla.as_str(), "vanilla");
1863 assert_eq!(OperationType::Merge.as_str(), "merge");
1864 assert_eq!(OperationType::Erase.as_str(), "erase");
1865 assert_eq!(OperationType::RemoteCopy.as_str(), "remote_copy");
1866 assert_eq!(OperationType::JoinReduce.as_str(), "join_reduce");
1867 }
1868
1869 #[test]
1870 fn merge_mode_wire_names() {
1871 assert_eq!(MergeMode::Unordered.as_str(), "unordered");
1872 assert_eq!(MergeMode::Ordered.as_str(), "ordered");
1873 assert_eq!(MergeMode::Sorted.as_str(), "sorted");
1874 }
1875
1876 #[test]
1879 fn a_merge_spec_names_one_output() {
1880 let out = render(&MergeSpec::new(["//tmp/a", "//tmp/b"], "//tmp/all").to_yson());
1881
1882 assert!(
1883 out.contains(r#"input_table_paths=["//tmp/a";"//tmp/b"]"#),
1884 "{out}"
1885 );
1886 assert!(out.contains(r#"output_table_path="//tmp/all""#), "{out}");
1887 assert!(
1888 out.contains("mode=unordered"),
1889 "the cheapest mode is the default, and it is sent rather than \
1890 assumed: {out}"
1891 );
1892 assert!(!out.contains("merge_by"), "{out}");
1893 }
1894
1895 #[test]
1896 fn a_sorted_merge_carries_its_key() {
1897 let spec = MergeSpec::new(["//tmp/a"], "//tmp/all")
1898 .with_mode(MergeMode::Sorted)
1899 .with_merge_by(["host", "day"])
1900 .with_combine_chunks(true)
1901 .with_job_count(4);
1902 let out = render(&spec.to_yson());
1903
1904 assert!(out.contains("mode=sorted"), "{out}");
1905 assert!(out.contains("merge_by=[host;day]"), "{out}");
1908 assert!(out.contains("combine_chunks=%true"), "{out}");
1909 assert!(out.contains("job_count=4"), "{out}");
1910 let _ = spec;
1911 }
1912
1913 #[test]
1918 fn a_sorted_merge_may_leave_its_key_to_the_cluster() {
1919 let out = render(
1920 &MergeSpec::new(["//tmp/a"], "//tmp/all")
1921 .with_mode(MergeMode::Sorted)
1922 .to_yson(),
1923 );
1924
1925 assert!(out.contains("mode=sorted"), "{out}");
1926 assert!(
1927 !out.contains("merge_by"),
1928 "an absent key is the request to infer one: {out}"
1929 );
1930 }
1931
1932 #[test]
1935 fn a_key_set_through_the_escape_hatch_is_rendered() {
1936 let out = render(
1937 &MergeSpec::new(["//tmp/a"], "//tmp/all")
1938 .with_mode(MergeMode::Sorted)
1939 .with_raw("merge_by", list([string("host")]))
1940 .to_yson(),
1941 );
1942 assert!(out.contains("merge_by=[host]"), "{out}");
1943 }
1944
1945 #[test]
1948 fn an_erase_spec_names_the_table_once() {
1949 let out = render(&EraseSpec::new("//tmp/log[#0:#10]").to_yson());
1950
1951 assert_eq!(out, r#"{table_path="//tmp/log[#0:#10]"}"#);
1952 }
1953
1954 #[test]
1955 fn an_erase_spec_can_ask_for_compaction() {
1956 let out = render(
1957 &EraseSpec::new("//tmp/log")
1958 .with_combine_chunks(true)
1959 .to_yson(),
1960 );
1961 assert!(out.contains("combine_chunks=%true"), "{out}");
1962 }
1963
1964 #[test]
1965 fn a_remote_copy_spec_names_the_source_cluster() {
1966 let spec = RemoteCopySpec::new("hahn", ["//tmp/theirs"], "//tmp/ours")
1967 .with_network_name("fastbone")
1968 .with_copy_attributes(true)
1969 .with_attribute_keys(["expiration_time"]);
1970 let out = render(&spec.to_yson());
1971
1972 assert!(out.contains("cluster_name=hahn"), "{out}");
1973 assert!(
1974 out.contains(r#"input_table_paths=["//tmp/theirs"]"#),
1975 "{out}"
1976 );
1977 assert!(out.contains(r#"output_table_path="//tmp/ours""#), "{out}");
1978 assert!(out.contains("network_name=fastbone"), "{out}");
1979 assert!(out.contains("copy_attributes=%true"), "{out}");
1980 assert!(out.contains("attribute_keys=[expiration_time]"), "{out}");
1981 }
1982
1983 #[test]
1984 fn the_new_specs_take_raw_fields_too() {
1985 let merge = render(
1986 &MergeSpec::new(["//i"], "//o")
1987 .with_raw("schema_inference_mode", string("from_output"))
1988 .to_yson(),
1989 );
1990 assert!(
1991 merge.contains("schema_inference_mode=from_output"),
1992 "{merge}"
1993 );
1994
1995 let erase = render(
1996 &EraseSpec::new("//t")
1997 .with_raw("schema_inference_mode", string("auto"))
1998 .to_yson(),
1999 );
2000 assert!(erase.contains("schema_inference_mode=auto"), "{erase}");
2001
2002 let copy = render(
2003 &RemoteCopySpec::new("c", ["//i"], "//o")
2004 .with_raw("allow_unfrozen_input_tables", boolean(true))
2005 .to_yson(),
2006 );
2007 assert!(copy.contains("allow_unfrozen_input_tables=%true"), "{copy}");
2008 }
2009
2010 #[test]
2013 fn reduce_puts_key_switch_under_job_io() {
2014 let out =
2015 render(&ReduceSpec::new("./wc reduce", ["//sorted"], ["//out"], ["word"]).to_yson());
2016
2017 assert!(
2018 out.contains("job_io={control_attributes={enable_key_switch=%true}}"),
2019 "{out}"
2020 );
2021 assert!(
2022 !out.contains("reduce_job_io"),
2023 "reduce_job_io belongs to map-reduce, not to reduce: {out}"
2024 );
2025 }
2026
2027 #[test]
2028 fn a_reduce_spec_carries_what_the_operation_needs() {
2029 let spec = ReduceSpec::new("./wc reduce", ["//tmp/sorted"], ["//tmp/counts"], ["word"])
2030 .with_local_file("//tmp/wc")
2031 .with_memory_limit(1024)
2032 .with_job_count(2);
2033 let out = render(&spec.to_yson());
2034
2035 assert!(out.contains(r#"command="./wc reduce""#), "{out}");
2036 assert!(out.contains(r#"file_paths=["//tmp/wc"]"#), "{out}");
2037 assert!(out.contains("memory_limit=1024"), "{out}");
2038 assert!(out.contains("reduce_by=[word]"), "{out}");
2039 assert!(out.contains("job_count=2"), "{out}");
2040 assert!(out.contains("input_format=<format=binary>yson"), "{out}");
2041 }
2042
2043 #[test]
2047 fn reduce_sort_by_is_only_sent_when_set() {
2048 let plain = render(&ReduceSpec::new("./r", ["//in"], ["//out"], ["k"]).to_yson());
2049 assert!(!plain.contains("sort_by"), "{plain}");
2050
2051 let asked = render(
2052 &ReduceSpec::new("./r", ["//in"], ["//out"], ["k"])
2053 .with_sort_by(["k", "ts"])
2054 .to_yson(),
2055 );
2056 assert!(asked.contains("sort_by=[k;ts]"), "{asked}");
2057 }
2058
2059 #[test]
2060 fn reduce_table_index_is_off_unless_asked_for() {
2061 let plain = render(&ReduceSpec::new("./r", ["//a", "//b"], ["//o"], ["k"]).to_yson());
2062 assert!(!plain.contains("enable_input_table_index"), "{plain}");
2063
2064 let asked = render(
2065 &ReduceSpec::new("./r", ["//a", "//b"], ["//o"], ["k"])
2066 .with_input_table_index()
2067 .to_yson(),
2068 );
2069 assert!(asked.contains("enable_input_table_index=%true"), "{asked}");
2070 }
2071
2072 #[test]
2075 fn sort_writes_one_table_through_a_singular_field() {
2076 let out = render(&SortSpec::new(["//a", "//b"], "//sorted", ["key", "sub"]).to_yson());
2077
2078 assert!(out.contains(r#"output_table_path="//sorted""#), "{out}");
2079 assert!(!out.contains("output_table_paths"), "{out}");
2080 assert!(out.contains(r#"input_table_paths=["//a";"//b"]"#), "{out}");
2081 assert!(out.contains("sort_by=[key;sub]"), "{out}");
2082 }
2083
2084 #[test]
2085 fn a_sort_spec_has_no_user_job() {
2086 let out = render(&SortSpec::new(["//a"], "//sorted", ["key"]).to_yson());
2087 assert!(
2088 !out.contains("command"),
2089 "the cluster sorts, not a job: {out}"
2090 );
2091 assert!(!out.contains("input_format"), "{out}");
2092 }
2093
2094 #[test]
2095 fn a_vanilla_spec_describes_its_tasks() {
2096 let out = render(
2097 &VanillaSpec::new(
2098 VanillaTask::new("worker", "./my_job", 4)
2099 .with_local_file("//tmp/my_job")
2100 .with_outputs(["//tmp/results"])
2101 .with_memory_limit(1024),
2102 )
2103 .with_task(VanillaTask::new("master", "./my_job master", 1))
2104 .with_raw("max_failed_job_count", int(1))
2105 .to_yson(),
2106 );
2107
2108 assert!(out.contains("tasks={"), "{out}");
2109 assert!(out.contains("worker={"), "{out}");
2110 assert!(out.contains("master={"), "{out}");
2111 assert!(out.contains("job_count=4"), "{out}");
2112 assert!(out.contains("job_count=1"), "{out}");
2113 assert!(
2114 out.contains(r#"output_table_paths=["//tmp/results"]"#),
2115 "{out}"
2116 );
2117 assert!(out.contains("max_failed_job_count=1"), "{out}");
2118 assert!(!out.contains("input_table_paths"), "{out}");
2120 }
2121
2122 #[test]
2123 fn reduce_can_select_skiff_for_both_directions() {
2124 let out = render(
2125 &ReduceSpec::new("./worker", ["//in"], ["//out"], ["key"])
2126 .with_skiff_formats(skiff_format("reduce_input"), skiff_format("reduce_output"))
2127 .to_yson(),
2128 );
2129
2130 assert!(out.contains("input_format=<table_skiff_schemas="), "{out}");
2131 assert!(out.contains("output_format=<table_skiff_schemas="), "{out}");
2132 assert!(out.contains("name=reduce_input"), "{out}");
2133 assert!(out.contains("name=reduce_output"), "{out}");
2134 assert!(!out.contains("format=binary"), "{out}");
2135 assert!(
2138 out.contains("control_attributes={enable_key_switch=%true}"),
2139 "{out}"
2140 );
2141 }
2142
2143 #[test]
2144 fn a_vanilla_task_can_select_skiff_output_only() {
2145 let out = render(
2146 &VanillaSpec::new(
2147 VanillaTask::new("worker", "./my_job", 1)
2148 .with_outputs(["//tmp/results"])
2149 .with_skiff_output_format(skiff_format("result")),
2150 )
2151 .to_yson(),
2152 );
2153
2154 assert!(out.contains("output_format=<table_skiff_schemas="), "{out}");
2155 assert!(out.contains("name=result"), "{out}");
2156 assert!(out.contains("input_format=<format=binary>yson"), "{out}");
2159 }
2160
2161 #[test]
2162 fn two_tasks_with_one_name_are_caught_before_the_cluster_sees_them() {
2163 let spec = VanillaSpec::new(VanillaTask::new("worker", "./j shard-a", 4))
2167 .with_task(VanillaTask::new("worker", "./j shard-b", 4));
2168
2169 assert_eq!(spec.duplicate_task(), Some("worker"));
2170
2171 let out = render(&spec.to_yson());
2172 assert!(!out.contains("shard-a"), "the first task is gone: {out}");
2173 }
2174
2175 #[test]
2176 fn tasks_with_distinct_names_are_fine() {
2177 let spec = VanillaSpec::new(VanillaTask::new("worker", "./j", 4))
2178 .with_task(VanillaTask::new("master", "./j master", 1));
2179 assert_eq!(spec.duplicate_task(), None);
2180 }
2181
2182 #[test]
2185 fn a_task_without_outputs_says_so() {
2186 let out = render(&VanillaSpec::new(VanillaTask::new("t", "./j", 1)).to_yson());
2187 assert!(out.contains("output_table_paths=[]"), "{out}");
2188 }
2189
2190 #[test]
2191 fn gang_options_go_through_raw() {
2192 let out = render(
2193 &VanillaSpec::new(
2194 VanillaTask::new("worker", "./j", 3).with_raw("gang_options", map::<&str>([])),
2195 )
2196 .to_yson(),
2197 );
2198 assert!(out.contains("gang_options={}"), "{out}");
2199 }
2200
2201 #[test]
2202 fn sort_tuning_goes_through_raw() {
2203 let out = render(
2204 &SortSpec::new(["//a"], "//sorted", ["key"])
2205 .with_raw("partition_count", int(4))
2206 .to_yson(),
2207 );
2208 assert!(out.contains("partition_count=4"), "{out}");
2209 }
2210}