Datafusion invalid projection

dask_planner/src/sql/logical.rs

@@ -9,8 +9,9 @@ mod filter;
 mod join;
 mod limit;
 mod offset;
-pub mod projection;
+mod projection;
 mod sort;
+mod table_scan;
 mod union;
 
 use datafusion_common::{Column, DFSchemaRef, DataFusionError, Result};
@@ -111,6 +112,11 @@ impl PyLogicalPlan {
         to_py_plan(self.current_node.as_ref())
     }
 
+    /// LogicalPlan::TableScan as PyTableScan
+    pub fn table_scan(&self) -> PyResult<table_scan::PyTableScan> {
+        to_py_plan(self.current_node.as_ref())
+    }
+
     /// Gets the "input" for the current LogicalPlan
     pub fn get_inputs(&mut self) -> PyResult<Vec<PyLogicalPlan>> {
         let mut py_inputs: Vec<PyLogicalPlan> = Vec::new();

dask_planner/src/sql/logical/table_scan.rs

@@ -0,0 +1,45 @@
+use crate::sql::exceptions::py_type_err;
+use crate::sql::logical;
+use datafusion_expr::logical_plan::TableScan;
+use pyo3::prelude::*;
+
+#[pyclass(name = "TableScan", module = "dask_planner", subclass)]
+#[derive(Clone)]
+pub struct PyTableScan {
+    pub(crate) table_scan: TableScan,
+}
+
+#[pymethods]
+impl PyTableScan {
+    #[pyo3(name = "getTableScanProjects")]
+    fn scan_projects(&mut self) -> PyResult<Vec<String>> {
+        match &self.table_scan.projection {
+            Some(indices) => {
+                let schema = self.table_scan.source.schema();
+                Ok(indices
+                    .iter()
+                    .map(|i| schema.field(*i).name().to_string())
+                    .collect())
+            }
+            None => Ok(vec![]),
+        }
+    }
+
+    /// If the 'TableScan' contains columns that should be projected during the
+    /// read return True, otherwise return False
+    #[pyo3(name = "containsProjections")]
+    fn contains_projections(&self) -> bool {
+        self.table_scan.projection.is_some()
+    }
+}
+
+impl TryFrom<logical::LogicalPlan> for PyTableScan {
+    type Error = PyErr;
+
+    fn try_from(logical_plan: logical::LogicalPlan) -> Result<Self, Self::Error> {
+        match logical_plan {
+            logical::LogicalPlan::TableScan(table_scan) => Ok(PyTableScan { table_scan }),
+            _ => Err(py_type_err("unexpected plan")),
+        }
+    }
+}

dask_planner/src/sql/types/rel_data_type.rs

@@ -76,8 +76,7 @@ impl RelDataType {
         self.field_list.clone()
     }
 
-    /// Returns the names of the fields in a struct type. The field count
-    /// is equal to the size of the returned list.
+    /// Returns the names of all of the columns in a given DaskTable
     #[pyo3(name = "getFieldNames")]
     pub fn field_names(&self) -> Vec<String> {
         assert!(!self.field_list.is_empty());

dask_sql/physical/rel/logical/table_scan.py

@@ -33,6 +33,9 @@ def convert(
         # There should not be any input. This is the first step.
         self.assert_inputs(rel, 0)
 
+        # Rust table_scan instance handle
+        table_scan = rel.table_scan()
+
         # The table(s) we need to return
         table = rel.getTable()
 
@@ -48,11 +51,15 @@ def convert(
         df = dc.df
         cc = dc.column_container
 
-        # Make sure we only return the requested columns
-        row_type = table.getRowType()
-        field_specifications = [str(f) for f in row_type.getFieldNames()]
-        cc = cc.limit_to(field_specifications)
+        # If the 'TableScan' instance contains projected columns only retrieve those columns
+        # otherwise get all projected columns from the 'Projection' instance, which is contained
+        # in the 'RelDataType' instance, aka 'row_type'
+        if table_scan.containsProjections():
+            field_specifications = table_scan.getTableScanProjects()
+        else:
+            field_specifications = [str(f) for f in table.getRowType().getFieldNames()]
 
+        cc = cc.limit_to(field_specifications)
         cc = self.fix_column_to_row_type(cc, rel.getRowType())
         dc = DataContainer(df, cc)
         dc = self.fix_dtype_to_row_type(dc, rel.getRowType())

tests/integration/test_select.py

@@ -222,3 +222,32 @@ def test_case_when_no_else(c):
     expected_df = pd.DataFrame({"C": [None, 1, 1, 1, None]})
 
     assert_eq(actual_df, expected_df)
+
+
+def test_singular_column_projection_simple(c):
+    df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+    c.create_table("df", df)
+
+    wildcard_result = c.sql("SELECT * from df")
+    single_col_result = c.sql("SELECT b from df")
+
+    assert_eq(wildcard_result["b"], single_col_result["b"])
+
+
+@pytest.mark.parametrize(
+    "input_cols",
+    [
+        ["a"],
+        ["a", "b"],
+        ["a", "d"],
+        ["d", "a"],
+        ["a", "b", "d"],
+    ],
+)
+def test_multiple_column_projection(c, input_cols):
+    projection_list = ", ".join(input_cols)
+    result = c.sql(f"SELECT {projection_list} from parquet_ddf")
+
+    # There are 5 columns in the table, ensure only specified ones are read
+    assert_eq(len(result.columns), len(input_cols))
+    assert all(x in input_cols for x in result.columns)