Skip to content

Commit

Permalink
[CHORE] Change IOConfig to be serialized into binary instead of JSON (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
kevinzwang authored Nov 23, 2024
1 parent 68248a9 commit 2ae53f4
Show file tree
Hide file tree
Showing 5 changed files with 8 additions and 41 deletions.
2 changes: 0 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

7 changes: 0 additions & 7 deletions daft/daft/__init__.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -638,13 +638,6 @@ class IOConfig:
gcs: GCSConfig | None = None,
http: HTTPConfig | None = None,
): ...
@staticmethod
def from_json(input: str) -> IOConfig:
"""
Recreate an IOConfig from a JSON string.
"""
...

def replace(
self,
s3: S3Config | None = None,
Expand Down
8 changes: 0 additions & 8 deletions daft/io/config.py

This file was deleted.

4 changes: 1 addition & 3 deletions src/common/io-config/Cargo.toml
Original file line number Diff line number Diff line change
@@ -1,17 +1,15 @@
[dependencies]
aws-credential-types = {version = "0.55.3"}
chrono = {workspace = true}
common-error = {path = "../error", default-features = false}
common-py-serde = {path = "../py-serde", default-features = false}
derive_more = {workspace = true}
pyo3 = {workspace = true, optional = true}
secrecy = {version = "0.8.0", features = ["alloc"], default-features = false}
serde = {workspace = true}
serde_json = {workspace = true}
typetag = {workspace = true}

[features]
python = ["dep:pyo3", "common-error/python", "common-py-serde/python"]
python = ["dep:pyo3", "common-py-serde/python"]

[lints]
workspace = true
Expand Down
28 changes: 7 additions & 21 deletions src/common/io-config/src/python.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,9 @@ use aws_credential_types::{
provider::{error::CredentialsError, ProvideCredentials},
Credentials,
};
use common_error::DaftError;
use common_py_serde::{deserialize_py_object, serialize_py_object};
use common_py_serde::{
deserialize_py_object, impl_bincode_py_state_serialization, serialize_py_object,
};
use pyo3::prelude::*;
use serde::{Deserialize, Serialize};

Expand Down Expand Up @@ -131,8 +132,8 @@ pub struct GCSConfig {
/// Example:
/// >>> io_config = IOConfig(s3=S3Config(key_id="xxx", access_key="xxx", num_tries=10), azure=AzureConfig(anonymous=True), gcs=GCSConfig(...))
/// >>> daft.read_parquet(["s3://some-path", "az://some-other-path", "gs://path3"], io_config=io_config)
#[derive(Clone, Default)]
#[pyclass]
#[derive(Clone, Default, Serialize, Deserialize)]
#[pyclass(module = "daft.daft")]
pub struct IOConfig {
pub config: config::IOConfig,
}
Expand Down Expand Up @@ -234,23 +235,6 @@ impl IOConfig {
})
}

#[staticmethod]
pub fn from_json(input: &str) -> PyResult<Self> {
let config: config::IOConfig = serde_json::from_str(input).map_err(DaftError::from)?;
Ok(config.into())
}

pub fn __reduce__(&self, py: Python) -> PyResult<(PyObject, (String,))> {
let io_config_module = py.import_bound(pyo3::intern!(py, "daft.io.config"))?;
let json_string = serde_json::to_string(&self.config).map_err(DaftError::from)?;
Ok((
io_config_module
.getattr(pyo3::intern!(py, "_io_config_from_json"))?
.into(),
(json_string,),
))
}

pub fn __hash__(&self) -> PyResult<u64> {
use std::{collections::hash_map::DefaultHasher, hash::Hash};

Expand All @@ -260,6 +244,8 @@ impl IOConfig {
}
}

impl_bincode_py_state_serialization!(IOConfig);

#[pymethods]
impl S3Config {
#[allow(clippy::too_many_arguments)]
Expand Down

0 comments on commit 2ae53f4

Please sign in to comment.