add indexed fields support to python api (apache#1502)

nl5887 · web-flow · commit 7374b18f1dca · 2021-12-31T10:40:25.000+08:00
* add nested struct support to python

implements nested structs `col("a")['b']`

* add test for indexed fields
diff --git a/python/datafusion/tests/test_dataframe.py b/python/datafusion/tests/test_dataframe.py
@@ -35,6 +35,19 @@ def df():
     return ctx.create_dataframe([[batch]])
 
 
+@pytest.fixture
+def struct_df():
+    ctx = ExecutionContext()
+
+    # create a RecordBatch and a new DataFrame from it
+    batch = pa.RecordBatch.from_arrays(
+        [pa.array([{"c": 1}, {"c": 2}, {"c": 3}]), pa.array([4, 5, 6])],
+        names=["a", "b"],
+    )
+
+    return ctx.create_dataframe([[batch]])
+
+
 def test_select(df):
     df = df.select(
         column("a") + column("b"),
@@ -153,3 +166,16 @@ def test_get_dataframe(tmp_path):
 
     df = ctx.table("csv")
     assert isinstance(df, DataFrame)
+
+
+def test_struct_select(struct_df):
+    df = struct_df.select(
+        column("a")["c"] + column("b"),
+        column("a")["c"] - column("b"),
+    )
+
+    # execute and collect the first (and only) batch
+    result = df.collect()[0]
+
+    assert result.column(0) == pa.array([5, 7, 9])
+    assert result.column(1) == pa.array([-3, -3, -3])
diff --git a/python/src/expression.rs b/python/src/expression.rs
@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use pyo3::PyMappingProtocol;
 use pyo3::{basic::CompareOp, prelude::*, PyNumberProtocol, PyObjectProtocol};
 use std::convert::{From, Into};
 
@@ -133,3 +134,14 @@ impl PyExpr {
         expr.into()
     }
 }
+
+#[pyproto]
+impl PyMappingProtocol for PyExpr {
+    fn __getitem__(&self, key: &str) -> PyResult<PyExpr> {
+        Ok(Expr::GetIndexedField {
+            expr: Box::new(self.expr.clone()),
+            key: ScalarValue::Utf8(Some(key.to_string()).to_owned()),
+        }
+        .into())
+    }
+}
