Skip to main content

datafusion_datasource_arrow/
file_format.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! [`ArrowFormat`]: Apache Arrow [`FileFormat`] abstractions
19//!
20//! Works with files following the [Arrow IPC format](https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format)
21
22use 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
70/// Initial writing buffer size. Note this is just a size hint for efficiency. It
71/// will grow beyond the set value if needed.
72const INITIAL_BUFFER_BYTES: usize = 1048576;
73
74/// If the buffered Arrow data exceeds this size, it is flushed to object store
75const BUFFER_FLUSH_BYTES: usize = 1024000;
76
77/// Factory struct used to create [`ArrowFormat`]
78#[derive(Default, Debug)]
79pub struct ArrowFormatFactory;
80
81impl ArrowFormatFactory {
82    /// Creates an instance of [ArrowFormatFactory]
83    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        // Removes the dot, i.e. ".parquet" -> "parquet"
105        DEFAULT_ARROW_EXTENSION[1..].to_string()
106    }
107}
108
109/// Arrow [`FileFormat`] implementation.
110#[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                            // not in the file format, but FileReader read some bytes
152                            // while trying to parse the file and so we need to rewind
153                            // it to the beginning of the file
154                            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        // Preserve projection from the original file source
214        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
248/// Implements [`FileSink`] for Arrow IPC files
249struct 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
381// Custom implementation of inferring schema. Should eventually be moved upstream to arrow-rs.
382// See <https://github.com/apache/arrow-rs/issues/5021>
383
384const 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    // IPC streaming format.
391    // See https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
392    //
393    //   <SCHEMA>
394    //   <DICTIONARY 0>
395    //   ...
396    //   <DICTIONARY k - 1>
397    //   <RECORD BATCH 0>
398    //   ...
399    //   <DICTIONARY x DELTA>
400    //   ...
401    //   <DICTIONARY y DELTA>
402    //   ...
403    //   <RECORD BATCH n - 1>
404    //   <EOS [optional]: 0xFFFFFFFF 0x00000000>
405
406    // The streaming format is made up of a sequence of encapsulated messages.
407    // See https://arrow.apache.org/docs/format/Columnar.html#encapsulated-message-format
408    //
409    //   <continuation: 0xFFFFFFFF>  (added in v0.15.0)
410    //   <metadata_size: int32>
411    //   <metadata_flatbuffer: bytes>
412    //   <padding>
413    //   <message body>
414    //
415    // The first message is the schema.
416
417    // IPC file format is a wrapper around the streaming format with indexing information.
418    // See https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format
419    //
420    //   <magic number "ARROW1">
421    //   <empty padding bytes [to 8 byte boundary]>
422    //   <STREAMING FORMAT with EOS>
423    //   <FOOTER>
424    //   <FOOTER SIZE: int32>
425    //   <magic number "ARROW1">
426
427    // For the purposes of this function, the arrow "preamble" is the magic number, padding,
428    // and the continuation marker. 16 bytes covers the preamble and metadata length
429    // no matter which version or format is used.
430    let bytes = extend_bytes_to_n_length_from_stream(vec![], 16, &mut stream).await?;
431
432    // The preamble length is everything before the metadata length
433    let preamble_len = if bytes[0..6] == ARROW_MAGIC {
434        // File format starts with magic number "ARROW1"
435        if bytes[8..12] == CONTINUATION_MARKER {
436            // Continuation marker was added in v0.15.0
437            12
438        } else {
439            // File format before v0.15.0
440            8
441        }
442    } else if bytes[0..4] == CONTINUATION_MARKER {
443        // Stream format after v0.15.0 starts with continuation marker
444        4
445    } else {
446        // Stream format before v0.15.0 does not have a preamble
447        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); // mangle end to show we don't need to read whole file
643            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            // Test chunk sizes where too small so we keep having to read more bytes
660            // And when large enough that first read contains all we need
661            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); // should cause error that file shorter than expected
687            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        // currently errors because it tries to read 0..6 from an empty file
778        assert!(result.is_err(), "Empty file should error");
779
780        Ok(())
781    }
782}