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::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
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 = 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        // Preserve projection from the original file source
213        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
247/// Implements [`FileSink`] for Arrow IPC files
248struct 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
380// Custom implementation of inferring schema. Should eventually be moved upstream to arrow-rs.
381// See <https://github.com/apache/arrow-rs/issues/5021>
382
383const 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    // IPC streaming format.
390    // See https://arrow.apache.org/docs/format/Columnar.html#ipc-streaming-format
391    //
392    //   <SCHEMA>
393    //   <DICTIONARY 0>
394    //   ...
395    //   <DICTIONARY k - 1>
396    //   <RECORD BATCH 0>
397    //   ...
398    //   <DICTIONARY x DELTA>
399    //   ...
400    //   <DICTIONARY y DELTA>
401    //   ...
402    //   <RECORD BATCH n - 1>
403    //   <EOS [optional]: 0xFFFFFFFF 0x00000000>
404
405    // The streaming format is made up of a sequence of encapsulated messages.
406    // See https://arrow.apache.org/docs/format/Columnar.html#encapsulated-message-format
407    //
408    //   <continuation: 0xFFFFFFFF>  (added in v0.15.0)
409    //   <metadata_size: int32>
410    //   <metadata_flatbuffer: bytes>
411    //   <padding>
412    //   <message body>
413    //
414    // The first message is the schema.
415
416    // IPC file format is a wrapper around the streaming format with indexing information.
417    // See https://arrow.apache.org/docs/format/Columnar.html#ipc-file-format
418    //
419    //   <magic number "ARROW1">
420    //   <empty padding bytes [to 8 byte boundary]>
421    //   <STREAMING FORMAT with EOS>
422    //   <FOOTER>
423    //   <FOOTER SIZE: int32>
424    //   <magic number "ARROW1">
425
426    // For the purposes of this function, the arrow "preamble" is the magic number, padding,
427    // and the continuation marker. 16 bytes covers the preamble and metadata length
428    // no matter which version or format is used.
429    let bytes = extend_bytes_to_n_length_from_stream(vec![], 16, &mut stream).await?;
430
431    // The preamble length is everything before the metadata length
432    let preamble_len = if bytes[0..6] == ARROW_MAGIC {
433        // File format starts with magic number "ARROW1"
434        if bytes[8..12] == CONTINUATION_MARKER {
435            // Continuation marker was added in v0.15.0
436            12
437        } else {
438            // File format before v0.15.0
439            8
440        }
441    } else if bytes[0..4] == CONTINUATION_MARKER {
442        // Stream format after v0.15.0 starts with continuation marker
443        4
444    } else {
445        // Stream format before v0.15.0 does not have a preamble
446        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); // mangle end to show we don't need to read whole file
647            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            // Test chunk sizes where too small so we keep having to read more bytes
664            // And when large enough that first read contains all we need
665            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); // should cause error that file shorter than expected
691            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        // currently errors because it tries to read 0..6 from an empty file
782        assert!(result.is_err(), "Empty file should error");
783
784        Ok(())
785    }
786}