preliminary plumbing for unambiguous join

Author: ntjohnson1

Cargo.lock

@@ -39,10 +39,10 @@ pyo3 = { version = "0.24", features = ["extension-module", "abi3", "abi3-py39"]
 pyo3-async-runtimes = { version = "0.24", features = ["tokio-runtime"]}
 pyo3-log = "0.12.4"
 arrow = { version = "55.1.0", features = ["pyarrow"] }
-datafusion = { version = "48.0.0", features = ["avro", "unicode_expressions"] }
-datafusion-substrait = { version = "48.0.0", optional = true }
-datafusion-proto = { version = "48.0.0" }
-datafusion-ffi = { version = "48.0.0" }
+datafusion = { path="../datafusion/datafusion/core", features = ["avro", "unicode_expressions"] }
+datafusion-substrait = { path="../datafusion/datafusion/substrait", optional = true }
+datafusion-proto = { path="../datafusion/datafusion/proto" }
+datafusion-ffi = { path="../datafusion/datafusion/ffi" }
 prost = "0.13.1" # keep in line with `datafusion-substrait`
 uuid = { version = "1.16", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] }

Cargo.toml

@@ -34,6 +34,7 @@
 SqlType = common_internal.SqlType
 SqlView = common_internal.SqlView
 TableType = common_internal.TableType
+TableReference = common_internal.TableReference
 TableSource = common_internal.TableSource
 Constraints = common_internal.Constraints
 
@@ -51,6 +52,7 @@
     "SqlTable",
     "SqlType",
     "SqlView",
+    "TableReference",
     "TableSource",
     "TableType",
 ]

python/datafusion/common.py

@@ -40,6 +40,7 @@
 from datafusion._internal import DataFrame as DataFrameInternal
 from datafusion._internal import ParquetColumnOptions as ParquetColumnOptionsInternal
 from datafusion._internal import ParquetWriterOptions as ParquetWriterOptionsInternal
+from datafusion.common import TableReference
 from datafusion.expr import Expr, SortExpr, sort_or_default
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
@@ -1156,3 +1157,25 @@ def fill_null(self, value: Any, subset: list[str] | None = None) -> DataFrame:
             - For columns not in subset, the original column is kept unchanged
         """
         return DataFrame(self.df.fill_null(value, subset))
+
+    def fully_qualified_name(self, col: str) -> tuple[TableReference, pa.Field]:
+        """Get the fully qualified name of a column in the DataFrame.
+
+        Args:
+            col: The column name to get the fully qualified name for.
+
+        Returns:
+            Fully qualified column name as a string.
+        """
+        return self.df.fully_qualified_name(col)
+
+    def _drop(self, *qualified_columns: tuple[TableReference, pa.Field]) -> DataFrame:
+        """Drop columns from the DataFrame using fully qualified names.
+
+        Args:
+            qualified_columns: Fully qualified column names to drop.
+
+        Returns:
+            DataFrame with specified columns removed.
+        """
+        return DataFrame(self.df._drop(*qualified_columns))

python/datafusion/dataframe.py

@@ -272,6 +272,21 @@ def test_drop(df):
     assert result.column(1) == pa.array([4, 5, 6])
 
 
+def test__drop(df):
+    # TODO: implement deep copy for DataFrame?
+    ctx = SessionContext()
+
+    # create a RecordBatch and a new DataFrame from it
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([8, 5, 8])],
+        names=["a", "d", "e"],
+    )
+
+    other_df = ctx.from_arrow(batch)
+    df = df.join(other_df, on="a")._drop(df.fully_qualified_name("a"))
+    assert len(df.schema().names) == len(other_df.schema().names) * 2 - 1
+
+
 def test_limit(df):
     df = df.limit(1)
 

python/tests/test_dataframe.py

@@ -21,6 +21,7 @@ pub mod data_type;
 pub mod df_schema;
 pub mod function;
 pub mod schema;
+pub mod table_reference;
 
 /// Initializes the `common` module to match the pattern of `datafusion-common` https://docs.rs/datafusion-common/18.0.0/datafusion_common/index.html
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
@@ -39,5 +40,6 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_class::<schema::PyTableType>()?;
     m.add_class::<schema::PyTableSource>()?;
     m.add_class::<schema::PyConstraints>()?;
+    m.add_class::<table_reference::PyTableReference>()?;
     Ok(())
 }

src/common.rs

@@ -0,0 +1,111 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::common::TableReference;
+use pyo3::prelude::*;
+
+/// PyO3 requires that objects passed between Rust and Python implement the trait `PyClass`
+/// Since `TableReference` exists in another package we cannot make that happen here so we wrap
+/// `TableReference` as `PyTableReference` This exists solely to satisfy those constraints.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[pyclass(name = "TableReference", module = "datafusion.common", subclass)]
+pub struct PyTableReference {
+    pub table_reference: TableReference,
+}
+
+impl PyTableReference {
+    pub fn new(table_reference: TableReference) -> Self {
+        Self { table_reference }
+    }
+}
+
+impl From<PyTableReference> for TableReference {
+    fn from(py_table_reference: PyTableReference) -> TableReference {
+        py_table_reference.table_reference
+    }
+}
+
+impl From<TableReference> for PyTableReference {
+    fn from(table_reference: TableReference) -> PyTableReference {
+        PyTableReference { table_reference }
+    }
+}
+
+impl std::fmt::Display for PyTableReference {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.table_reference.fmt(f)
+    }
+}
+
+#[pymethods]
+impl PyTableReference {
+    /// Create a bare (unqualified) table reference
+    #[staticmethod]
+    pub fn bare(table: &str) -> Self {
+        Self::new(TableReference::bare(table))
+    }
+
+    /// Create a partial (schema.table) table reference
+    #[staticmethod]
+    pub fn partial(schema: &str, table: &str) -> Self {
+        Self::new(TableReference::partial(schema, table))
+    }
+
+    /// Create a full (catalog.schema.table) table reference
+    #[staticmethod]
+    pub fn full(catalog: &str, schema: &str, table: &str) -> Self {
+        Self::new(TableReference::full(catalog, schema, table))
+    }
+
+    /// Get the table name
+    pub fn table(&self) -> &str {
+        self.table_reference.table()
+    }
+
+    /// Get the schema name if present
+    pub fn schema(&self) -> Option<&str> {
+        self.table_reference.schema()
+    }
+
+    /// Get the catalog name if present
+    pub fn catalog(&self) -> Option<&str> {
+        self.table_reference.catalog()
+    }
+
+    /// Check if this is a bare table reference
+    pub fn is_bare(&self) -> bool {
+        matches!(self.table_reference, TableReference::Bare { .. })
+    }
+
+    /// Check if this is a partial table reference
+    pub fn is_partial(&self) -> bool {
+        matches!(self.table_reference, TableReference::Partial { .. })
+    }
+
+    /// Check if this is a full table reference
+    pub fn is_full(&self) -> bool {
+        matches!(self.table_reference, TableReference::Full { .. })
+    }
+
+    fn __str__(&self) -> String {
+        self.to_string()
+    }
+
+    fn __repr__(&self) -> String {
+        format!("TableReference('{self}')")
+    }
+}

src/common/table_reference.rs

@@ -21,14 +21,14 @@ use std::sync::Arc;
 
 use arrow::array::{new_null_array, RecordBatch, RecordBatchIterator, RecordBatchReader};
 use arrow::compute::can_cast_types;
+use arrow::datatypes::Field;
 use arrow::error::ArrowError;
 use arrow::ffi::FFI_ArrowSchema;
 use arrow::ffi_stream::FFI_ArrowArrayStream;
-use arrow::pyarrow::FromPyArrow;
 use datafusion::arrow::datatypes::Schema;
-use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
+use datafusion::arrow::pyarrow::{FromPyArrow, PyArrowType, ToPyArrow};
 use datafusion::arrow::util::pretty;
-use datafusion::common::UnnestOptions;
+use datafusion::common::{TableReference, UnnestOptions};
 use datafusion::config::{CsvOptions, ParquetColumnOptions, ParquetOptions, TableParquetOptions};
 use datafusion::dataframe::{DataFrame, DataFrameWriteOptions};
 use datafusion::datasource::TableProvider;
@@ -45,6 +45,7 @@ use pyo3::types::{PyCapsule, PyList, PyTuple, PyTupleMethods};
 use tokio::task::JoinHandle;
 
 use crate::catalog::PyTable;
+use crate::common::table_reference::PyTableReference;
 use crate::errors::{py_datafusion_err, to_datafusion_err, PyDataFusionError};
 use crate::expr::sort_expr::to_sort_expressions;
 use crate::physical_plan::PyExecutionPlan;
@@ -58,6 +59,81 @@ use crate::{
     expr::{sort_expr::PySortExpr, PyExpr},
 };
 
+/// A wrapper around Arc<str> that implements PyO3 traits for easier Python interop
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+pub struct PyArcStr(Arc<str>);
+
+impl PyArcStr {
+    pub fn new(s: &str) -> Self {
+        Self(Arc::from(s))
+    }
+
+    pub fn as_str(&self) -> &str {
+        &self.0
+    }
+
+    pub fn into_arc(self) -> Arc<str> {
+        self.0
+    }
+}
+
+impl From<Arc<str>> for PyArcStr {
+    fn from(arc: Arc<str>) -> Self {
+        Self(arc)
+    }
+}
+
+impl From<PyArcStr> for Arc<str> {
+    fn from(py_arc: PyArcStr) -> Self {
+        py_arc.0
+    }
+}
+
+impl From<&str> for PyArcStr {
+    fn from(s: &str) -> Self {
+        Self::new(s)
+    }
+}
+
+impl From<String> for PyArcStr {
+    fn from(s: String) -> Self {
+        Self(Arc::from(s))
+    }
+}
+
+impl std::fmt::Display for PyArcStr {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl<'py> pyo3::IntoPyObject<'py> for PyArcStr {
+    type Target = pyo3::types::PyString;
+    type Output = pyo3::Bound<'py, Self::Target>;
+    type Error = std::convert::Infallible;
+
+    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
+        Ok(pyo3::types::PyString::new(py, &self.0))
+    }
+}
+
+impl<'py> pyo3::IntoPyObject<'py> for &PyArcStr {
+    type Target = pyo3::types::PyString;
+    type Output = pyo3::Bound<'py, Self::Target>;
+    type Error = std::convert::Infallible;
+
+    fn into_pyobject(self, py: pyo3::Python<'py>) -> Result<Self::Output, Self::Error> {
+        Ok(pyo3::types::PyString::new(py, &self.0))
+    }
+}
+
+impl<'py> pyo3::FromPyObject<'py> for PyArcStr {
+    fn extract_bound(ob: &pyo3::Bound<'py, pyo3::PyAny>) -> pyo3::PyResult<Self> {
+        let s: String = ob.extract()?;
+        Ok(Self::new(&s))
+    }
+}
+
 // https://github.com/apache/datafusion-python/pull/1016#discussion_r1983239116
 // - we have not decided on the table_provider approach yet
 // this is an interim implementation
@@ -428,6 +504,17 @@ impl PyDataFrame {
         PyArrowType(self.df.schema().into())
     }
 
+    fn fully_qualified_name(&self, col: &str) -> PyResult<(PyTableReference, PyArrowType<Field>)> {
+        let result = self.df.schema().qualified_field_with_unqualified_name(col);
+        match result {
+            Ok(parts) => Ok((parts.0.unwrap().clone().into(), parts.1.clone().into())),
+            Err(err) => Err(PyValueError::new_err(format!(
+                "Error: {:?}",
+                err.to_string()
+            ))),
+        }
+    }
+
     /// Convert this DataFrame into a Table that can be used in register_table
     /// By convention, into_... methods consume self and return the new object.
     /// Disabling the clippy lint, so we can use &self
@@ -467,6 +554,20 @@ impl PyDataFrame {
         Ok(Self::new(df))
     }
 
+    #[pyo3(signature = (*_args))]
+    fn _drop(
+        &self,
+        _args: Vec<(PyTableReference, PyArrowType<Field>)>,
+    ) -> PyDataFusionResult<Self> {
+        // TODO need to finish plumbing through
+        let cols = _args
+            .iter()
+            .map(|(table, s)| (Some(table.clone().into()), s.0.clone()))
+            .collect::<Vec<(Option<TableReference>, Field)>>();
+        let df = self.df.as_ref().clone().drop_qualified_columns(&cols)?;
+        Ok(Self::new(df))
+    }
+
     fn filter(&self, predicate: PyExpr) -> PyDataFusionResult<Self> {
         let df = self.df.as_ref().clone().filter(predicate.into())?;
         Ok(Self::new(df))