[FEAT]: huggingface integration

daft/daft.pyi

@@ -449,6 +449,7 @@ class HTTPConfig:
     """
 
     user_agent: str | None
+    bearer_token: str | None
 
 class S3Config:
     """

src/common/io-config/src/http.rs

@@ -4,15 +4,19 @@ use std::fmt::Formatter;
 use serde::Deserialize;
 use serde::Serialize;
 
+use crate::ObfuscatedString;
+
 #[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct HTTPConfig {
     pub user_agent: String,
+    pub bearer_token: Option<ObfuscatedString>,
 }
 
 impl Default for HTTPConfig {
     fn default() -> Self {
         HTTPConfig {
             user_agent: "daft/0.0.1".to_string(), // NOTE: Ideally we grab the version of Daft, but that requires a dependency on daft-core
+            bearer_token: None,
         }
     }
 }
@@ -30,6 +34,17 @@ impl Display for HTTPConfig {
             "HTTPConfig
     user_agent: {}",
             self.user_agent,
-        )
+        )?;
+
+        if let Some(bearer_token) = &self.bearer_token {
+            write!(
+                f,
+                "
+    bearer_token: {}",
+                bearer_token
+            )
+        } else {
+            Ok(())
+        }
     }
 }

src/common/io-config/src/python.rs

@@ -136,7 +136,7 @@ pub struct IOConfig {
 ///     user_agent (str, optional): The value for the user-agent header, defaults to "daft/{__version__}" if not provided
 ///
 /// Example:
-///     >>> io_config = IOConfig(http=HTTPConfig(user_agent="my_application/0.0.1"))
+///     >>> io_config = IOConfig(http=HTTPConfig(user_agent="my_application/0.0.1", bearer_token="xxx"))
 ///     >>> daft.read_parquet("http://some-path", io_config=io_config)
 #[derive(Clone, Default)]
 #[pyclass]

src/daft-core/src/count_mode.rs

@@ -1,5 +1,5 @@
 #[cfg(feature = "python")]
-use pyo3::{exceptions::PyValueError, prelude::*, types::PyBytes, PyTypeInfo};
+use pyo3::{exceptions::PyValueError, prelude::*};
 use serde::{Deserialize, Serialize};
 use std::fmt::{Display, Formatter, Result};
 use std::str::FromStr;

src/daft-core/src/join.rs

@@ -6,10 +6,7 @@ use std::{
 use crate::impl_bincode_py_state_serialization;
 use common_error::{DaftError, DaftResult};
 #[cfg(feature = "python")]
-use pyo3::{
-    exceptions::PyValueError, pyclass, pymethods, types::PyBytes, PyObject, PyResult, PyTypeInfo,
-    Python, ToPyObject,
-};
+use pyo3::{exceptions::PyValueError, pyclass, pymethods, PyResult, ToPyObject};
 
 use serde::{Deserialize, Serialize};
 

src/daft-core/src/python/datatype.rs

@@ -6,8 +6,7 @@ use pyo3::{
     class::basic::CompareOp,
     exceptions::PyValueError,
     prelude::*,
-    types::{PyBytes, PyDict, PyString},
-    PyTypeInfo,
+    types::{PyDict, PyString},
 };
 use serde::{Deserialize, Serialize};
 

src/daft-core/src/python/field.rs

@@ -1,4 +1,4 @@
-use pyo3::{prelude::*, types::PyBytes, PyTypeInfo};
+use pyo3::prelude::*;
 use serde::{Deserialize, Serialize};
 
 use super::datatype::PyDataType;

src/daft-core/src/python/schema.rs

@@ -1,8 +1,6 @@
 use std::sync::Arc;
 
 use pyo3::prelude::*;
-use pyo3::types::PyBytes;
-use pyo3::PyTypeInfo;
 
 use serde::{Deserialize, Serialize};
 

src/daft-core/src/utils/mod.rs

@@ -21,25 +21,34 @@ macro_rules! impl_binary_trait_by_reference {
 macro_rules! impl_bincode_py_state_serialization {
     ($ty:ty) => {
         #[cfg(feature = "python")]
-        #[pymethods]
+        #[pyo3::pymethods]
         impl $ty {
-            pub fn __reduce__(&self, py: Python) -> PyResult<(PyObject, PyObject)> {
+            pub fn __reduce__(
+                &self,
+                py: pyo3::Python,
+            ) -> pyo3::PyResult<(pyo3::PyObject, pyo3::PyObject)> {
+                use pyo3::PyTypeInfo;
+
                 Ok((
                     Self::type_object(py)
                         .getattr("_from_serialized")?
                         .to_object(py),
-                    (
-                        PyBytes::new(py, &$crate::utils::bincode::serialize(&self).unwrap())
-                            .to_object(py),
+                    (pyo3::types::PyBytes::new(
+                        py,
+                        &$crate::utils::bincode::serialize(&self).unwrap(),
                     )
+                    .to_object(py),)
                         .to_object(py),
                 ))
             }
 
             #[staticmethod]
-            pub fn _from_serialized(py: Python, serialized: PyObject) -> PyResult<Self> {
+            pub fn _from_serialized(
+                py: pyo3::Python,
+                serialized: pyo3::PyObject,
+            ) -> pyo3::PyResult<Self> {
                 serialized
-                    .extract::<&PyBytes>(py)
+                    .extract::<&pyo3::types::PyBytes>(py)
                     .map(|s| $crate::utils::bincode::deserialize(s.as_bytes()).unwrap())
             }
         }

src/daft-csv/src/options.rs

@@ -5,10 +5,7 @@ use serde::{Deserialize, Serialize};
 use {
     daft_core::python::schema::PySchema,
     daft_dsl::python::PyExpr,
-    pyo3::{
-        pyclass, pyclass::CompareOp, pymethods, types::PyBytes, PyObject, PyResult, PyTypeInfo,
-        Python, ToPyObject,
-    },
+    pyo3::{pyclass, pyclass::CompareOp, pymethods, PyResult, ToPyObject},
 };
 
 /// Options for converting CSV data to Daft data.

src/daft-dsl/src/python.rs

@@ -23,7 +23,6 @@ use pyo3::{
     prelude::*,
     pyclass::CompareOp,
     types::{PyBool, PyBytes, PyFloat, PyInt, PyString},
-    PyTypeInfo,
 };
 
 #[pyfunction]

src/daft-io/Cargo.toml

@@ -15,6 +15,7 @@ azure_storage_blobs = {version = "0.17.0", features = ["enable_reqwest"], defaul
 bytes = {workspace = true}
 common-error = {path = "../common/error", default-features = false}
 common-io-config = {path = "../common/io-config", default-features = false}
+daft-core = {path = "../daft-core", default-features = false}
 futures = {workspace = true}
 globset = "0.4"
 google-cloud-storage = {version = "0.15.0", default-features = false, features = ["default-tls", "auth"]}
@@ -29,22 +30,23 @@ openssl-sys = {version = "0.9.102", features = ["vendored"]}
 pyo3 = {workspace = true, optional = true}
 rand = "0.8.5"
 regex = {version = "1.10.4"}
+serde = {workspace = true}
 snafu = {workspace = true}
 tokio = {workspace = true}
 tokio-stream = {workspace = true}
 url = {workspace = true}
 
 [dependencies.reqwest]
 default-features = false
-features = ["stream", "native-tls"]
+features = ["stream", "native-tls", "json"]
 version = "0.11.18"
 
 [dev-dependencies]
 md5 = "0.7.0"
 tempfile = "3.8.1"
 
 [features]
-python = ["dep:pyo3", "common-error/python", "common-io-config/python"]
+python = ["dep:pyo3", "common-error/python", "common-io-config/python", "daft-core/python"]
 
 [package]
 edition = {workspace = true}

src/daft-io/src/azure_blob.rs

@@ -15,7 +15,7 @@ use crate::{
     object_io::{FileMetadata, FileType, LSResult, ObjectSource},
     stats::IOStatsRef,
     stream_utils::io_stats_on_bytestream,
-    GetResult,
+    FileFormat, GetResult,
 };
 use common_io_config::AzureConfig;
 
@@ -577,6 +577,7 @@ impl ObjectSource for AzureBlobSource {
         page_size: Option<i32>,
         limit: Option<usize>,
         io_stats: Option<IOStatsRef>,
+        _file_format: Option<FileFormat>,
     ) -> super::Result<BoxStream<'static, super::Result<FileMetadata>>> {
         use crate::object_store_glob::glob;
 

src/daft-io/src/file_format.rs

@@ -0,0 +1,58 @@
+use std::str::FromStr;
+
+use common_error::{DaftError, DaftResult};
+use daft_core::impl_bincode_py_state_serialization;
+#[cfg(feature = "python")]
+use pyo3::prelude::*;
+
+use serde::{Deserialize, Serialize};
+
+/// Format of a file, e.g. Parquet, CSV, JSON.
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Copy)]
+#[cfg_attr(feature = "python", pyclass(module = "daft.daft"))]
+pub enum FileFormat {
+    Parquet,
+    Csv,
+    Json,
+    Database,
+    Python,
+}
+
+#[cfg(feature = "python")]
+#[pymethods]
+impl FileFormat {
+    fn ext(&self) -> &'static str {
+        match self {
+            Self::Parquet => "parquet",
+            Self::Csv => "csv",
+            Self::Json => "json",
+            Self::Database => "db",
+            Self::Python => "py",
+        }
+    }
+}
+
+impl FromStr for FileFormat {
+    type Err = DaftError;
+
+    fn from_str(file_format: &str) -> DaftResult<Self> {
+        use FileFormat::*;
+
+        if file_format.trim().eq_ignore_ascii_case("parquet") {
+            Ok(Parquet)
+        } else if file_format.trim().eq_ignore_ascii_case("csv") {
+            Ok(Csv)
+        } else if file_format.trim().eq_ignore_ascii_case("json") {
+            Ok(Json)
+        } else if file_format.trim().eq_ignore_ascii_case("database") {
+            Ok(Database)
+        } else {
+            Err(DaftError::TypeError(format!(
+                "FileFormat {} not supported!",
+                file_format
+            )))
+        }
+    }
+}
+
+impl_bincode_py_state_serialization!(FileFormat);

src/daft-io/src/google_cloud.rs

@@ -23,6 +23,7 @@ use crate::object_io::LSResult;
 use crate::object_io::ObjectSource;
 use crate::stats::IOStatsRef;
 use crate::stream_utils::io_stats_on_bytestream;
+use crate::FileFormat;
 use crate::GetResult;
 use common_io_config::GCSConfig;
 
@@ -436,6 +437,7 @@ impl ObjectSource for GCSSource {
         page_size: Option<i32>,
         limit: Option<usize>,
         io_stats: Option<IOStatsRef>,
+        _file_format: Option<FileFormat>,
     ) -> super::Result<BoxStream<'static, super::Result<FileMetadata>>> {
         use crate::object_store_glob::glob;
 

src/daft-io/src/http.rs

@@ -15,6 +15,7 @@ use crate::{
     object_io::{FileMetadata, FileType, LSResult},
     stats::IOStatsRef,
     stream_utils::io_stats_on_bytestream,
+    FileFormat,
 };
 
 use super::object_io::{GetResult, ObjectSource};
@@ -140,7 +141,7 @@ fn _get_file_metadata_from_html(path: &str, text: &str) -> super::Result<Vec<Fil
 }
 
 pub(crate) struct HttpSource {
-    client: reqwest::Client,
+    pub(crate) client: reqwest::Client,
 }
 
 impl From<Error> for super::Error {
@@ -276,6 +277,7 @@ impl ObjectSource for HttpSource {
         _page_size: Option<i32>,
         limit: Option<usize>,
         io_stats: Option<IOStatsRef>,
+        _file_format: Option<FileFormat>,
     ) -> super::Result<BoxStream<'static, super::Result<FileMetadata>>> {
         use crate::object_store_glob::glob;