Set of small features

docs/source/user-guide/common-operations/expressions.rst

@@ -60,6 +60,43 @@ examples for the and, or, and not operations.
  heavy_red_units = (col("color") == lit("red")) & (col("weight") > lit(42))
  not_red_units = ~(col("color") == lit("red"))
 
+Arrays
+------
+
+For columns that contain arrays of values, you can access individual elements of the array by index
+using bracket indexing. This is similar to callling the function
+:py:func:`datafusion.functions.array_element`, except that array indexing using brackets is 0 based,
+similar to Python arrays and ``array_element`` is 1 based indexing to be compatible with other SQL
+approaches.
+
+.. ipython:: python
+
+ from datafusion import SessionContext, col
+
+ ctx = SessionContext()
+ df = ctx.from_pydict({"a": [[1, 2, 3], [4, 5, 6]]})
+ df.select(col("a")[0].alias("a0"))
+
+
+.. warning::
+
+ Indexing an element of an array via ``[]`` starts at index 0 whereas
+ :py:func:`~datafusion.functions.array_element` starts at index 1.
+
+Structs
+-------
+
+Columns that contain struct elements can be accessed using the bracket notation as if they were
+Python dictionary style objects. This expects a string key as the parameter passed.
+
+.. ipython:: python
+
+ ctx = SessionContext()
+ data = {"a": [{"size": 15, "color": "green"}, {"size": 10, "color": "blue"}]}
+ df = ctx.from_pydict(data)
+ df.select(col("a")["size"].alias("a_size"))
+
+
 Functions
 ---------
 

python/datafusion/dataframe.py

@@ -30,6 +30,7 @@
  import pandas as pd
  import polars as pl
  import pathlib
+ from typing import Callable
 
 from datafusion._internal import DataFrame as DataFrameInternal
 from datafusion.expr import Expr
@@ -72,6 +73,9 @@ def __repr__(self) -> str:
  """
  return self.df.__repr__()
 
+ def _repr_html_(self) -> str:
+ return self.df._repr_html_()
+
  def describe(self) -> DataFrame:
  """Return the statistics for this DataFrame.
 
@@ -539,3 +543,25 @@ def __arrow_c_stream__(self, requested_schema: pa.Schema) -> Any:
  Arrow PyCapsule object.
  """
  return self.df.__arrow_c_stream__(requested_schema)
+
+ def transform(self, func: Callable[..., DataFrame], *args: Any) -> DataFrame:
+ """Apply a function to the current DataFrame which returns another DataFrame.
+
+ This is useful for chaining together multiple functions. For example::
+
+ def add_3(df: DataFrame) -> DataFrame:
+ return df.with_column("modified", lit(3))
+
+ def within_limit(df: DataFrame, limit: int) -> DataFrame:
+ return df.filter(col("a") < lit(limit)).distinct()
+
+ df = df.transform(modify_df).transform(within_limit, 4)
+
+ Args:
+ func: A callable function that takes a DataFrame as it's first argument
+ args: Zero or more arguments to pass to `func`
+
+ Returns:
+ DataFrame: After applying func to the original dataframe.
+ """
+ return func(self, *args)

python/datafusion/expr.py

@@ -22,7 +22,11 @@
 
 from __future__ import annotations
 
-from ._internal import expr as expr_internal, LogicalPlan
+from ._internal import (
+ expr as expr_internal,
+ LogicalPlan,
+ functions as functions_internal,
+)
 from datafusion.common import NullTreatment, RexType, DataTypeMap
 from typing import Any, Optional
 import pyarrow as pa
@@ -257,8 +261,17 @@ def __invert__(self) -> Expr:
  """Binary not (~)."""
  return Expr(self.expr.__invert__())
 
- def __getitem__(self, key: str) -> Expr:
- """For struct data types, return the field indicated by ``key``."""
+ def __getitem__(self, key: str | int) -> Expr:
+ """Retrieve sub-object.
+
+ If ``key`` is a string, returns the subfield of the struct.
+ If ``key`` is an integer, retrieves the element in the array. Note that the
+ element index begins at ``0``, unlike `array_element` which begines at ``1``.
+ """
+ if isinstance(key, int):
+ return Expr(
+ functions_internal.array_element(self.expr, Expr.literal(key + 1).expr)
+ )
  return Expr(self.expr.__getitem__(key))
 
  def __eq__(self, rhs: Any) -> Expr:

python/datafusion/functions.py

@@ -1023,7 +1023,7 @@ def array(*args: Expr) -> Expr:
 
  This is an alias for :py:func:`make_array`.
  """
- return make_array(args)
+ return make_array(*args)
 
 
 def range(start: Expr, stop: Expr, step: Expr) -> Expr:

python/datafusion/tests/test_dataframe.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import os
+from typing import Any
 
 import pyarrow as pa
 from pyarrow.csv import write_csv
@@ -970,3 +971,34 @@ def test_dataframe_export(df) -> None:
  except Exception:
  failed_convert = True
  assert failed_convert
+
+
+def test_dataframe_transform(df):
+ def add_string_col(df_internal) -> DataFrame:
+ return df_internal.with_column("string_col", literal("string data"))
+
+ def add_with_parameter(df_internal, value: Any) -> DataFrame:
+ return df_internal.with_column("new_col", literal(value))
+
+ df = df.transform(add_string_col).transform(add_with_parameter, 3)
+
+ result = df.to_pydict()
+
+ assert result["a"] == [1, 2, 3]
+ assert result["string_col"] == ["string data" for _i in range(0, 3)]
+ assert result["new_col"] == [3 for _i in range(0, 3)]
+
+
+def test_dataframe_repr_html(df) -> None:
+ output = df._repr_html_()
+
+ ref_html = """<table border='1'>
+ <tr><th>a</td><th>b</td><th>c</td></tr>
+ <tr><td>1</td><td>4</td><td>8</td></tr>
+ <tr><td>2</td><td>5</td><td>5</td></tr>
+ <tr><td>3</td><td>6</td><td>8</td></tr>
+ </table>
+ """
+
+ # Ignore whitespace just to make this test look cleaner
+ assert output.replace(" ", "") == ref_html.replace(" ", "")

python/datafusion/tests/test_expr.py

@@ -169,3 +169,26 @@ def traverse_logical_plan(plan):
  == '[Expr(Utf8("dfa")), Expr(Utf8("ad")), Expr(Utf8("dfre")), Expr(Utf8("vsa"))]'
  )
  assert not variant.negated()
+
+
+def test_expr_getitem() -> None:
+ ctx = SessionContext()
+ data = {
+ "array_values": [[1, 2, 3], [4, 5], [6], []],
+ "struct_values": [
+ {"name": "Alice", "age": 15},
+ {"name": "Bob", "age": 14},
+ {"name": "Charlie", "age": 13},
+ {"name": None, "age": 12},
+ ],
+ }
+ df = ctx.from_pydict(data, name="table1")
+
+ names = df.select(col("struct_values")["name"].alias("name")).collect()
+ names = [r.as_py() for rs in names for r in rs["name"]]
+
+ array_values = df.select(col("array_values")[1].alias("value")).collect()
+ array_values = [r.as_py() for rs in array_values for r in rs["value"]]
+
+ assert names == ["Alice", "Bob", "Charlie", None]
+ assert array_values == [2, 5, None, None]

src/dataframe.rs

@@ -23,6 +23,7 @@ use arrow::compute::can_cast_types;
 use arrow::error::ArrowError;
 use arrow::ffi::FFI_ArrowSchema;
 use arrow::ffi_stream::FFI_ArrowArrayStream;
+use arrow::util::display::{ArrayFormatter, FormatOptions};
 use datafusion::arrow::datatypes::Schema;
 use datafusion::arrow::pyarrow::{PyArrowType, ToPyArrow};
 use datafusion::arrow::util::pretty;
@@ -95,6 +96,51 @@ impl PyDataFrame {
  }
  }
 
+ fn _repr_html_(&self, py: Python) -> PyResult<String> {
+ let mut html_str = "<table border='1'>\n".to_string();
+
+ let df = self.df.as_ref().clone().limit(0, Some(10))?;
+ let batches = wait_for_future(py, df.collect())?;
+
+ if batches.is_empty() {
+ html_str.push_str("</table>\n");
+ return Ok(html_str);
+ }
+
+ let schema = batches[0].schema();
+
+ let mut header = Vec::new();
+ for field in schema.fields() {
+ header.push(format!("<th>{}</td>", field.name()));
+ }
+ let header_str = header.join("");
+ html_str.push_str(&format!("<tr>{}</tr>\n", header_str));
+
+ for batch in batches {
+ let formatters = batch
+ .columns()
+ .iter()
+ .map(|c| ArrayFormatter::try_new(c.as_ref(), &FormatOptions::default()))
+ .map(|c| {
+ c.map_err(|e| PyValueError::new_err(format!("Error: {:?}", e.to_string())))
+ })
+ .collect::<Result<Vec<_>, _>>()?;
+
+ for row in 0..batch.num_rows() {
+ let mut cells = Vec::new();
+ for formatter in &formatters {
+ cells.push(format!("<td>{}</td>", formatter.value(row)));
+ }
+ let row_str = cells.join("");
+ html_str.push_str(&format!("<tr>{}</tr>\n", row_str));
+ }
+ }
+
+ html_str.push_str("</table>\n");
+
+ Ok(html_str)
+ }
+
  /// Calculate summary statistics for a DataFrame
  fn describe(&self, py: Python) -> PyResult<Self> {
  let df = self.df.as_ref().clone();