feat(python): handle pyarrow/fsspec filesystems

[wjones127]

Description

Will allow passing in PyArrow and fsspec filesystems.

Related Issue(s)

• closes Python: Finish filesystem bindings #570

Documentation

[kylebrooks-8451]

@wjones127 - We had an internal project that used PyArrow, fsspec, and DataFusion that had this same issue. I'm attaching the fsspec ObjectStore implementation here, it uses an outdated DataFusion ObjectStore crate but it does implement some things that are missing here like list. I think this might be useful to this PR, if not, ignore this.

/// Implements a Datafusion ObjectStore that delegates to a Python fsspec filesystem
/// This allows us to use Python fsspec filesystems until native Rust / Datafusion filesystem are availible
use pyo3::prelude::*;
use pyo3::types::{PyType, PyDict, IntoPyDict, PyBytes};
use pyo3::exceptions::PyValueError;

use async_trait::async_trait;
use futures::{stream, AsyncRead};

use std::sync::Arc;
use std::io::{Read, Cursor, Error, ErrorKind};

use datafusion_data_access::{FileMeta, SizedFile};
use datafusion_data_access::object_store::{FileMetaStream, ListEntryStream, ObjectReader, ObjectStore};
use datafusion_data_access::Result as DFResult;

use chrono::{DateTime, Utc, NaiveDateTime};

pub(crate) fn python_datetime_to_rust(dt: &PyAny) -> Result<DateTime<Utc>, YeetError> {
    let timestamp: f64 = dt.call_method0("timestamp")?.extract()?;
    let naive_dt = NaiveDateTime::from_timestamp(timestamp.floor() as i64, (timestamp.fract() * 1e9) as u32);
    Ok(DateTime::from_utc(naive_dt, Utc))
}

/// `ObjectStore` and `ObjectReader` implementation for fsspec Python based filesystems
#[derive(Debug, Clone)]
#[pyclass(name = "FSSpecFileSystem", module = "yeet", subclass)]
pub(crate) struct FSSpecFileSystem {
    fs: PyObject,
}

#[pymethods]
impl FSSpecFileSystem {
    // Creates a Python fsspec filesystem as a Datafusion ObjectStore
    #[new(fs)]
    fn new(fs: &PyAny, py: Python) -> PyResult<Self> {
        // Ensure that we were passed an instance of fsspec.spec.AbstractFileSystem
        let fsspec_module = PyModule::import(py, "fsspec.spec")?;
        let afs_type: &PyType = fsspec_module.getattr("AbstractFileSystem")?.downcast()?;
        match fs.is_instance(afs_type)? {
            true => Ok(FSSpecFileSystem { fs: fs.into() }),
            false => Err(PyValueError::new_err("fs argument must be a fsspec.spec.AbstractFileSystem object")),
        }
    }

    #[getter]
    fn fs(&self, py: Python) -> PyResult<PyObject> {
        Ok(self.fs.clone_ref(py))
    }
}

impl FSSpecFileSystem {
    pub fn clone_ref(&self, py: Python) -> Self {
        FSSpecFileSystem { fs: self.fs.clone_ref(py) }
    }
}

fn create_file_metadata(path: String, object: &PyAny) -> Result<FileMeta, PyErr> {
    let size: u64 = object.get_item("size")?.extract().unwrap_or(0);
    let last_modified = object.get_item("last_modified").ok().and_then(|py_dt| python_datetime_to_rust(py_dt).ok());
    Ok(FileMeta { sized_file: SizedFile { path, size }, last_modified })
}

#[async_trait]
impl ObjectStore for FSSpecFileSystem {
    /// Returns all the files and directories in path `prefix`
    async fn list_file(&self, prefix: &str) -> DFResult<FileMetaStream> {
        let objects: DFResult<Vec<DFResult<FileMeta>>> = {
            Python::with_gil(|py| {
                let fs: &PyAny = self.fs.as_ref(py);
                let objects: Result<&PyAny, PyErr> = fs.call_method(
                    "find",
                    (prefix,),
                    Some(vec![("maxdepth", None), ("withdirs", Some(false)), ("detail", Some(true))].into_py_dict(py)),
                );
                let objects: &PyAny = objects.map_err(|err| Error::new(ErrorKind::Other, Box::new(err)))?;
                let objects: &PyDict =
                    objects.downcast().map_err(|err| Error::new(ErrorKind::Other, Box::new(PyErr::from(err))))?;

                Ok(objects
                    .iter()
                    .map(|(path, object)| {
                        let path: String = path.extract()?;
                        create_file_metadata(path, object).map_err(|err| Error::new(ErrorKind::Other, Box::new(err)))
                    })
                    .collect())
            })
        };
        let result = stream::iter(objects?);
        Ok(Box::pin(result))
    }

    /// Returns all the files in `prefix` if the `prefix` is already a leaf dir,
    /// or all paths between the `prefix` and the first occurrence of the `delimiter` if it is provided.
    async fn list_dir(&self, _prefix: &str, _delimiter: Option<String>) -> DFResult<ListEntryStream> {
        todo!()
    }

    /// Get object reader for one file
    fn file_reader(&self, file: SizedFile) -> DFResult<Arc<dyn ObjectReader>> {
        let f = Python::with_gil(|py| -> PyResult<PyObject> {
            let fs: &PyAny = self.fs.as_ref(py);
            let f: &PyAny = fs.call_method("open", (file.path.clone(),), Some(vec![("mode", "rb")].into_py_dict(py)))?;
            Ok(f.into())
        })
        .map_err(|err| Error::new(ErrorKind::Other, Box::new(err)))?;
        Ok(Arc::new(FSSpecFileReader { f, file }))
    }
}

struct FSSpecFileReader {
    f: PyObject,
    file: SizedFile,
}

#[async_trait]
impl ObjectReader for FSSpecFileReader {
    /// Get reader for a part [start, start + length] in the file asynchronously
    async fn chunk_reader(&self, _start: u64, _length: usize) -> DFResult<Box<dyn AsyncRead>> {
        todo!("implement once async file readers are available (arrow-rs#78, arrow-rs#111)")
    }

    /// Get reader for a part [start, start + length] in the file
    fn sync_chunk_reader(&self, start: u64, length: usize) -> DFResult<Box<dyn Read + Send + Sync>> {
        Python::with_gil(|py| {
            let f: &PyAny = self.f.as_ref(py);
            f.call_method1("seek", (start,)).map_err(|err| Error::new(ErrorKind::Other, Box::new(err)))?;
            let data: &PyBytes = f
                .call_method1("read", (length,))?
                .downcast()
                .map_err(PyErr::from)
                .map_err(|err| Error::new(ErrorKind::Other, Box::new(err)))?;
            let data: Box<dyn Read + Send + Sync> = Box::new(Cursor::new(data.as_bytes().to_owned()));
            Ok(data)
        })
    }

    /// Get the size of the file
    fn length(&self) -> u64 {
        self.file.size
    }
}