Merge branch 'main' into fill-null

Author: kosiew

.github/workflows/test.yaml

@@ -71,7 +71,7 @@ jobs:
 
       - name: Run Clippy
         if: ${{ matrix.python-version == '3.10' && matrix.toolchain == 'stable' }}
-        run: cargo clippy --all-targets --all-features -- -D clippy::all -A clippy::redundant_closure
+        run: cargo clippy --all-targets --all-features -- -D clippy::all -D warnings -A clippy::redundant_closure
 
       - name: Install dependencies and build
         uses: astral-sh/setup-uv@v5

.pre-commit-config.yaml

@@ -40,7 +40,7 @@ repos:
           - id: rust-clippy
             name: Rust clippy
             description: Run cargo clippy on files included in the commit. clippy should be installed before-hand.
-            entry: cargo clippy --all-targets --all-features -- -Dclippy::all -Aclippy::redundant_closure
+            entry: cargo clippy --all-targets --all-features -- -Dclippy::all -D warnings -Aclippy::redundant_closure
             pass_filenames: false
             types: [file, rust]
             language: system

python/datafusion/functions.py

@@ -252,6 +252,7 @@
     "to_hex",
     "to_timestamp",
     "to_timestamp_micros",
+    "to_timestamp_nanos",
     "to_timestamp_millis",
     "to_timestamp_seconds",
     "to_unixtime",

python/tests/test_functions.py

@@ -871,6 +871,7 @@ def test_temporal_functions(df):
         f.to_timestamp_millis(literal("2023-09-07 05:06:14.523952")),
         f.to_timestamp_micros(literal("2023-09-07 05:06:14.523952")),
         f.extract(literal("day"), column("d")),
+        f.to_timestamp_nanos(literal("2023-09-07 05:06:14.523952")),
     )
     result = df.collect()
     assert len(result) == 1
@@ -909,6 +910,9 @@ def test_temporal_functions(df):
         [datetime(2023, 9, 7, 5, 6, 14, 523952)] * 3, type=pa.timestamp("us")
     )
     assert result.column(10) == pa.array([31, 26, 2], type=pa.int32())
+    assert result.column(11) == pa.array(
+        [datetime(2023, 9, 7, 5, 6, 14, 523952)] * 3, type=pa.timestamp("ns")
+    )
 
 
 def test_arrow_cast(df):

src/config.rs

@@ -47,14 +47,14 @@ impl PyConfig {
     }
 
     /// Get a configuration option
-    pub fn get(&mut self, key: &str, py: Python) -> PyResult<PyObject> {
+    pub fn get<'py>(&mut self, key: &str, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
         let options = self.config.to_owned();
         for entry in options.entries() {
             if entry.key == key {
-                return Ok(entry.value.into_py(py));
+                return Ok(entry.value.into_pyobject(py)?);
             }
         }
-        Ok(None::<String>.into_py(py))
+        Ok(None::<String>.into_pyobject(py)?)
     }
 
     /// Set a configuration option
@@ -66,10 +66,10 @@ impl PyConfig {
 
     /// Get all configuration options
     pub fn get_all(&mut self, py: Python) -> PyResult<PyObject> {
-        let dict = PyDict::new_bound(py);
+        let dict = PyDict::new(py);
         let options = self.config.to_owned();
         for entry in options.entries() {
-            dict.set_item(entry.key, entry.value.clone().into_py(py))?;
+            dict.set_item(entry.key, entry.value.clone().into_pyobject(py)?)?;
         }
         Ok(dict.into())
     }

src/context.rs

@@ -458,8 +458,8 @@ impl PySessionContext {
         let py = data.py();
 
         // Instantiate pyarrow Table object & convert to Arrow Table
-        let table_class = py.import_bound("pyarrow")?.getattr("Table")?;
-        let args = PyTuple::new_bound(py, &[data]);
+        let table_class = py.import("pyarrow")?.getattr("Table")?;
+        let args = PyTuple::new(py, &[data])?;
         let table = table_class.call_method1("from_pylist", args)?;
 
         // Convert Arrow Table to datafusion DataFrame
@@ -478,8 +478,8 @@ impl PySessionContext {
         let py = data.py();
 
         // Instantiate pyarrow Table object & convert to Arrow Table
-        let table_class = py.import_bound("pyarrow")?.getattr("Table")?;
-        let args = PyTuple::new_bound(py, &[data]);
+        let table_class = py.import("pyarrow")?.getattr("Table")?;
+        let args = PyTuple::new(py, &[data])?;
         let table = table_class.call_method1("from_pydict", args)?;
 
         // Convert Arrow Table to datafusion DataFrame
@@ -533,8 +533,8 @@ impl PySessionContext {
         let py = data.py();
 
         // Instantiate pyarrow Table object & convert to Arrow Table
-        let table_class = py.import_bound("pyarrow")?.getattr("Table")?;
-        let args = PyTuple::new_bound(py, &[data]);
+        let table_class = py.import("pyarrow")?.getattr("Table")?;
+        let args = PyTuple::new(py, &[data])?;
         let table = table_class.call_method1("from_pandas", args)?;
 
         // Convert Arrow Table to datafusion DataFrame

src/dataframe.rs

@@ -545,12 +545,12 @@ impl PyDataFrame {
     /// Convert to Arrow Table
     /// Collect the batches and pass 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_pyobject(py)?.to_object(py);
+        let batches = self.collect(py)?.into_pyobject(py)?;
+        let schema = self.schema().into_pyobject(py)?;
 
         // Instantiate pyarrow Table object and use its from_batches method
-        let table_class = py.import_bound("pyarrow")?.getattr("Table")?;
-        let args = PyTuple::new_bound(py, &[batches, schema]);
+        let table_class = py.import("pyarrow")?.getattr("Table")?;
+        let args = PyTuple::new(py, &[batches, schema])?;
         let table: PyObject = table_class.call_method1("from_batches", args)?.into();
         Ok(table)
     }
@@ -585,8 +585,7 @@ impl PyDataFrame {
 
         let ffi_stream = FFI_ArrowArrayStream::new(reader);
         let stream_capsule_name = CString::new("arrow_array_stream").unwrap();
-        PyCapsule::new_bound(py, ffi_stream, Some(stream_capsule_name))
-            .map_err(PyDataFusionError::from)
+        PyCapsule::new(py, ffi_stream, Some(stream_capsule_name)).map_err(PyDataFusionError::from)
     }
 
     fn execute_stream(&self, py: Python) -> PyDataFusionResult<PyRecordBatchStream> {
@@ -649,8 +648,8 @@ impl PyDataFrame {
     /// Collect the batches, pass to Arrow Table & then convert to polars DataFrame
     fn to_polars(&self, py: Python<'_>) -> PyResult<PyObject> {
         let table = self.to_arrow_table(py)?;
-        let dataframe = py.import_bound("polars")?.getattr("DataFrame")?;
-        let args = PyTuple::new_bound(py, &[table]);
+        let dataframe = py.import("polars")?.getattr("DataFrame")?;
+        let args = PyTuple::new(py, &[table])?;
         let result: PyObject = dataframe.call1(args)?.into();
         Ok(result)
     }
@@ -673,7 +672,7 @@ fn print_dataframe(py: Python, df: DataFrame) -> PyDataFusionResult<()> {
 
     // Import the Python 'builtins' module to access the print function
     // Note that println! does not print to the Python debug console and is not visible in notebooks for instance
-    let print = py.import_bound("builtins")?.getattr("print")?;
+    let print = py.import("builtins")?.getattr("print")?;
     print.call1((result,))?;
     Ok(())
 }

src/dataset.rs

@@ -48,7 +48,7 @@ impl Dataset {
     // Creates a Python PyArrow.Dataset
     pub fn new(dataset: &Bound<'_, PyAny>, py: Python) -> PyResult<Self> {
         // Ensure that we were passed an instance of pyarrow.dataset.Dataset
-        let ds = PyModule::import_bound(py, "pyarrow.dataset")?;
+        let ds = PyModule::import(py, "pyarrow.dataset")?;
         let ds_attr = ds.getattr("Dataset")?;
         let ds_type = ds_attr.downcast::<PyType>()?;
         if dataset.is_instance(ds_type)? {

src/dataset_exec.rs

@@ -104,7 +104,7 @@ impl DatasetExec {
             })
             .transpose()?;
 
-        let kwargs = PyDict::new_bound(py);
+        let kwargs = PyDict::new(py);
 
         kwargs.set_item("columns", columns.clone())?;
         kwargs.set_item(
@@ -121,7 +121,7 @@ impl DatasetExec {
                 .0,
         );
 
-        let builtins = Python::import_bound(py, "builtins")?;
+        let builtins = Python::import(py, "builtins")?;
         let pylist = builtins.getattr("list")?;
 
         // Get the fragments or partitions of the dataset
@@ -198,7 +198,7 @@ impl ExecutionPlan for DatasetExec {
             let dataset_schema = dataset
                 .getattr("schema")
                 .map_err(|err| InnerDataFusionError::External(Box::new(err)))?;
-            let kwargs = PyDict::new_bound(py);
+            let kwargs = PyDict::new(py);
             kwargs
                 .set_item("columns", self.columns.clone())
                 .map_err(|err| InnerDataFusionError::External(Box::new(err)))?;
@@ -223,7 +223,7 @@ impl ExecutionPlan for DatasetExec {
             let record_batches: Bound<'_, PyIterator> = scanner
                 .call_method0("to_batches")
                 .map_err(|err| InnerDataFusionError::External(Box::new(err)))?
-                .iter()
+                .try_iter()
                 .map_err(|err| InnerDataFusionError::External(Box::new(err)))?;
 
             let record_batches = PyArrowBatchesAdapter {

src/errors.rs

@@ -91,3 +91,7 @@ pub fn py_datafusion_err(e: impl Debug) -> PyErr {
 pub fn py_unsupported_variant_err(e: impl Debug) -> PyErr {
     PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{e:?}"))
 }
+
+pub fn to_datafusion_err(e: impl Debug) -> InnerDataFusionError {
+    InnerDataFusionError::Execution(format!("{e:?}"))
+}