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::display::FileGroupDisplay;
41use datafusion_datasource::file::FileSource;
42use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder};
43use datafusion_datasource::sink::{DataSink, DataSinkExec};
44use datafusion_datasource::write::{
45 ObjectWriterBuilder, SharedBuffer, get_writer_schema,
46};
47use datafusion_datasource::{TableSchema, TableSchemaBuilder};
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 = TableSchemaBuilder::from(conf.file_schema())
201 .with_table_partition_cols(conf.table_partition_cols().clone())
202 .build();
203
204 let mut source: Arc<dyn FileSource> =
205 match is_object_in_arrow_ipc_file_format(object_store, object_location).await
206 {
207 Ok(true) => Arc::new(ArrowSource::new_file_source(table_schema)),
208 Ok(false) => Arc::new(ArrowSource::new_stream_file_source(table_schema)),
209 Err(e) => Err(e)?,
210 };
211
212 if let Some(projection) = conf.file_source.projection()
214 && let Some(new_source) = source.try_pushdown_projection(projection)?
215 {
216 source = new_source;
217 }
218
219 let config = FileScanConfigBuilder::from(conf)
220 .with_source(source)
221 .build();
222
223 Ok(DataSourceExec::from_data_source(config))
224 }
225
226 async fn create_writer_physical_plan(
227 &self,
228 input: Arc<dyn ExecutionPlan>,
229 _state: &dyn Session,
230 conf: FileSinkConfig,
231 order_requirements: Option<LexRequirement>,
232 ) -> Result<Arc<dyn ExecutionPlan>> {
233 if conf.insert_op != InsertOp::Append {
234 return not_impl_err!("Overwrites are not implemented yet for Arrow format");
235 }
236
237 let sink = Arc::new(ArrowFileSink::new(conf));
238
239 Ok(Arc::new(DataSinkExec::new(input, sink, order_requirements)) as _)
240 }
241
242 fn file_source(&self, table_schema: TableSchema) -> Arc<dyn FileSource> {
243 Arc::new(ArrowSource::new_file_source(table_schema))
244 }
245}
246
247struct ArrowFileSink {
249 config: FileSinkConfig,
250}
251
252impl ArrowFileSink {
253 fn new(config: FileSinkConfig) -> Self {
254 Self { config }
255 }
256}
257
258#[async_trait]
259impl FileSink for ArrowFileSink {
260 fn config(&self) -> &FileSinkConfig {
261 &self.config
262 }
263
264 async fn spawn_writer_tasks_and_join(
265 &self,
266 context: &Arc<TaskContext>,
267 demux_task: SpawnedTask<Result<()>>,
268 mut file_stream_rx: DemuxedStreamReceiver,
269 object_store: Arc<dyn ObjectStore>,
270 ) -> Result<u64> {
271 let mut file_write_tasks: JoinSet<std::result::Result<usize, DataFusionError>> =
272 JoinSet::new();
273
274 let ipc_options =
275 IpcWriteOptions::try_new(64, false, arrow_ipc::MetadataVersion::V5)?
276 .try_with_compression(Some(CompressionType::LZ4_FRAME))?;
277 while let Some((path, mut rx)) = file_stream_rx.recv().await {
278 let shared_buffer = SharedBuffer::new(INITIAL_BUFFER_BYTES);
279 let mut arrow_writer = arrow_ipc::writer::FileWriter::try_new_with_options(
280 shared_buffer.clone(),
281 &get_writer_schema(&self.config),
282 ipc_options.clone(),
283 )?;
284 let mut object_store_writer = ObjectWriterBuilder::new(
285 FileCompressionType::UNCOMPRESSED,
286 &path,
287 Arc::clone(&object_store),
288 )
289 .with_buffer_size(Some(
290 context
291 .session_config()
292 .options()
293 .execution
294 .objectstore_writer_buffer_size,
295 ))
296 .build()?;
297 file_write_tasks.spawn(async move {
298 let mut row_count = 0;
299 while let Some(batch) = rx.recv().await {
300 row_count += batch.num_rows();
301 arrow_writer.write(&batch)?;
302 let mut buff_to_flush = shared_buffer.buffer.try_lock().unwrap();
303 if buff_to_flush.len() > BUFFER_FLUSH_BYTES {
304 object_store_writer
305 .write_all(buff_to_flush.as_slice())
306 .await?;
307 buff_to_flush.clear();
308 }
309 }
310 arrow_writer.finish()?;
311 let final_buff = shared_buffer.buffer.try_lock().unwrap();
312
313 object_store_writer.write_all(final_buff.as_slice()).await?;
314 object_store_writer.shutdown().await?;
315 Ok(row_count)
316 });
317 }
318
319 let mut row_count = 0;
320 while let Some(result) = file_write_tasks.join_next().await {
321 match result {
322 Ok(r) => {
323 row_count += r?;
324 }
325 Err(e) => {
326 if e.is_panic() {
327 std::panic::resume_unwind(e.into_panic());
328 } else {
329 unreachable!();
330 }
331 }
332 }
333 }
334
335 demux_task
336 .join_unwind()
337 .await
338 .map_err(|e| DataFusionError::ExecutionJoin(Box::new(e)))??;
339 Ok(row_count as u64)
340 }
341}
342
343impl Debug for ArrowFileSink {
344 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
345 f.debug_struct("ArrowFileSink").finish()
346 }
347}
348
349impl DisplayAs for ArrowFileSink {
350 fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
351 match t {
352 DisplayFormatType::Default | DisplayFormatType::Verbose => {
353 write!(f, "ArrowFileSink(file_groups=",)?;
354 FileGroupDisplay(&self.config.file_group).fmt_as(t, f)?;
355 write!(f, ")")
356 }
357 DisplayFormatType::TreeRender => {
358 writeln!(f, "format: arrow")?;
359 write!(f, "file={}", self.config.original_url)
360 }
361 }
362 }
363}
364
365#[async_trait]
366impl DataSink for ArrowFileSink {
367 fn schema(&self) -> &SchemaRef {
368 self.config.output_schema()
369 }
370
371 async fn write_all(
372 &self,
373 data: SendableRecordBatchStream,
374 context: &Arc<TaskContext>,
375 ) -> Result<u64> {
376 FileSink::write_all(self, data, context).await
377 }
378}
379
380const ARROW_MAGIC: [u8; 6] = *b"ARROW1";
384const CONTINUATION_MARKER: [u8; 4] = [0xff; 4];
385
386async fn infer_stream_schema(
387 mut stream: BoxStream<'static, object_store::Result<Bytes>>,
388) -> Result<SchemaRef> {
389 let bytes = extend_bytes_to_n_length_from_stream(vec![], 16, &mut stream).await?;
430
431 let preamble_len = if bytes[0..6] == ARROW_MAGIC {
433 if bytes[8..12] == CONTINUATION_MARKER {
435 12
437 } else {
438 8
440 }
441 } else if bytes[0..4] == CONTINUATION_MARKER {
442 4
444 } else {
445 0
447 };
448
449 let meta_len_bytes: [u8; 4] = bytes[preamble_len..preamble_len + 4]
450 .try_into()
451 .map_err(|err| {
452 ArrowError::ParseError(format!(
453 "Unable to read IPC message metadata length: {err:?}"
454 ))
455 })?;
456
457 let meta_len = i32::from_le_bytes([
458 meta_len_bytes[0],
459 meta_len_bytes[1],
460 meta_len_bytes[2],
461 meta_len_bytes[3],
462 ]);
463
464 if meta_len < 0 {
465 return Err(ArrowError::ParseError(
466 "IPC message metadata length is negative".to_string(),
467 )
468 .into());
469 }
470
471 let bytes = extend_bytes_to_n_length_from_stream(
472 bytes,
473 preamble_len + 4 + (meta_len as usize),
474 &mut stream,
475 )
476 .await?;
477
478 let message = root_as_message(&bytes[preamble_len + 4..]).map_err(|err| {
479 ArrowError::ParseError(format!("Unable to read IPC message metadata: {err:?}"))
480 })?;
481 let fb_schema = message.header_as_schema().ok_or_else(|| {
482 ArrowError::IpcError("Unable to read IPC message schema".to_string())
483 })?;
484 let schema = fb_to_schema(fb_schema);
485
486 Ok(Arc::new(schema))
487}
488
489async fn extend_bytes_to_n_length_from_stream(
490 bytes: Vec<u8>,
491 n: usize,
492 stream: &mut BoxStream<'static, object_store::Result<Bytes>>,
493) -> Result<Vec<u8>> {
494 if bytes.len() >= n {
495 return Ok(bytes);
496 }
497
498 let mut buf = bytes;
499
500 while let Some(b) = stream.next().await.transpose()? {
501 buf.extend_from_slice(&b);
502
503 if buf.len() >= n {
504 break;
505 }
506 }
507
508 if buf.len() < n {
509 return Err(ArrowError::ParseError(
510 "Unexpected end of byte stream for Arrow IPC file".to_string(),
511 )
512 .into());
513 }
514
515 Ok(buf)
516}
517
518async fn is_object_in_arrow_ipc_file_format(
519 store: Arc<dyn ObjectStore>,
520 object_location: &Path,
521) -> Result<bool> {
522 let get_opts = GetOptions {
523 range: Some(GetRange::Bounded(0..6)),
524 ..Default::default()
525 };
526 let bytes = store
527 .get_opts(object_location, get_opts)
528 .await?
529 .bytes()
530 .await?;
531 Ok(bytes.len() >= 6 && bytes[0..6] == ARROW_MAGIC)
532}
533
534#[cfg(test)]
535mod tests {
536 use super::*;
537
538 use std::any::Any;
539
540 use chrono::DateTime;
541 use datafusion_common::DFSchema;
542 use datafusion_common::config::TableOptions;
543 use datafusion_execution::config::SessionConfig;
544 use datafusion_execution::runtime_env::RuntimeEnv;
545 use datafusion_expr::execution_props::ExecutionProps;
546 use datafusion_expr::registry::ExtensionTypeRegistryRef;
547 use datafusion_expr::{
548 AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF,
549 };
550 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
551 use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList};
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 fn catalog_list(&self) -> Arc<dyn CatalogProviderList> {
579 Arc::new(EmptyCatalogProviderList)
580 }
581
582 async fn create_physical_plan(
583 &self,
584 _logical_plan: &LogicalPlan,
585 ) -> Result<Arc<dyn ExecutionPlan>> {
586 unimplemented!()
587 }
588
589 fn create_physical_expr(
590 &self,
591 _expr: Expr,
592 _df_schema: &DFSchema,
593 ) -> Result<Arc<dyn PhysicalExpr>> {
594 unimplemented!()
595 }
596
597 fn scalar_functions(&self) -> &HashMap<String, Arc<ScalarUDF>> {
598 unimplemented!()
599 }
600
601 fn higher_order_functions(&self) -> &HashMap<String, Arc<HigherOrderUDF>> {
602 unimplemented!()
603 }
604
605 fn aggregate_functions(&self) -> &HashMap<String, Arc<AggregateUDF>> {
606 unimplemented!()
607 }
608
609 fn window_functions(&self) -> &HashMap<String, Arc<WindowUDF>> {
610 unimplemented!()
611 }
612
613 fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef {
614 unimplemented!()
615 }
616
617 fn runtime_env(&self) -> &Arc<RuntimeEnv> {
618 &self.runtime_env
619 }
620
621 fn execution_props(&self) -> &ExecutionProps {
622 unimplemented!()
623 }
624
625 fn as_any(&self) -> &dyn Any {
626 unimplemented!()
627 }
628
629 fn table_options(&self) -> &TableOptions {
630 unimplemented!()
631 }
632
633 fn table_options_mut(&mut self) -> &mut TableOptions {
634 unimplemented!()
635 }
636
637 fn task_ctx(&self) -> Arc<TaskContext> {
638 unimplemented!()
639 }
640 }
641
642 #[tokio::test]
643 async fn test_infer_schema_stream() -> Result<()> {
644 for file in ["example.arrow", "example_stream.arrow"] {
645 let mut bytes = std::fs::read(format!("tests/data/{file}"))?;
646 bytes.truncate(bytes.len() - 20); let location = Path::parse(file)?;
648 let in_memory_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
649 in_memory_store.put(&location, bytes.into()).await?;
650
651 let state = MockSession::new();
652 let object_meta = ObjectMeta {
653 location,
654 last_modified: DateTime::default(),
655 size: u64::MAX,
656 e_tag: None,
657 version: None,
658 };
659
660 let arrow_format = ArrowFormat {};
661 let expected = vec!["f0: Int64", "f1: Utf8", "f2: Boolean"];
662
663 for chunk_size in [7, 3000] {
666 let store =
667 Arc::new(ChunkedStore::new(in_memory_store.clone(), chunk_size));
668 let inferred_schema = arrow_format
669 .infer_schema(
670 &state,
671 &(store.clone() as Arc<dyn ObjectStore>),
672 std::slice::from_ref(&object_meta),
673 )
674 .await?;
675 let actual_fields = inferred_schema
676 .fields()
677 .iter()
678 .map(|f| format!("{}: {:?}", f.name(), f.data_type()))
679 .collect::<Vec<_>>();
680 assert_eq!(expected, actual_fields);
681 }
682 }
683 Ok(())
684 }
685
686 #[tokio::test]
687 async fn test_infer_schema_short_stream() -> Result<()> {
688 for file in ["example.arrow", "example_stream.arrow"] {
689 let mut bytes = std::fs::read(format!("tests/data/{file}"))?;
690 bytes.truncate(20); let location = Path::parse(file)?;
692 let in_memory_store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
693 in_memory_store.put(&location, bytes.into()).await?;
694
695 let state = MockSession::new();
696 let object_meta = ObjectMeta {
697 location,
698 last_modified: DateTime::default(),
699 size: u64::MAX,
700 e_tag: None,
701 version: None,
702 };
703
704 let arrow_format = ArrowFormat {};
705
706 let store = Arc::new(ChunkedStore::new(in_memory_store.clone(), 7));
707 let err = arrow_format
708 .infer_schema(
709 &state,
710 &(store.clone() as Arc<dyn ObjectStore>),
711 std::slice::from_ref(&object_meta),
712 )
713 .await;
714
715 assert!(err.is_err());
716 assert_eq!(
717 "Arrow error: Parser error: Unexpected end of byte stream for Arrow IPC file",
718 err.unwrap_err().to_string().lines().next().unwrap()
719 );
720 }
721
722 Ok(())
723 }
724
725 #[tokio::test]
726 async fn test_format_detection_file_format() -> Result<()> {
727 let store = Arc::new(InMemory::new());
728 let path = Path::from("test.arrow");
729
730 let file_bytes = std::fs::read("tests/data/example.arrow")?;
731 store.put(&path, file_bytes.into()).await?;
732
733 let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?;
734 assert!(is_file, "Should detect file format");
735 Ok(())
736 }
737
738 #[tokio::test]
739 async fn test_format_detection_stream_format() -> Result<()> {
740 let store = Arc::new(InMemory::new());
741 let path = Path::from("test_stream.arrow");
742
743 let stream_bytes = std::fs::read("tests/data/example_stream.arrow")?;
744 store.put(&path, stream_bytes.into()).await?;
745
746 let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?;
747
748 assert!(!is_file, "Should detect stream format (not file)");
749
750 Ok(())
751 }
752
753 #[tokio::test]
754 async fn test_format_detection_corrupted_file() -> Result<()> {
755 let store = Arc::new(InMemory::new());
756 let path = Path::from("corrupted.arrow");
757
758 store
759 .put(&path, Bytes::from(vec![0x43, 0x4f, 0x52, 0x41]).into())
760 .await?;
761
762 let is_file = is_object_in_arrow_ipc_file_format(store.clone(), &path).await?;
763
764 assert!(
765 !is_file,
766 "Corrupted file should not be detected as file format"
767 );
768
769 Ok(())
770 }
771
772 #[tokio::test]
773 async fn test_format_detection_empty_file() -> Result<()> {
774 let store = Arc::new(InMemory::new());
775 let path = Path::from("empty.arrow");
776
777 store.put(&path, Bytes::new().into()).await?;
778
779 let result = is_object_in_arrow_ipc_file_format(store.clone(), &path).await;
780
781 assert!(result.is_err(), "Empty file should error");
783
784 Ok(())
785 }
786}