[BUG] Block on parquet schema future in estimate_size_bytes (#3484)

Co-authored-by: Colin Ho <[email protected]>
Eventual-Inc · Dec 4, 2024 · 83470e0 · 83470e0
1 parent 95d7a43
commit 83470e0
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/src/daft-scan/src/python.rs b/src/daft-scan/src/python.rs
@@ -524,13 +524,15 @@ pub mod pylib {
         columns: Option<Vec<String>>,
         has_metadata: Option<bool>,
     ) -> PyResult<usize> {
-        let (schema, metadata) = daft_parquet::read::read_parquet_schema(
-            uri,
-            default::Default::default(),
-            None,
-            default::Default::default(),
-            None,
-        )?;
+        let io_runtime = common_runtime::get_io_runtime(true);
+        let (schema, metadata) =
+            io_runtime.block_on_current_thread(daft_parquet::read::read_parquet_schema(
+                uri,
+                default::Default::default(),
+                None,
+                default::Default::default(),
+                None,
+            ))?;
         let data_source = DataSource::File {
             path: uri.to_string(),
             chunk_spec: None,