From d682ad1dfe729174067dc1184b4d41cb404aa66a Mon Sep 17 00:00:00 2001
From: Trent Hauck <trent@trenthauck.com>
Date: Fri, 13 Oct 2023 10:48:52 -0700
Subject: [PATCH] feat: improve setup

---
 Makefile                |  2 +-
 src/execution_result.rs |  8 +++-----
 src/session_context.rs  | 12 +++++++++++-
 3 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index 2727a7f..d63d391 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,5 @@
 build:
-	cargo build
+	cargo build --release
 	maturin develop --release
 
 test: build
diff --git a/src/execution_result.rs b/src/execution_result.rs
index 41db94f..271da4f 100644
--- a/src/execution_result.rs
+++ b/src/execution_result.rs
@@ -19,7 +19,7 @@ use arrow::{
     pyarrow::{PyArrowType, ToPyArrow},
 };
 use datafusion::prelude::DataFrame;
-use pyo3::{pyclass, pymethods, types::PyTuple, IntoPy, PyObject, PyResult, Python, ToPyObject};
+use pyo3::{pyclass, pymethods, types::PyTuple, PyObject, PyResult, Python, ToPyObject};
 
 use crate::{error, runtime::wait_for_future};
 
@@ -52,12 +52,11 @@ impl PyExecutionResult {
     /// Convert to Arrow Table
     fn to_arrow_table(&self, py: Python) -> PyResult<PyObject> {
         let batches = self.collect(py)?.to_object(py);
-        let schema: PyObject = self.schema().into_py(py);
 
         Python::with_gil(|py| {
             // Instantiate pyarrow Table object and use its from_batches method
             let table_class = py.import("pyarrow")?.getattr("Table")?;
-            let args = PyTuple::new(py, &[batches, schema]);
+            let args = PyTuple::new(py, &[batches]);
             let table: PyObject = table_class.call_method1("from_batches", args)?.into();
             Ok(table)
         })
@@ -66,11 +65,10 @@ impl PyExecutionResult {
     /// Convert to a Polars DataFrame
     fn to_polars(&self, py: Python) -> PyResult<PyObject> {
         let batches = self.collect(py)?.to_object(py);
-        let schema: PyObject = self.schema().into_py(py);
 
         Python::with_gil(|py| {
             let table_class = py.import("pyarrow")?.getattr("Table")?;
-            let args = PyTuple::new(py, &[batches, schema]);
+            let args = PyTuple::new(py, &[batches]);
             let table: PyObject = table_class.call_method1("from_batches", args)?.into();
 
             let table_class = py.import("polars")?.getattr("DataFrame")?;
diff --git a/src/session_context.rs b/src/session_context.rs
index c92c7a5..a16937f 100644
--- a/src/session_context.rs
+++ b/src/session_context.rs
@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use datafusion::prelude::SessionContext;
-use exon::ExonSessionExt;
+use exon::{ExonRuntimeEnvExt, ExonSessionExt};
 
 use pyo3::prelude::*;
 
@@ -41,12 +41,22 @@ impl ExonSessionContext {
         Ok(Self::default())
     }
 
+    /// Execute a SQL query and return the result as a [`PyExecutionResult`].
     fn sql(&mut self, query: &str, py: Python) -> PyResult<PyExecutionResult> {
         let result = self.ctx.sql(query);
         let df = wait_for_future(py, result).map_err(error::BioBearError::from)?;
 
         Ok(PyExecutionResult::new(df))
     }
+
+    /// Register an object store with the given URI.
+    fn register_object_store_from_url(&mut self, url: &str, py: Python) -> PyResult<()> {
+        let runtime = self.ctx.runtime_env();
+        let registration = runtime.exon_register_object_store_uri(url);
+        wait_for_future(py, registration).map_err(error::BioBearError::from)?;
+
+        Ok(())
+    }
 }
 
 #[pyfunction]