1#![cfg_attr(
23 not(any(feature = "file-format-avro", feature = "file-format-orc")),
24 allow(unreachable_code, unused_variables, dead_code)
25)]
26
27use super::{FileFormat, FormatOptions};
28use crate::error::FaucetError;
29use serde_json::Value;
30
31#[derive(Debug)]
33pub enum FileInput {
34 File(std::fs::File),
37 Bytes(Vec<u8>),
39}
40
41#[derive(Debug)]
42enum Anchor {
43 #[cfg(feature = "file-format-avro")]
44 Avro(apache_avro::Schema),
45 #[cfg(feature = "file-format-orc")]
46 Orc(arrow::datatypes::SchemaRef),
47}
48
49#[derive(Debug)]
51pub struct ContainerDecoder {
52 format: FileFormat,
53 #[cfg(feature = "file-format-orc")]
54 opts: FormatOptions,
55 #[cfg(feature = "file-format-avro")]
56 configured: bool,
57 anchor: Option<(String, Anchor)>,
58}
59
60impl ContainerDecoder {
61 pub fn new(format: FileFormat, opts: &FormatOptions) -> Result<Self, FaucetError> {
64 let anchor: Option<(String, Anchor)> = match format {
65 #[cfg(feature = "file-format-avro")]
66 FileFormat::Avro => opts
67 .avro
68 .parsed_schema()?
69 .map(|s| ("avro.schema".to_string(), Anchor::Avro(s))),
70 #[cfg(feature = "file-format-orc")]
71 FileFormat::Orc => None,
72 #[cfg(not(feature = "file-format-avro"))]
73 FileFormat::Avro => return Err(super::missing_feature(format, "file-format-avro")),
74 #[cfg(not(feature = "file-format-orc"))]
75 FileFormat::Orc => return Err(super::missing_feature(format, "file-format-orc")),
76 other => {
77 return Err(FaucetError::Config(format!(
78 "ContainerDecoder handles avro and orc, not `{}`",
79 other.as_str()
80 )));
81 }
82 };
83 Ok(Self {
84 format,
85 #[cfg(feature = "file-format-orc")]
86 opts: opts.clone(),
87 #[cfg(feature = "file-format-avro")]
88 configured: anchor.is_some(),
89 anchor,
90 })
91 }
92
93 pub fn format(&self) -> FileFormat {
95 self.format
96 }
97
98 pub fn decode_all(&mut self, name: &str, input: FileInput) -> Result<Vec<Value>, FaucetError> {
100 let mut out = Vec::new();
101 self.records(name, input, 0, &mut |c| {
102 out.extend(c);
103 Ok(())
104 })?;
105 Ok(out)
106 }
107
108 pub async fn decode_all_offloaded(
112 mut self,
113 name: String,
114 input: FileInput,
115 ) -> Result<(Self, Vec<Value>), FaucetError> {
116 tokio::task::spawn_blocking(move || {
117 let rows = self.decode_all(&name, input)?;
118 Ok((self, rows))
119 })
120 .await
121 .map_err(|e| FaucetError::Source(format!("container decode task failed: {e}")))?
122 }
123
124 #[cfg(feature = "arrow")]
126 #[allow(clippy::type_complexity)]
127 pub async fn decode_batches_offloaded(
128 mut self,
129 name: String,
130 input: FileInput,
131 batch_size: usize,
132 ) -> Result<(Self, Vec<arrow::array::RecordBatch>), FaucetError> {
133 tokio::task::spawn_blocking(move || {
134 let (_, batches) = self.decode_batches(&name, input, batch_size)?;
135 Ok((self, batches))
136 })
137 .await
138 .map_err(|e| FaucetError::Source(format!("container decode task failed: {e}")))?
139 }
140
141 #[cfg(feature = "arrow")]
143 pub fn decode_batches(
144 &mut self,
145 name: &str,
146 input: FileInput,
147 batch_size: usize,
148 ) -> Result<(arrow::datatypes::SchemaRef, Vec<arrow::array::RecordBatch>), FaucetError> {
149 let mut out = Vec::new();
150 let schema = self.batches(name, input, batch_size, &mut |b| {
151 out.push(b);
152 Ok(())
153 })?;
154 Ok((schema, out))
155 }
156
157 pub fn records(
160 &mut self,
161 name: &str,
162 input: FileInput,
163 chunk: usize,
164 f: &mut dyn FnMut(Vec<Value>) -> Result<(), FaucetError>,
165 ) -> Result<(), FaucetError> {
166 let chunk = if chunk == 0 { usize::MAX } else { chunk };
167 match self.format {
168 #[cfg(feature = "file-format-avro")]
169 FileFormat::Avro => {
170 let reader_schema = self.avro_reader_schema();
171 let check = self.avro_writer_check();
172 let result = with_reader(input, |r| {
173 super::avro::read_records_checked(
174 r,
175 reader_schema.as_ref(),
176 check.as_deref(),
177 chunk,
178 f,
179 )
180 });
181 self.finish_avro(name, reader_schema, result)
182 }
183 #[cfg(feature = "file-format-orc")]
184 FileFormat::Orc => {
185 let batch = if chunk == usize::MAX { 0 } else { chunk };
186 let mut pending: Vec<Value> = Vec::new();
187 self.orc_read(name, input, batch, &mut |b| {
188 let rows = crate::columnar::record_batch_to_values(&b)?;
189 if chunk == usize::MAX {
190 pending.extend(rows);
191 Ok(())
192 } else {
193 f(rows)
194 }
195 })?;
196 if !pending.is_empty() {
197 f(pending)?;
198 }
199 Ok(())
200 }
201 _ => unreachable!("rejected in new()"),
202 }
203 }
204
205 #[cfg(feature = "arrow")]
209 pub fn batches(
210 &mut self,
211 name: &str,
212 input: FileInput,
213 batch_size: usize,
214 f: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<(), FaucetError>,
215 ) -> Result<arrow::datatypes::SchemaRef, FaucetError> {
216 match self.format {
217 #[cfg(feature = "file-format-avro")]
218 FileFormat::Avro => {
219 let reader_schema = self.avro_reader_schema();
220 let check = self.avro_writer_check();
221 let result = with_reader(input, |r| {
222 super::avro::read_batches_checked(
223 r,
224 reader_schema.as_ref(),
225 check.as_deref(),
226 batch_size,
227 f,
228 )
229 });
230 let arrow = result.as_ref().ok().map(|(_, a)| a.clone());
231 self.finish_avro(name, reader_schema, result.map(|(s, _)| s))?;
232 Ok(arrow.expect("finish_avro succeeded only on Ok"))
233 }
234 #[cfg(feature = "file-format-orc")]
235 FileFormat::Orc => self.orc_read(name, input, batch_size, f),
236 _ => unreachable!("rejected in new()"),
237 }
238 }
239
240 #[cfg(feature = "file-format-avro")]
241 fn avro_reader_schema(&self) -> Option<apache_avro::Schema> {
242 match &self.anchor {
243 Some((_, Anchor::Avro(s))) => Some(s.clone()),
244 _ => None,
245 }
246 }
247
248 #[cfg(feature = "file-format-avro")]
253 #[allow(clippy::type_complexity)]
254 fn avro_writer_check(&self) -> Option<Box<super::avro::WriterCheck<'static>>> {
255 if self.configured {
256 return None;
257 }
258 let Some((_, Anchor::Avro(anchor))) = &self.anchor else {
259 return None;
260 };
261 let known: std::collections::HashSet<String> =
262 record_field_names(anchor).into_iter().collect();
263 Some(Box::new(move |writer: &apache_avro::Schema| {
264 let extra: Vec<String> = record_field_names(writer)
265 .into_iter()
266 .filter(|n| !known.contains(n))
267 .collect();
268 if extra.is_empty() {
269 Ok(())
270 } else {
271 Err(FaucetError::Source(format!(
272 "it has field(s) {} the reader schema lacks, which would be dropped — \
273 set `avro.schema` to a reader schema that includes them",
274 extra
275 .iter()
276 .map(|n| format!("`{n}`"))
277 .collect::<Vec<_>>()
278 .join(", ")
279 )))
280 }
281 }))
282 }
283
284 #[cfg(feature = "file-format-avro")]
285 fn finish_avro(
286 &mut self,
287 name: &str,
288 reader_schema: Option<apache_avro::Schema>,
289 result: Result<apache_avro::Schema, FaucetError>,
290 ) -> Result<(), FaucetError> {
291 match (result, &self.anchor) {
292 (Ok(schema), None) => {
293 self.anchor = Some((name.to_string(), Anchor::Avro(schema)));
294 Ok(())
295 }
296 (Ok(_), Some(_)) => Ok(()),
297 (Err(e), Some((first, _))) if reader_schema.is_some() => {
298 let against = if self.configured {
299 "the configured `avro.schema`".to_string()
300 } else {
301 format!("'{first}' (the first file's schema)")
302 };
303 Err(FaucetError::Source(format!(
304 "avro schema of '{name}' cannot be resolved against {against}: {e}"
305 )))
306 }
307 (Err(e), _) => Err(FaucetError::Source(format!("'{name}': {e}"))),
308 }
309 }
310
311 #[cfg(feature = "file-format-orc")]
312 fn orc_read(
313 &mut self,
314 name: &str,
315 input: FileInput,
316 batch_size: usize,
317 f: &mut dyn FnMut(arrow::array::RecordBatch) -> Result<(), FaucetError>,
318 ) -> Result<arrow::datatypes::SchemaRef, FaucetError> {
319 let input = match input {
320 FileInput::File(file) => super::orc::OrcInput::File(file),
321 FileInput::Bytes(b) => super::orc::OrcInput::Bytes(bytes::Bytes::from(b)),
322 };
323 let reference = match &self.anchor {
324 Some((first, Anchor::Orc(s))) => Some((first.clone(), s.clone())),
325 _ => None,
326 };
327 let mut check = |schema: &arrow::datatypes::SchemaRef| -> Result<(), FaucetError> {
328 match &reference {
329 Some((first, reference)) if !same_shape(reference, schema) => {
330 Err(schema_conflict(name, first, reference, schema))
331 }
332 _ => Ok(()),
333 }
334 };
335 let schema =
336 super::orc::read_batches_checked(input, &self.opts.orc, batch_size, &mut check, f)
337 .map_err(|e| prefix(name, e))?;
338 if self.anchor.is_none() {
339 self.anchor = Some((name.to_string(), Anchor::Orc(schema.clone())));
340 }
341 Ok(schema)
342 }
343}
344
345#[cfg(feature = "arrow")]
351pub fn columnar_pages<'a, F, Fut>(
352 names: Vec<String>,
353 concurrency: usize,
354 mut decoder: ContainerDecoder,
355 batch_size: usize,
356 fetch: F,
357) -> std::pin::Pin<
358 Box<dyn futures::Stream<Item = Result<crate::columnar::ColumnarPage, FaucetError>> + Send + 'a>,
359>
360where
361 F: Fn(String) -> Fut + Send + Sync + 'a,
362 Fut: std::future::Future<Output = Result<Vec<u8>, FaucetError>> + Send + 'a,
363{
364 use futures::StreamExt as _;
365 Box::pin(async_stream::try_stream! {
366 let fetch = &fetch;
367 let mut fetched = futures::stream::iter(names)
368 .map(|name| async move {
369 let body = fetch(name.clone()).await;
370 (name, body)
371 })
372 .buffered(concurrency.max(1));
373 while let Some((name, body)) = fetched.next().await {
374 let (d, batches) = decoder
375 .decode_batches_offloaded(name, FileInput::Bytes(body?), batch_size)
376 .await?;
377 decoder = d;
378 for batch in batches {
379 yield crate::columnar::ColumnarPage::new(batch, None);
380 }
381 }
382 })
383}
384
385#[cfg(feature = "file-format-avro")]
386fn record_field_names(schema: &apache_avro::Schema) -> Vec<String> {
387 match schema {
388 apache_avro::Schema::Record(r) => r.fields.iter().map(|f| f.name.clone()).collect(),
389 _ => Vec::new(),
390 }
391}
392
393#[cfg(feature = "file-format-avro")]
394fn with_reader<T>(
395 input: FileInput,
396 f: impl FnOnce(&mut dyn std::io::Read) -> Result<T, FaucetError>,
397) -> Result<T, FaucetError> {
398 match input {
399 FileInput::File(file) => f(&mut std::io::BufReader::new(file)),
400 FileInput::Bytes(b) => f(&mut &b[..]),
401 }
402}
403
404#[cfg(feature = "file-format-orc")]
405fn prefix(name: &str, e: FaucetError) -> FaucetError {
406 match e {
407 FaucetError::Source(m) if !m.starts_with('\'') => {
408 FaucetError::Source(format!("'{name}': {m}"))
409 }
410 other => other,
411 }
412}
413
414#[cfg(feature = "arrow")]
418pub fn same_shape(a: &arrow::datatypes::Schema, b: &arrow::datatypes::Schema) -> bool {
419 a.fields().len() == b.fields().len()
420 && a.fields()
421 .iter()
422 .zip(b.fields().iter())
423 .all(|(x, y)| x.name() == y.name() && x.data_type() == y.data_type())
424}
425
426#[cfg(feature = "arrow")]
429pub fn schema_conflict(
430 file: &str,
431 first: &str,
432 reference: &arrow::datatypes::Schema,
433 schema: &arrow::datatypes::Schema,
434) -> FaucetError {
435 let detail = reference
436 .fields()
437 .iter()
438 .zip(schema.fields().iter())
439 .find(|(a, b)| a.name() != b.name() || a.data_type() != b.data_type())
440 .map(|(a, b)| {
441 format!(
442 "field `{}` ({}) vs `{}` ({})",
443 a.name(),
444 a.data_type(),
445 b.name(),
446 b.data_type()
447 )
448 })
449 .unwrap_or_else(|| {
450 format!(
451 "{} vs {} fields",
452 reference.fields().len(),
453 schema.fields().len()
454 )
455 });
456 FaucetError::Source(format!(
457 "schema of '{file}' conflicts with '{first}' (the first file's schema): {detail}"
458 ))
459}
460
461#[cfg(test)]
462mod tests {
463 use super::*;
464 #[cfg(feature = "file-format-avro")]
465 use serde_json::json;
466
467 #[cfg(feature = "file-format-avro")]
468 fn avro(records: &[Value]) -> Vec<u8> {
469 super::super::avro::encode(records, &Default::default()).expect("encode")
470 }
471
472 #[cfg(all(feature = "file-format-avro", feature = "arrow"))]
473 #[tokio::test]
474 async fn columnar_pages_decode_in_listing_order() {
475 use futures::StreamExt as _;
476 let bodies: std::collections::HashMap<String, Vec<u8>> = [
477 ("a".to_string(), avro(&[json!({"id": 1}), json!({"id": 2})])),
478 ("b".to_string(), avro(&[json!({"id": 3})])),
479 ]
480 .into_iter()
481 .collect();
482 let decoder = ContainerDecoder::new(FileFormat::Avro, &FormatOptions::default()).unwrap();
483 let pages: Vec<_> = columnar_pages(
484 vec!["a".into(), "b".into(), "missing".into()],
485 2,
486 decoder,
487 1,
488 |n| {
489 let body = bodies.get(&n).cloned();
490 async move { body.ok_or_else(|| FaucetError::Source(format!("no {n}"))) }
491 },
492 )
493 .collect()
494 .await;
495 assert_eq!(pages.len(), 4);
496 assert_eq!(
497 pages[..3]
498 .iter()
499 .map(|p| p.as_ref().unwrap().num_rows())
500 .sum::<usize>(),
501 3
502 );
503 assert!(pages[3].as_ref().is_err());
504 }
505
506 #[cfg(all(feature = "file-format-avro", feature = "arrow"))]
507 #[tokio::test]
508 async fn offloaded_decodes_keep_the_anchor() {
509 let d = ContainerDecoder::new(FileFormat::Avro, &FormatOptions::default()).unwrap();
510 let (d, rows) = d
511 .decode_all_offloaded("a".into(), FileInput::Bytes(avro(&[json!({"id": 1})])))
512 .await
513 .unwrap();
514 assert_eq!(rows, vec![json!({"id": 1})]);
515 let (d, batches) = d
516 .decode_batches_offloaded("b".into(), FileInput::Bytes(avro(&[json!({"id": 2})])), 0)
517 .await
518 .unwrap();
519 assert_eq!(batches[0].num_rows(), 1);
520 let err = d
521 .decode_all_offloaded(
522 "c".into(),
523 FileInput::Bytes(avro(&[json!({"id": 1, "z": 1})])),
524 )
525 .await
526 .expect_err("anchored on a");
527 assert!(err.to_string().contains("`z`"), "{err}");
528 }
529
530 #[test]
531 fn only_container_formats_are_accepted() {
532 let err =
533 ContainerDecoder::new(FileFormat::Csv, &FormatOptions::default()).expect_err("csv");
534 assert!(err.to_string().contains("avro and orc"), "{err}");
535 }
536
537 #[cfg(feature = "file-format-avro")]
538 #[test]
539 fn later_avro_files_resolve_against_the_first() {
540 let mut d = ContainerDecoder::new(FileFormat::Avro, &FormatOptions::default()).unwrap();
541 assert_eq!(d.format(), FileFormat::Avro);
542 let mut rows = Vec::new();
543 d.records(
544 "a.avro",
545 FileInput::Bytes(avro(&[json!({"id": 1, "x": "a"})])),
546 0,
547 &mut |c| {
548 rows.extend(c);
549 Ok(())
550 },
551 )
552 .unwrap();
553 let wider = avro(&[json!({"id": 2, "x": "b", "extra": true})]);
556 let err = d
557 .records("b.avro", FileInput::Bytes(wider), 1, &mut |c| {
558 rows.extend(c);
559 Ok(())
560 })
561 .expect_err("added field");
562 let msg = err.to_string();
563 assert!(msg.contains("b.avro") && msg.contains("`extra`"), "{msg}");
564 d.records(
566 "b2.avro",
567 FileInput::Bytes(avro(&[json!({"id": 2, "x": "b"})])),
568 1,
569 &mut |c| {
570 rows.extend(c);
571 Ok(())
572 },
573 )
574 .unwrap();
575 assert_eq!(
576 rows,
577 vec![json!({"id": 1, "x": "a"}), json!({"id": 2, "x": "b"})]
578 );
579 let err = d
580 .records(
581 "c.avro",
582 FileInput::Bytes(avro(&[json!({"id": "text"})])),
583 0,
584 &mut |_| Ok(()),
585 )
586 .expect_err("conflict");
587 let msg = err.to_string();
588 assert!(msg.contains("c.avro") && msg.contains("a.avro"), "{msg}");
589 }
590
591 #[cfg(feature = "file-format-avro")]
592 #[test]
593 fn a_configured_reader_schema_is_named_in_conflicts() {
594 let opts = FormatOptions {
595 avro: super::super::AvroOptions {
596 schema: Some(
597 json!({"type": "record", "name": "faucet_record", "fields": [
598 {"name": "id", "type": "long"}
599 ]}),
600 ),
601 ..Default::default()
602 },
603 ..Default::default()
604 };
605 let mut d = ContainerDecoder::new(FileFormat::Avro, &opts).unwrap();
606 let err = d
607 .records(
608 "x.avro",
609 FileInput::Bytes(avro(&[json!({"other": 1})])),
610 0,
611 &mut |_| Ok(()),
612 )
613 .expect_err("unresolvable");
614 assert!(err.to_string().contains("configured"), "{err}");
615 let err = d
616 .records("y.avro", FileInput::Bytes(b"junk".to_vec()), 0, &mut |_| {
617 Ok(())
618 })
619 .expect_err("junk");
620 assert!(err.to_string().contains("y.avro"), "{err}");
621 }
622
623 #[cfg(all(feature = "file-format-avro", feature = "arrow"))]
624 #[test]
625 fn avro_batches_and_local_files() {
626 let dir = tempfile::tempdir().unwrap();
627 let path = dir.path().join("a.avro");
628 std::fs::write(&path, avro(&[json!({"id": 1}), json!({"id": 2})])).unwrap();
629 let mut d = ContainerDecoder::new(FileFormat::Avro, &FormatOptions::default()).unwrap();
630 let mut n = 0;
631 let schema = d
632 .batches(
633 "a.avro",
634 FileInput::File(std::fs::File::open(&path).unwrap()),
635 1,
636 &mut |b| {
637 n += b.num_rows();
638 Ok(())
639 },
640 )
641 .unwrap();
642 assert_eq!(n, 2);
643 assert_eq!(schema.field(0).name(), "id");
644 }
645
646 #[cfg(feature = "file-format-orc")]
647 #[test]
648 fn orc_files_must_share_a_schema() {
649 const FIXTURE: &[u8] = include_bytes!("../../tests/fixtures/orc/people.orc");
650 let mut d = ContainerDecoder::new(FileFormat::Orc, &FormatOptions::default()).unwrap();
651 let mut rows = Vec::new();
652 d.records("a.orc", FileInput::Bytes(FIXTURE.to_vec()), 0, &mut |c| {
653 rows.extend(c);
654 Ok(())
655 })
656 .unwrap();
657 assert_eq!(rows.len(), 3);
658 let mut chunks = 0;
659 d.records("b.orc", FileInput::Bytes(FIXTURE.to_vec()), 2, &mut |_| {
660 chunks += 1;
661 Ok(())
662 })
663 .unwrap();
664 assert_eq!(chunks, 2);
665 let mut projected = ContainerDecoder::new(
666 FileFormat::Orc,
667 &FormatOptions {
668 orc: super::super::OrcOptions {
669 columns: Some(vec!["id".into()]),
670 },
671 ..Default::default()
672 },
673 )
674 .unwrap();
675 let s = projected
676 .batches("p.orc", FileInput::Bytes(FIXTURE.to_vec()), 0, &mut |_| {
677 Ok(())
678 })
679 .unwrap();
680 let full = d.anchor.as_ref().map(|(_, a)| match a {
682 Anchor::Orc(s) => s.clone(),
683 #[allow(unreachable_patterns)]
684 _ => unreachable!(),
685 });
686 let err = schema_conflict("p.orc", "a.orc", &full.unwrap(), &s);
687 assert!(err.to_string().contains("p.orc") && err.to_string().contains("a.orc"));
688 d.anchor = Some(("a.orc".into(), Anchor::Orc(s)));
689 let err = d
690 .records("c.orc", FileInput::Bytes(FIXTURE.to_vec()), 0, &mut |_| {
691 Ok(())
692 })
693 .expect_err("conflict");
694 assert!(err.to_string().contains("c.orc"), "{err}");
695 let err = d
696 .records("d.orc", FileInput::Bytes(b"junk".to_vec()), 0, &mut |_| {
697 Ok(())
698 })
699 .expect_err("junk");
700 assert!(err.to_string().contains("d.orc"), "{err}");
701 }
702
703 #[cfg(feature = "arrow")]
704 #[test]
705 fn schema_conflict_reports_a_field_count_difference() {
706 use arrow::datatypes::{DataType, Field, Schema};
707 let a = Schema::new(vec![Field::new("x", DataType::Int64, true)]);
708 let b = Schema::new(vec![
709 Field::new("x", DataType::Int64, true),
710 Field::new("y", DataType::Int64, true),
711 ]);
712 let msg = schema_conflict("b", "a", &a, &b).to_string();
713 assert!(msg.contains("1 vs 2"), "{msg}");
714 }
715}