datafusion-python

# DataFusion Python DataFrame API Guide ## What Is DataFusion? DataFusion is an **in-process query engine** built on Apache Arrow. It is not a database -- there is no server, no connection string, and no external dependencies. You create a `SessionContext`, point it at data (Parquet, CSV, JSON, Arrow IPC, Pandas, Polars, or raw Python dicts/lists), and run queries using either SQL or the DataFrame API described below. All data flows through **Apache Arrow**. The canonical Python implementation is PyArrow (`pyarrow.RecordBatch` / `pyarrow.Table`), but any library that conforms to the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) can interoperate with DataFusion. ## Core Abstractions | Abstraction | Role | Key import | |---|---|---| | `SessionContext` | Entry point. Loads data, runs SQL, produces DataFrames. | `from datafusion import SessionContext` | | `DataFrame` | Lazy query builder. Each method returns a new DataFrame. | Returned by context methods | | `Expr` | Expression tree node (column ref, literal, function call, ...). | `from datafusion import col, lit` | | `functions` | 290+ built-in scalar, aggregate, and window functions. | `from datafusion import functions as F` | | `functions.spark` | PySpark-compatible function surface (parameter names match `pyspark.sql.functions`). | `from datafusion.functions import spark` | ## Import Conventions ```python from datafusion import SessionContext, col, lit from datafusion import functions as F from datafusion.functions import spark # only when porting pyspark code ``` ## Data Loading ```python ctx = SessionContext() # From files df = ctx.read_parquet("path/to/data.parquet") df = ctx.read_csv("path/to/data.csv") df = ctx.read_json("path/to/data.json") # From Python objects df = ctx.from_pydict({"a": [1, 2, 3], "b": ["x", "y", "z"]}) df = ctx.from_pylist([{"a": 1, "b": "x"}, {"a": 2, "b": "y"}]) df = ctx.from_pandas(pandas_df) df = ctx.from_polars(polars_df) df = ctx.from_arrow(arrow_table) df = ctx.read_batch(record_batch) # one pa.RecordBatch, no named table df = ctx.read_batches([batch1, batch2]) # several pa.RecordBatch # From SQL df = ctx.sql("SELECT a, b FROM my_table WHERE a > 1") ``` To make a DataFrame queryable by name in SQL, register it first: ```python ctx.register_parquet("my_table", "path/to/data.parquet") ctx.register_csv("my_table", "path/to/data.csv") ``` ## DataFrame Operations Quick Reference Every method returns a **new** DataFrame (immutable/lazy). Chain them fluently. ### Projection ```python df.select("a", "b") # preferred: plain names as strings df.select(col("a"), (col("b") + 1).alias("b_plus_1")) # use col()/Expr only when you need an expression df.with_column("new_col", col("a") + lit(10)) # add one column df.with_columns( col("a").alias("x"), y=col("b") + lit(1), # named keyword form ) df.drop("unwanted_col") df.with_column_renamed("old_name", "new_name") ``` When a column is referenced by name alone, pass the name as a string rather than wrapping it in `col()`. Reach for `col()` only when the projection needs arithmetic, aliasing, casting, or another expression operation. **Case sensitivity**: both `select("Name")` and `col("Name")` lowercase the identifier. For a column whose real name has uppercase letters, embed double quotes inside the string: `select('"MyCol"')` or `col('"MyCol"')`. Without the inner quotes the lookup will fail with `No field named mycol`. ### Filtering ```python df.filter(col("a") > 10) df.filter(col("a") > 10, col("b") == "x") # multiple = AND df.filter("a > 10") # SQL expression string ``` Raw Python values on the right-hand side of a comparison are auto-wrapped into literals by the `Expr` operators, so prefer `col("a") > 10` over `col("a") > lit(10)`. See the Comparisons section and pitfall #2 for the full rule. ### Aggregation ```python # GROUP BY a, compute sum(b) and count(*) df.aggregate(["a"], [F.sum(col("b")), F.count(col("a"))]) # HAVING equivalent: use the filter keyword on the aggregate function df.aggregate( ["region"], [F.sum(col("sales"), filter=col("sales") > 1000).alias("large_sales")], ) ``` As with `select()`, group keys can be passed as plain name strings. Reach for `col(...)` only when the grouping expression needs arithmetic, aliasing, casting, or another expression operation. Most aggregate functions accept an optional `filter` keyword argument. When provided, only rows where the filter expression is true contribute to the aggregate. ### Sorting ```python df.sort("a") # ascending (plain name, preferred) df.sort(col("a")) # ascending via col() df.sort(col("a").sort(ascending=False)) # descending df.sort(col("a").sort(nulls_first=False)) # override null placement df.sort_by("a", "b") # ascending-only shortcut ``` As with `select()` and `aggregate()`, bare column references can be passed as plain name strings. A plain expression passed to `sort()` is already treated as ascending, so reach for `col(...).sort(...)` only when you need to override a default (descending order or null placement). Writing `col("a").sort(ascending=True)` is redundant. For ascending-only sorts with no null-placement override, `df.sort_by(...)` is a shorter alias for `df.sort(...)`. ### Joining ```python # Equi-join on shared column name df1.join(df2, on="key") df1.join(df2, on="key", how="left") # Different column names df1.join(df2, left_on="id", right_on="fk_id", how="inner") # Expression-based join (supports inequality predicates) df1.join_on(df2, col("a") == col("b"), how="inner") # Semi join: keep rows from left where a match exists in right (like EXISTS) df1.join(df2, on="key", how="semi") # Anti join: keep rows from left where NO match exists in right (like NOT EXISTS) df1.join(df2, on="key", how="anti") ``` Join types: `"inner"`, `"left"`, `"right"`, `"full"`, `"semi"`, `"anti"`. Inner is the default `how`. Prefer `df1.join(df2, on="key")` over `df1.join(df2, on="key", how="inner")` — drop `how=` unless you need a non-inner join type. When the two sides' join columns have different native names, use `left_on=`/`right_on=` with the original names rather than aliasing one side to match the other — see pitfall #7. ### Window Functions ```python from datafusion import WindowFrame # Row number partitioned by group, ordered by value df.window( F.row_number( partition_by=[col("group")], order_by=[col("value")], ).alias("rn") ) # Using a Window object for reuse from datafusion.expr import Window win = Window( partition_by=[col("group")], order_by=[col("value").sort(ascending=True)], ) df.select( col("group"), col("value"), F.sum(col("value")).over(win).alias("running_total"), ) # With explicit frame bounds win = Window( partition_by=[col("group")], order_by=[col("value").sort(ascending=True)], window_frame=WindowFrame("rows", 0, None), # current row to unbounded following ) ``` ### Set Operations ```python df1.union(df2) # UNION ALL (by position) df1.union(df2, distinct=True) # UNION DISTINCT df1.union_by_name(df2) # match columns by name, not position df1.intersect(df2) # INTERSECT ALL df1.intersect(df2, distinct=True) # INTERSECT (distinct) df1.except_all(df2) # EXCEPT ALL df1.except_all(df2, distinct=True) # EXCEPT (distinct) ``` ### Limit and Offset ```python df.limit(10) # first 10 rows df.limit(10, offset=20) # skip 20, then take 10 ``` ### Deduplication ```python df.distinct() # remove duplicate rows df.distinct_on( # keep first row per group (like DISTINCT ON in Postgres) [col("a")], # uniqueness columns [col("a"), col("b")], # output columns [col("b").sort(ascending=True)], # which row to keep ) ``` ## Executing and Collecting Results DataFrames are lazy until you collect. ```python df.show() # print formatted table to stdout batches = df.collect() # list[pa.RecordBatch] arr = df.collect_column("col_name") # pa.Array | pa.ChunkedArray (single column) table = df.to_arrow_table() # pa.Table pandas_df = df.to_pandas() # pd.DataFrame polars_df = df.to_polars() # pl.DataFrame py_dict = df.to_pydict() # dict[str, list] py_list = df.to_pylist() # list[dict] count = df.count() # int df = df.cache() # materialize in memory, return DataFrame ``` ### Date and Timestamp Type Conversion The Python type returned by `to_pydict()` / `to_pylist()` depends on the Arrow column type, and the mapping is inherited from PyArrow: | Arrow type | Python type returned | |---|---| | `timestamp(s)` / `(ms)` / `(us)` | `datetime.datetime` | | `timestamp(ns)` | `pandas.Timestamp` | | `date32` / `date64` | `datetime.date` | | `duration(s)` / `(ms)` / `(us)` | `datetime.timedelta` | | `duration(ns)` | `pandas.Timedelta` | The nanosecond-precision fallback to pandas types is the main surprise: pandas is not a hard dependency of `datafusion`, but PyArrow reaches for it when `datetime.datetime` / `datetime.timedelta` would lose precision (stdlib types only go to microseconds). If you need plain stdlib types, cast to a coarser unit before collecting, e.g. `df.select(col("ts").cast(pa.timestamp("us")))`. `df.to_pandas()` has its own footgun for dates: pandas has no pure-date dtype, so a `date32`/`date64` column comes back as an `object` column of `datetime.date` values rather than `datetime64[ns]`. If downstream code expects a datetime column, cast on the DataFusion side first: `col("ship_date").cast(pa.timestamp("ns"))`. ### Streaming Results Prefer streaming over `collect()` when the result is too large to materialize in memory, when you want to start processing before the query finishes, or when you may break out of the loop early. `execute_stream()` pulls one `RecordBatch` at a time from the execution plan rather than buffering the whole result up front. ```python # Single-partition stream; batch is a datafusion.RecordBatch stream = df.execute_stream() for batch in stream: process(batch.to_pyarrow()) # convert to pa.RecordBatch if needed # DataFrame is iterable directly (delegates to execute_stream) for batch in df: process(batch.to_pyarrow()) # One stream per partition, for parallel consumption for stream in df.execute_stream_partitioned(): for batch in stream: process(batch.to_pyarrow()) ``` Async iteration is also supported via `async for batch in df: ...` (or `df.execute_stream()`), which is useful when batches are interleaved with other I/O. ### Caching Intermediate Results `df.cache()` materializes a DataFrame as an in-memory table and returns a new DataFrame backed by it. Reach for it when the same intermediate result feeds multiple downstream queries — without `cache()`, each branch re-executes the full upstream plan (re-reading files, recomputing filters/aggregates). ```python base = ( ctx.read_parquet("orders.parquet") .filter(col("status") == "shipped") .cache() # materialize once, reuse below ) by_region = base.aggregate(["region"], [F.sum(col("amount")).alias("total")])