polars_python/functions/
io.rs

1use std::io::BufReader;
2
3#[cfg(any(feature = "ipc", feature = "parquet"))]
4use polars::prelude::ArrowSchema;
5use polars::prelude::CloudScheme;
6use pyo3::prelude::*;
7use pyo3::types::PyDict;
8
9use crate::conversion::Wrap;
10use crate::error::PyPolarsErr;
11use crate::file::{EitherRustPythonFile, get_either_file};
12use crate::io::cloud_options::OptPyCloudOptions;
13
14#[cfg(feature = "ipc")]
15#[pyfunction]
16pub fn read_ipc_schema(py: Python<'_>, py_f: Py<PyAny>) -> PyResult<Bound<'_, PyDict>> {
17    use arrow::io::ipc::read::read_file_metadata;
18    let metadata = match get_either_file(py_f, false)? {
19        EitherRustPythonFile::Rust(r) => {
20            read_file_metadata(&mut BufReader::new(r)).map_err(PyPolarsErr::from)?
21        },
22        EitherRustPythonFile::Py(mut r) => read_file_metadata(&mut r).map_err(PyPolarsErr::from)?,
23    };
24
25    let dict = PyDict::new(py);
26    fields_to_pydict(&metadata.schema, &dict)?;
27    Ok(dict)
28}
29
30#[cfg(feature = "parquet")]
31#[pyfunction]
32pub fn read_parquet_metadata(
33    py: Python,
34    py_f: Py<PyAny>,
35    storage_options: OptPyCloudOptions,
36    credential_provider: Option<Py<PyAny>>,
37) -> PyResult<Py<PyDict>> {
38    use std::io::Cursor;
39
40    use polars_error::feature_gated;
41    use polars_parquet::read::read_metadata;
42    use polars_parquet::read::schema::read_custom_key_value_metadata;
43
44    use crate::file::{PythonScanSourceInput, get_python_scan_source_input};
45
46    let metadata = match get_python_scan_source_input(py_f, false)? {
47        PythonScanSourceInput::Buffer(buf) => {
48            read_metadata(&mut Cursor::new(buf)).map_err(PyPolarsErr::from)?
49        },
50        PythonScanSourceInput::Path(p) => {
51            let cloud_options = storage_options.extract_opt_cloud_options(
52                CloudScheme::from_path(p.as_str()),
53                credential_provider,
54            )?;
55
56            if p.has_scheme() {
57                feature_gated!("cloud", {
58                    use polars::prelude::ParquetObjectStore;
59                    use polars_error::PolarsResult;
60
61                    py.detach(|| {
62                        polars_core::runtime::ASYNC.block_on(async {
63                            let mut reader =
64                                ParquetObjectStore::from_uri(p, cloud_options.as_ref(), None)
65                                    .await?;
66                            let result = reader.get_metadata().await?;
67                            PolarsResult::Ok((**result).clone())
68                        })
69                    })
70                })
71                .map_err(PyPolarsErr::from)?
72            } else {
73                let file =
74                    polars_utils::io::open_file(p.as_std_path()).map_err(PyPolarsErr::from)?;
75                read_metadata(&mut BufReader::new(file)).map_err(PyPolarsErr::from)?
76            }
77        },
78        PythonScanSourceInput::File(f) => {
79            read_metadata(&mut BufReader::new(f)).map_err(PyPolarsErr::from)?
80        },
81    };
82
83    let key_value_metadata = read_custom_key_value_metadata(metadata.key_value_metadata());
84    let dict = PyDict::new(py);
85    for (key, value) in key_value_metadata.into_iter() {
86        dict.set_item(key.as_str(), value.as_str())?;
87    }
88    Ok(dict.unbind())
89}
90
91/// Decode a parquet footer, optionally apply `FileMetadata::pruned`, then
92/// bincode-encode and return the byte length of the wire form.
93///
94/// Exposed for out-of-tree measurement of the IR-plan-borne metadata wire
95/// form (the `bincode(FileMetadata)` blob shipped to workers in distributed
96/// execution); no caller in py-polars itself.
97///
98/// `projection = None` ⇒ encode the full `FileMetadata`. `projection =
99/// Some(cols)` ⇒ apply `pruned(cols, predicate)`. Local files only.
100#[cfg(all(feature = "parquet", feature = "json"))]
101#[pyfunction]
102pub fn _bench_parquet_metadata_bincode_size(
103    path: &str,
104    projection: Option<Vec<String>>,
105    predicate: Vec<String>,
106) -> PyResult<usize> {
107    use polars_parquet::read::read_metadata;
108    use polars_utils::pl_serialize;
109    use polars_utils::pl_str::PlSmallStr;
110
111    let file = std::fs::File::open(path).map_err(|e| PyPolarsErr::Other(e.to_string()))?;
112    let metadata = read_metadata(&mut BufReader::new(file)).map_err(PyPolarsErr::from)?;
113
114    // Match the IR-plan serializer's framing format.
115    let bytes = match projection {
116        None => {
117            pl_serialize::serialize_to_bytes::<_, false>(&metadata).map_err(PyPolarsErr::from)?
118        },
119        Some(keep) => {
120            let keep_pl: Vec<PlSmallStr> = keep.into_iter().map(PlSmallStr::from).collect();
121            let pred_pl: Vec<PlSmallStr> = predicate.into_iter().map(PlSmallStr::from).collect();
122            let pruned = metadata
123                .pruned(&keep_pl, &pred_pl)
124                .map_err(|e| PyPolarsErr::Other(e.to_string()))?;
125            pl_serialize::serialize_to_bytes::<_, false>(&pruned).map_err(PyPolarsErr::from)?
126        },
127    };
128    Ok(bytes.len())
129}
130
131/// Decode a parquet footer, apply `FileMetadata::pruned(projection, predicate)`,
132/// and return the result as a JSON string. Format-agnostic custom serde lets
133/// the same wire DTOs emit JSON for inspection or bincode for dispatch.
134///
135/// Used by py-polars tests to assert structural prune behavior (only kept
136/// columns survive, stats only on predicate columns). Local files only.
137#[cfg(all(feature = "parquet", feature = "json"))]
138#[pyfunction]
139pub fn _parquet_metadata_pruned_json(
140    path: &str,
141    projection: Vec<String>,
142    predicate: Vec<String>,
143) -> PyResult<String> {
144    use polars_parquet::read::read_metadata;
145    use polars_utils::pl_str::PlSmallStr;
146
147    let file = std::fs::File::open(path).map_err(|e| PyPolarsErr::Other(e.to_string()))?;
148    let metadata = read_metadata(&mut BufReader::new(file)).map_err(PyPolarsErr::from)?;
149
150    let keep: Vec<PlSmallStr> = projection.into_iter().map(PlSmallStr::from).collect();
151    let pred: Vec<PlSmallStr> = predicate.into_iter().map(PlSmallStr::from).collect();
152    let pruned = metadata
153        .pruned(&keep, &pred)
154        .map_err(|e| PyPolarsErr::Other(e.to_string()))?;
155
156    serde_json::to_string(&pruned).map_err(|e| PyPolarsErr::Other(e.to_string()).into())
157}
158
159#[cfg(any(feature = "ipc", feature = "parquet"))]
160fn fields_to_pydict(schema: &ArrowSchema, dict: &Bound<'_, PyDict>) -> PyResult<()> {
161    for field in schema.iter_values() {
162        let dt = Wrap(polars::prelude::DataType::from_arrow_field(field));
163        dict.set_item(field.name.as_str(), &dt)?;
164    }
165    Ok(())
166}
167
168#[cfg(feature = "clipboard")]
169#[pyfunction]
170pub fn read_clipboard_string() -> PyResult<String> {
171    use arboard;
172    let mut clipboard =
173        arboard::Clipboard::new().map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
174    let result = clipboard
175        .get_text()
176        .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
177    Ok(result)
178}
179
180#[cfg(feature = "clipboard")]
181#[pyfunction]
182pub fn write_clipboard_string(s: &str) -> PyResult<()> {
183    use arboard;
184    let mut clipboard =
185        arboard::Clipboard::new().map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
186    clipboard
187        .set_text(s)
188        .map_err(|e| PyPolarsErr::Other(format!("{e}")))?;
189    Ok(())
190}
polars_python/functions/io.rs

polars_python/functions/
io.rs
