Change tpch validation to use exec_sql_on_tables

Cargo.toml

@@ -86,6 +86,9 @@ tonic-build = { version = "0.8", default-features = false, features = [
 ] }
 url = "2"
 
+[dev-dependencies]
+tempfile = "3.17"
+
 [lib]
 name = "datafusion_ray"
 crate-type = ["cdylib", "rlib"]

datafusion_ray/__init__.py

@@ -20,6 +20,6 @@
 except ImportError:
     import importlib_metadata
 
-from .core import RayContext, prettify, runtime_env, RayStagePool
+from .core import RayContext, exec_sql_on_tables, prettify, runtime_env, RayStagePool
 
 __version__ = importlib_metadata.version(__name__)

datafusion_ray/core.py

@@ -31,6 +31,7 @@
 from datafusion_ray._datafusion_ray_internal import (
     RayContext as RayContextInternal,
     RayDataFrame as RayDataFrameInternal,
+    exec_sql_on_tables,
     prettify,
 )
 
@@ -465,6 +466,9 @@ def stages(self):
 
         return self._stages
 
+    def schema(self):
+        return self.df.schema()
+
     def execution_plan(self):
         return self.df.execution_plan()
 
@@ -479,7 +483,7 @@ def collect(self) -> list[pa.RecordBatch]:
             t1 = time.time()
             self.stages()
             t2 = time.time()
-            log.debug(f"creating stages took {t2 -t1}s")
+            log.debug(f"creating stages took {t2 - t1}s")
 
             last_stage_id = max([stage.stage_id for stage in self._stages])
             log.debug(f"last stage is {last_stage_id}")
@@ -553,7 +557,9 @@ def __init__(
         s = time.time()
         call_sync(wait_for([start_ref], "RayContextSupervisor start"))
         e = time.time()
-        log.info(f"RayContext::__init__ waiting for supervisor to be ready took {e-s}s")
+        log.info(
+            f"RayContext::__init__ waiting for supervisor to be ready took {e - s}s"
+        )
 
     def register_parquet(self, name: str, path: str):
         self.ctx.register_parquet(name, path)

src/dataframe.rs

@@ -266,6 +266,10 @@ impl RayDataFrame {
         Ok(PyLogicalPlan::new(self.df.logical_plan().clone()))
     }
 
+    fn schema(&self, py: Python) -> PyResult<PyObject> {
+        self.df.schema().as_arrow().to_pyarrow(py)
+    }
+
     fn optimized_logical_plan(&self) -> PyResult<PyLogicalPlan> {
         Ok(PyLogicalPlan::new(self.df.clone().into_optimized_plan()?))
     }

src/lib.rs

@@ -44,6 +44,7 @@ fn _datafusion_ray_internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<dataframe::PyDataFrameStage>()?;
     m.add_class::<stage_service::StageService>()?;
     m.add_function(wrap_pyfunction!(util::prettify, m)?)?;
+    m.add_function(wrap_pyfunction!(util::exec_sql_on_tables, m)?)?;
     Ok(())
 }