More moving around of utils to clean up

Author: timsaucer

Cargo.lock

@@ -22,13 +22,13 @@ use arrow::array::{Array, ArrayRef};
 use arrow::datatypes::{Field, FieldRef};
 use arrow::ffi::{FFI_ArrowArray, FFI_ArrowSchema};
 use arrow::pyarrow::ToPyArrow;
+use datafusion_python_util::validate_pycapsule;
 use pyo3::ffi::c_str;
 use pyo3::prelude::{PyAnyMethods, PyCapsuleMethods};
 use pyo3::types::PyCapsule;
 use pyo3::{Bound, PyAny, PyResult, Python, pyclass, pymethods};
 
 use crate::errors::PyDataFusionResult;
-use crate::utils::validate_pycapsule;
 
 /// A Python object which implements the Arrow PyCapsule for importing
 /// into other libraries.

crates/core/src/array.rs

@@ -30,19 +30,20 @@ use datafusion::datasource::TableProvider;
 use datafusion_ffi::catalog_provider::FFI_CatalogProvider;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_ffi::schema_provider::FFI_SchemaProvider;
+use datafusion_python_util::{
+    create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, validate_pycapsule,
+    wait_for_future,
+};
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::PyKeyError;
 use pyo3::ffi::c_str;
 use pyo3::prelude::*;
 use pyo3::types::PyCapsule;
 
+use crate::context::PySessionContext;
 use crate::dataset::Dataset;
 use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err};
 use crate::table::PyTable;
-use crate::utils::{
-    create_logical_extension_capsule, extract_logical_extension_codec, validate_pycapsule,
-    wait_for_future,
-};
 
 #[pyclass(
     from_py_object,
@@ -710,6 +711,17 @@ fn extract_schema_provider_from_pyobj(
     Ok(provider)
 }
 
+fn extract_logical_extension_codec(
+    py: Python,
+    obj: Option<Bound<PyAny>>,
+) -> PyResult<Arc<FFI_LogicalExtensionCodec>> {
+    let obj = match obj {
+        Some(obj) => obj,
+        None => PySessionContext::global_ctx()?.into_bound_py_any(py)?,
+    };
+    ffi_logical_codec_from_pycapsule(obj).map(Arc::new)
+}
+
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<PyCatalog>()?;
     m.add_class::<PySchema>()?;

crates/core/src/catalog.rs

@@ -52,6 +52,10 @@ use datafusion_ffi::catalog_provider_list::FFI_CatalogProviderList;
 use datafusion_ffi::execution::FFI_TaskContextProvider;
 use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec;
 use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec;
+use datafusion_python_util::{
+    create_logical_extension_capsule, ffi_logical_codec_from_pycapsule, get_global_ctx,
+    get_tokio_runtime, spawn_future, validate_pycapsule, wait_for_future,
+};
 use object_store::ObjectStore;
 use pyo3::IntoPyObjectExt;
 use pyo3::exceptions::{PyKeyError, PyValueError};
@@ -82,10 +86,6 @@ use crate::udaf::PyAggregateUDF;
 use crate::udf::PyScalarUDF;
 use crate::udtf::PyTableFunction;
 use crate::udwf::PyWindowUDF;
-use crate::utils::{
-    create_logical_extension_capsule, extract_logical_extension_codec, get_global_ctx,
-    get_tokio_runtime, spawn_future, validate_pycapsule, wait_for_future,
-};
 
 /// Configuration options for a SessionContext
 #[pyclass(
@@ -1187,8 +1187,7 @@ impl PySessionContext {
         &self,
         codec: Bound<'py, PyAny>,
     ) -> PyDataFusionResult<Self> {
-        let py = codec.py();
-        let logical_codec = extract_logical_extension_codec(py, Some(codec))?;
+        let logical_codec = Arc::new(ffi_logical_codec_from_pycapsule(codec)?);
 
         Ok({
             Self {

crates/core/src/context.rs

@@ -41,6 +41,7 @@ use datafusion::logical_expr::SortExpr;
 use datafusion::logical_expr::dml::InsertOp;
 use datafusion::parquet::basic::{BrotliLevel, Compression, GzipLevel, ZstdLevel};
 use datafusion::prelude::*;
+use datafusion_python_util::{is_ipython_env, spawn_future, validate_pycapsule, wait_for_future};
 use futures::{StreamExt, TryStreamExt};
 use parking_lot::Mutex;
 use pyo3::PyErr;
@@ -58,7 +59,6 @@ use crate::physical_plan::PyExecutionPlan;
 use crate::record_batch::{PyRecordBatchStream, poll_next_batch};
 use crate::sql::logical::PyLogicalPlan;
 use crate::table::{PyTable, TempViewTable};
-use crate::utils::{is_ipython_env, spawn_future, validate_pycapsule, wait_for_future};
 
 /// File-level static CStr for the Arrow array stream capsule name.
 static ARROW_ARRAY_STREAM_NAME: &CStr = cstr!("arrow_array_stream");

crates/core/src/dataframe.rs

@@ -62,7 +62,6 @@ mod udaf;
 mod udf;
 pub mod udtf;
 mod udwf;
-pub mod utils;
 
 #[cfg(feature = "mimalloc")]
 #[global_allocator]

crates/core/src/lib.rs

@@ -20,14 +20,14 @@ use std::sync::Arc;
 use datafusion::arrow::pyarrow::ToPyArrow;
 use datafusion::arrow::record_batch::RecordBatch;
 use datafusion::physical_plan::SendableRecordBatchStream;
+use datafusion_python_util::wait_for_future;
 use futures::StreamExt;
 use pyo3::exceptions::{PyStopAsyncIteration, PyStopIteration};
 use pyo3::prelude::*;
 use pyo3::{PyAny, PyResult, Python, pyclass, pymethods};
 use tokio::sync::Mutex;
 
 use crate::errors::PyDataFusionError;
-use crate::utils::wait_for_future;
 
 #[pyclass(name = "RecordBatch", module = "datafusion", subclass, frozen)]
 pub struct PyRecordBatch {

crates/core/src/record_batch.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use datafusion_python_util::wait_for_future;
 use datafusion_substrait::logical_plan::{consumer, producer};
 use datafusion_substrait::serializer;
 use datafusion_substrait::substrait::proto::Plan;
@@ -25,7 +26,6 @@ use pyo3::types::PyBytes;
 use crate::context::PySessionContext;
 use crate::errors::{PyDataFusionError, PyDataFusionResult, py_datafusion_err, to_datafusion_err};
 use crate::sql::logical::PyLogicalPlan;
-use crate::utils::wait_for_future;
 
 #[pyclass(
     from_py_object,

crates/core/src/substrait.rs

@@ -27,13 +27,13 @@ use datafusion::datasource::{TableProvider, TableType};
 use datafusion::logical_expr::{Expr, LogicalPlanBuilder, TableProviderFilterPushDown};
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion::prelude::DataFrame;
+use datafusion_python_util::table_provider_from_pycapsule;
 use pyo3::IntoPyObjectExt;
 use pyo3::prelude::*;
 
 use crate::context::PySessionContext;
 use crate::dataframe::PyDataFrame;
 use crate::dataset::Dataset;
-use crate::utils::table_provider_from_pycapsule;
 
 /// This struct is used as a common method for all TableProviders,
 /// whether they refer to an FFI provider, an internally known

crates/core/src/table.rs

@@ -27,14 +27,14 @@ use datafusion::logical_expr::{
     Accumulator, AccumulatorFactoryFunction, AggregateUDF, AggregateUDFImpl, create_udaf,
 };
 use datafusion_ffi::udaf::FFI_AggregateUDF;
+use datafusion_python_util::{parse_volatility, validate_pycapsule};
 use pyo3::ffi::c_str;
 use pyo3::prelude::*;
 use pyo3::types::{PyCapsule, PyTuple};
 
 use crate::common::data_type::PyScalarValue;
 use crate::errors::{PyDataFusionResult, py_datafusion_err, to_datafusion_err};
 use crate::expr::PyExpr;
-use crate::utils::{parse_volatility, validate_pycapsule};
 
 #[derive(Debug)]
 struct RustAccumulator {