1use std::collections::HashMap;
23use std::fmt::{self, Debug};
24use std::io::{Seek, SeekFrom};
25use std::sync::Arc;
26
27use arrow::datatypes::{Schema, SchemaRef};
28use arrow::error::ArrowError;
29use arrow::ipc::convert::fb_to_schema;
30use arrow::ipc::reader::{FileReader, StreamReader};
31use arrow::ipc::writer::IpcWriteOptions;
32use arrow::ipc::{CompressionType, root_as_message};
33use datafusion_common::error::Result;
34use datafusion_common::parsers::CompressionTypeVariant;
35use datafusion_common::{
36 DEFAULT_ARROW_EXTENSION, DataFusionError, GetExt, Statistics,
37 internal_datafusion_err, not_impl_err,
38};
39use datafusion_common_runtime::{JoinSet, SpawnedTask};
40use datafusion_datasource::TableSchema;
41use datafusion_datasource::display::FileGroupDisplay;
42use datafusion_datasource::file::FileSource;
43use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
44use datafusion_datasource::sink::{DataSink, DataSinkExec};
45use datafusion_datasource::write::{
46 ObjectWriterBuilder, SharedBuffer, get_writer_schema,
47};
48use datafusion_execution::{SendableRecordBatchStream, TaskContext};
49use datafusion_expr::dml::InsertOp;
50use datafusion_physical_expr_common::sort_expr::LexRequirement;
51
52use crate::source::ArrowSource;
53use async_trait::async_trait;
54use bytes::Bytes;
55use datafusion_datasource::file_compression_type::FileCompressionType;
56use datafusion_datasource::file_format::{FileFormat, FileFormatFactory};
57use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig};
58use datafusion_datasource::source::DataSourceExec;
59use datafusion_datasource::write::demux::DemuxedStreamReceiver;
60use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan};
61use datafusion_session::Session;
62use futures::StreamExt;
63use futures::stream::BoxStream;
64use object_store::{
65 GetOptions, GetRange, GetResultPayload, ObjectMeta, ObjectStore, ObjectStoreExt,
66 path::Path,
67};
68use tokio::io::AsyncWriteExt;
69
70const INITIAL_BUFFER_BYTES: usize = 1048576;
73
74const BUFFER_FLUSH_BYTES: usize = 1024000;
76
77#[derive(Default, Debug)]
79pub struct ArrowFormatFactory;
80
81impl ArrowFormatFactory {
82 pub fn new() -> Self {
84 Self {}
85 }
86}
87
88impl FileFormatFactory for ArrowFormatFactory {
89 fn create(
90 &self,
91 _state: &dyn Session,
92 _format_options: &HashMap<String, String>,
93 ) -> Result<Arc<dyn FileFormat>> {
94 Ok(Arc::new(ArrowFormat))
95 }
96
97 fn default(&self) -> Arc<dyn FileFormat> {
98 Arc::new(ArrowFormat)
99 }
100}
101
102impl GetExt for ArrowFormatFactory {
103 fn get_ext(&self) -> String {
104 DEFAULT_ARROW_EXTENSION[1..].to_string()
106 }
107}
108
109#[derive(Default, Debug)]
111pub struct ArrowFormat;
112
113#[async_trait]
114impl FileFormat for ArrowFormat {
115 fn get_ext(&self) -> String {
116 ArrowFormatFactory::new().get_ext()
117 }
118
119 fn get_ext_with_compression(
120 &self,
121 file_compression_type: &FileCompressionType,
122 ) -> Result<String> {
123 let ext = self.get_ext();
124 match file_compression_type.get_variant() {
125 CompressionTypeVariant::UNCOMPRESSED => Ok(ext),
126 _ => Err(internal_datafusion_err!(
127 "Arrow FileFormat does not support compression."
128 )),
129 }
130 }
131
132 fn compression_type(&self) -> Option<FileCompressionType> {
133 None
134 }
135
136 async fn infer_schema(
137 &self,
138 _state: &dyn Session,
139 store: &Arc<dyn ObjectStore>,
140 objects: &[ObjectMeta],
141 ) -> Result<SchemaRef> {
142 let mut schemas = vec![];
143 for object in objects {
144 let r = store.as_ref().get(&object.location).await?;
145 let schema = match r.payload {
146 #[cfg(not(target_arch = "wasm32"))]
147 GetResultPayload::File(mut file, _) => {
148 match FileReader::try_new(&mut file, None) {
149 Ok(reader) => reader.schema(),
150 Err(file_error) => {
151 file.seek(SeekFrom::Start(0))?;
155 match StreamReader::try_new(&mut file, None) {
156 Ok(reader) => reader.schema(),
157 Err(stream_error) => {
158 return Err(internal_datafusion_err!(
159 "Failed to parse Arrow file as either file format or stream format. File format error: {file_error}. Stream format error: {stream_error}"
160 ));
161 }
162 }
163 }
164 }
165 }
166 GetResultPayload::Stream(stream) => infer_stream_schema(stream).await?,
167 };
168 schemas.push(Arc::unwrap_or_clone(schema));
169 }
170 let merged_schema = Schema::try_merge(schemas)?;
171 Ok(Arc::new(merged_schema))
172 }
173
174 async fn infer_stats(
175 &self,
176 _state: &dyn Session,
177 _store: &Arc<dyn ObjectStore>,
178 table_schema: SchemaRef,
179 _object: &ObjectMeta,
180 ) -> Result<Statistics> {
181 Ok(Statistics::new_unknown(&table_schema))
182 }
183
184 async fn create_physical_plan(
185 &self,
186 state: &dyn Session,
187 conf: FileScanConfig,
188 ) -> Result<Arc<dyn ExecutionPlan>> {
189 let object_store = state.runtime_env().object_store(&conf.object_store_url)?;
190 let object_location = &conf
191 .file_groups
192 .first()
193 .ok_or_else(|| internal_datafusion_err!("No files found in file group"))?
194 .files()
195 .first()
196 .ok_or_else(|| internal_datafusion_err!("No files found in file group"))?
197 .object_meta
198 .location;
199
200 let table_schema = TableSchema::new(
201 Arc::clone(conf.file_schema()),
202 conf.table_partition_cols().clone(),
203 );
204
205 let mut source: Arc<dyn FileSource> =
206 match is_object_in_arrow_ipc_file_format(object_store, object_location).await
207 {
208 Ok(true) => Arc::new(ArrowSource::new_file_source(table_schema)),
209 Ok(false) => Arc::new(ArrowSource::new_stream_file_source(table_schema)),
210 Err(e) => Err(e)?,
211 };
212
213 if let Some(projection) = conf.file_source.projection()
215 && let Some(new_source) = source.try_pushdown_projection(projection)?
216 {
217 source = new_source;
218 }
219
220 let config = FileScanConfigBuilder::from(conf)
221 .with_source(source)
222 .build();
223
224 Ok(DataSourceExec::from_data_source(config))
225 }
226
227 async fn create_writer_physical_plan(
228 &self,
229 input: Arc<dyn ExecutionPlan>,
230 _state: &dyn Session,
231 conf: FileSinkConfig,
232 order_requirements: Option<LexRequirement>,
233 ) -> Result<Arc<dyn ExecutionPlan>> {
234 if conf.insert_op != InsertOp::Append {
235 return not_impl_err!("Overwrites are not implemented yet for Arrow format");
236 }
237
238 let sink = Arc::new(ArrowFileSink::new(conf));
239
240 Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
241 }
242
243 fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
244 Arc::new(ArrowSource::new_file_source(table_schema))
245 }
246}
247
248struct ArrowFileSink {
250 config: FileSinkConfig,
251}
252
253impl ArrowFileSink {
254 fn new(config: FileSinkConfig) -> Self {
255 Self { config }
256 }
257}
258
259#[async_trait]
260impl FileSink for ArrowFileSink {
261 fn config(&self) -> &FileSinkConfig {
262 &self.config
263 }
264
265 async fn spawn_writer_tasks_and_join(
266 &self,
267 context: &Arc<TaskContext>,
268 demux_task: SpawnedTask<Result<()>>,
269 mut file_stream_rx: DemuxedStreamReceiver,
270 object_store: Arc<dyn ObjectStore>,
271 ) -> Result<u64> {
272 let mut file_write_tasks: JoinSet<std::result::Result<usize, DataFusionError>> =
273 JoinSet::new();
274
275 let ipc_options =
276 IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)?
277 .try_with_compression(Some(CompressionType::LZ4_FRAME))?;
278 while let Some((path, mut rx)) = file_stream_rx.recv().await {
279 let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES);
280 let mut arrow_writer = arrow_ipc::writer::FileWriter::try_new_with_options(
281 shared_buffer.clone(),
282 &get_writer_schema(&self.config),
283 ipc_options.clone(),
284 )?;
285 let mut object_store_writer = ObjectWriterBuilder::new(
286 FileCompressionType::UNCOMPRESSED,
287 &path,
288 Arc::clone(&object_store),
289 )
290 .with_buffer_size(Some(
291 context
292 .session_config()
293 .options()
294 .execution
295 .objectstore_writer_buffer_size,
296 ))
297 .build()?;
298 file_write_tasks.spawn(async move {
299 let mut row_count = 0;
300 while let Some(batch) = rx.recv().await {
301 row_count += batch.num_rows();
302 arrow_writer.write(&batch)?;
303 let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap();
304 if buff_to_flush.len() > BUFFER_FLUSH_BYTES {
305 object_store_writer
306 .write_all(buff_to_flush.as_slice())
307 .await?;
308 buff_to_flush.clear();
309 }
310 }
311 arrow_writer.finish()?;
312 let final_buff = shared_buffer.buffer.try_lock().unwrap();
313
314 object_store_writer.write_all(final_buff.as_slice()).await?;
315 object_store_writer.shutdown().await?;
316 Ok(row_count)
317 });
318 }
319
320 let mut row_count = 0;
321 while let Some(result) = file_write_tasks.join_next().await {
322 match result {
323 Ok(r) => {
324 row_count += r?;
325 }
326 Err(e) => {
327 if e.is_panic() {
328 std::panic::resume_unwind(e.into_panic());
329 } else {
330 unreachable!();
331 }
332 }
333 }
334 }
335
336 demux_task
337 .join_unwind()
338 .await
339 .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
340 Ok(row_count as u64)
341 }
342}
343
344impl Debug for ArrowFileSink {
345 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
346 f.debug_struct("ArrowFileSink").finish()
347 }
348}
349
350impl DisplayAs for ArrowFileSink {
351 fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
352 match t {
353 DisplayFormatType::Default | DisplayFormatType::Verbose => {
354 write!(f, "ArrowFileSink(file_groups=",)?;
355 FileGroupDisplay(&self.config.file_group).fmt_as(t, f)?;
356 write!(f, ")")
357 }
358 DisplayFormatType::TreeRender => {
359 writeln!(f, "format: arrow")?;
360 write!(f, "file={}", &self.config.original_url)
361 }
362 }
363 }
364}
365
366#[async_trait]
367impl DataSink for ArrowFileSink {
368 fn schema(&self) -> &SchemaRef {
369 self.config.output_schema()
370 }
371
372 async fn write_all(
373 &self,
374 data: SendableRecordBatchStream,
375 context: &Arc<TaskContext>,
376 ) -> Result<u64> {
377 FileSink::write_all(self, data, context).await
378 }
379}
380
381const ARROW_MAGIC: [u8; 6] = [b'A', b'R', b'R', b'O', b'W', b'1'];
385const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
386
387async fn infer_stream_schema(
388 mut stream: BoxStream<'static, object_store::Result<Bytes>>,
389) -> Result<SchemaRef> {
390 let bytes = extend_bytes_to_n_length_from_stream(vec![], 16, &mut stream).await?;
431
432 let preamble_len = if bytes[0..6] == ARROW_MAGIC {
434 if bytes[8..12] == CONTINUATION_MARKER {
436 12
438 } else {
439 8
441 }
442 } else if bytes[0..4] == CONTINUATION_MARKER {
443 4
445 } else {
446 0
448 };
449
450 let meta_len_bytes: [u8; 4] = bytes[preamble_len..preamble_len + 4]
451 .try_into()
452 .map_err(|err| {
453 ArrowError::ParseError(format!(
454 "Unable to read IPC message metadata length: {err:?}"
455 ))
456 })?;
457
458 let meta_len = i32::from_le_bytes([
459 meta_len_bytes[0],
460 meta_len_bytes[1],
461 meta_len_bytes[2],
462 meta_len_bytes[3],
463 ]);
464
465 if meta_len < 0 {
466 return Err(ArrowError::ParseError(
467 "IPC message metadata length is negative".to_string(),
468 )
469 .into());
470 }
471
472 let bytes = extend_bytes_to_n_length_from_stream(
473 bytes,
474 preamble_len + 4 + (meta_len as usize),
475 &mut stream,
476 )
477 .await?;
478
479 let message = root_as_message(&bytes[preamble_len + 4..]).map_err(|err| {
480 ArrowError::ParseError(format!("Unable to read IPC message metadata: {err:?}"))
481 })?;
482 let fb_schema = message.header_as_schema().ok_or_else(|| {
483 ArrowError::IpcError("Unable to read IPC message schema".to_string())
484 })?;
485 let schema = fb_to_schema(fb_schema);
486
487 Ok(Arc::new(schema))
488}
489
490async fn extend_bytes_to_n_length_from_stream(
491 bytes: Vec<u8>,
492 n: usize,
493 stream: &mut BoxStream<'static, object_store::Result<Bytes>>,
494) -> Result<Vec<u8>> {
495 if bytes.len() >= n {
496 return Ok(bytes);
497 }
498
499 let mut buf = bytes;
500
501 while let Some(b) = stream.next().await.transpose()? {
502 buf.extend_from_slice(&b);
503
504 if buf.len() >= n {
505 break;
506 }
507 }
508
509 if buf.len() < n {
510 return Err(ArrowError::ParseError(
511 "Unexpected end of byte stream for Arrow IPC file".to_string(),
512 )
513 .into());
514 }
515
516 Ok(buf)
517}
518
519async fn is_object_in_arrow_ipc_file_format(
520 store: Arc<dyn ObjectStore>,
521 object_location: &Path,
522) -> Result<bool> {
523 let get_opts = GetOptions {
524 range: Some(GetRange::Bounded(0..6)),
525 ..Default::default()
526 };
527 let bytes = store
528 .get_opts(object_location, get_opts)
529 .await?
530 .bytes()
531 .await?;
532 Ok(bytes.len() >= 6 && bytes[0..6] == ARROW_MAGIC)
533}
534
535#[cfg(test)]
536mod tests {
537 use super::*;
538
539 use std::any::Any;
540
541 use chrono::DateTime;
542 use datafusion_common::DFSchema;
543 use datafusion_common::config::TableOptions;
544 use datafusion_execution::config::SessionConfig;
545 use datafusion_execution::runtime_env::RuntimeEnv;
546 use datafusion_expr::execution_props::ExecutionProps;
547 use datafusion_expr::registry::ExtensionTypeRegistryRef;
548 use datafusion_expr::{
549 AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
550 };
551 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
552 use object_store::{chunked::ChunkedStore, memory::InMemory};
553
554 struct MockSession {
555 config: SessionConfig,
556 runtime_env: Arc<RuntimeEnv>,
557 }
558
559 impl MockSession {
560 fn new() -> Self {
561 Self {
562 config: SessionConfig::new(),
563 runtime_env: Arc::new(RuntimeEnv::default()),
564 }
565 }
566 }
567
568 #[async_trait::async_trait]
569 impl Session for MockSession {
570 fn session_id(&self) -> &str {
571 unimplemented!()
572 }
573
574 fn config(&self) -> &SessionConfig {
575 &self.config
576 }
577
578 async fn create_physical_plan(
579 &self,
580 _logical_plan: &LogicalPlan,
581 ) -> Result<Arc<dyn ExecutionPlan>> {
582 unimplemented!()
583 }
584
585 fn create_physical_expr(
586 &self,
587 _expr: Expr,
588 _df_schema: &DFSchema,
589 ) -> Result<Arc<dyn PhysicalExpr>> {
590 unimplemented!()
591 }
592
593 fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> {
594 unimplemented!()
595 }
596
597 fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>> {
598 unimplemented!()
599 }
600
601 fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
602 unimplemented!()
603 }
604
605 fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
606 unimplemented!()
607 }
608
609 fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef {
610 unimplemented!()
611 }
612
613 fn runtime_env(&self) -> &Arc<RuntimeEnv> {
614 &self.runtime_env
615 }
616
617 fn execution_props(&self) -> &ExecutionProps {
618 unimplemented!()
619 }
620
621 fn as_any(&self) -> &dyn Any {
622 unimplemented!()
623 }
624
625 fn table_options(&self) -> &TableOptions {
626 unimplemented!()
627 }
628
629 fn table_options_mut(&mut self) -> &mut TableOptions {
630 unimplemented!()
631 }
632
633 fn task_ctx(&self) -> Arc<TaskContext> {
634 unimplemented!()
635 }
636 }
637
638 #[tokio::test]
639 async fn test_infer_schema_stream() -> Result<()> {
640 for file in ["example.arrow", "example_stream.arrow"] {
641 let mut bytes = std::fs::read(format!("tests/data/{file}"))?;
642 bytes.truncate(bytes.len() - 20); let location = Path::parse(file)?;
644 let in_memory_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
645 in_memory_store.put(&location, bytes.into()).await?;
646
647 let state = MockSession::new();
648 let object_meta = ObjectMeta {
649 location,
650 last_modified: DateTime::default(),
651 size: u64::MAX,
652 e_tag: None,
653 version: None,
654 };
655
656 let arrow_format = ArrowFormat {};
657 let expected = vec!["f0: Int64", "f1: Utf8", "f2: Boolean"];
658
659 for chunk_size in [7, 3000] {
662 let store =
663 Arc::new(ChunkedStore::new(in_memory_store.clone(), chunk_size));
664 let inferred_schema = arrow_format
665 .infer_schema(
666 &state,
667 &(store.clone() as Arc<dyn ObjectStore>),
668 std::slice::from_ref(&object_meta),
669 )
670 .await?;
671 let actual_fields = inferred_schema
672 .fields()
673 .iter()
674 .map(|f| format!("{}: {:?}", f.name(), f.data_type()))
675 .collect::<Vec<_>>();
676 assert_eq!(expected, actual_fields);
677 }
678 }
679 Ok(())
680 }
681
682 #[tokio::test]
683 async fn test_infer_schema_short_stream() -> Result<()> {
684 for file in ["example.arrow", "example_stream.arrow"] {
685 let mut bytes = std::fs::read(format!("tests/data/{file}"))?;
686 bytes.truncate(20); let location = Path::parse(file)?;
688 let in_memory_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
689 in_memory_store.put(&location, bytes.into()).await?;
690
691 let state = MockSession::new();
692 let object_meta = ObjectMeta {
693 location,
694 last_modified: DateTime::default(),
695 size: u64::MAX,
696 e_tag: None,
697 version: None,
698 };
699
700 let arrow_format = ArrowFormat {};
701
702 let store = Arc::new(ChunkedStore::new(in_memory_store.clone(), 7));
703 let err = arrow_format
704 .infer_schema(
705 &state,
706 &(store.clone() as Arc<dyn ObjectStore>),
707 std::slice::from_ref(&object_meta),
708 )
709 .await;
710
711 assert!(err.is_err());
712 assert_eq!(
713 "Arrow error: Parser error: Unexpected end of byte stream for Arrow IPC file",
714 err.unwrap_err().to_string().lines().next().unwrap()
715 );
716 }
717
718 Ok(())
719 }
720
721 #[tokio::test]
722 async fn test_format_detection_file_format() -> Result<()> {
723 let store = Arc::new(InMemory::new());
724 let path = Path::from("test.arrow");
725
726 let file_bytes = std::fs::read("tests/data/example.arrow")?;
727 store.put(&path, file_bytes.into()).await?;
728
729 let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?;
730 assert!(is_file, "Should detect file format");
731 Ok(())
732 }
733
734 #[tokio::test]
735 async fn test_format_detection_stream_format() -> Result<()> {
736 let store = Arc::new(InMemory::new());
737 let path = Path::from("test_stream.arrow");
738
739 let stream_bytes = std::fs::read("tests/data/example_stream.arrow")?;
740 store.put(&path, stream_bytes.into()).await?;
741
742 let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?;
743
744 assert!(!is_file, "Should detect stream format (not file)");
745
746 Ok(())
747 }
748
749 #[tokio::test]
750 async fn test_format_detection_corrupted_file() -> Result<()> {
751 let store = Arc::new(InMemory::new());
752 let path = Path::from("corrupted.arrow");
753
754 store
755 .put(&path, Bytes::from(vec![0x43, 0x4f, 0x52, 0x41]).into())
756 .await?;
757
758 let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?;
759
760 assert!(
761 !is_file,
762 "Corrupted file should not be detected as file format"
763 );
764
765 Ok(())
766 }
767
768 #[tokio::test]
769 async fn test_format_detection_empty_file() -> Result<()> {
770 let store = Arc::new(InMemory::new());
771 let path = Path::from("empty.arrow");
772
773 store.put(&path, Bytes::new().into()).await?;
774
775 let result = is_object_in_arrow_ipc_file_format(store.clone(), &path).await;
776
777 assert!(result.is_err(), "Empty file should error");
779
780 Ok(())
781 }
782}