Datafusion Python

from datafusion import SessionContext, col, lit
from datafusion import functions as F
from datafusion.functions import spark   # only when porting pyspark code

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")

ctx.register_parquet("my_table", "path/to/data.parquet")
ctx.register_csv("my_table", "path/to/data.csv")

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")

df.filter(col("a") > 10)
df.filter(col("a") > 10, col("b") == "x")   # multiple = AND
df.filter("a > 10")                          # SQL expression string

# 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")],
)

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

# 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")

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
)

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)

df.limit(10)            # first 10 rows
df.limit(10, offset=20) # skip 20, then take 10

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
)

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

# 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())

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")])
by_customer = base.aggregate(["customer"], [F.sum(col("amount")).alias("total")])

df.write_parquet("output.parquet")
df.write_csv("output.csv")
df.write_json("output.json")

col("column_name")              # reference a column
lit(42)                          # integer literal
lit("hello")                     # string literal
lit(3.14)                        # float literal
lit(pa.scalar(value))            # PyArrow scalar (preserves Arrow type)

col("price") * col("quantity")            # multiplication
col("a") + lit(1)                          # addition
col("a") - col("b")                        # subtraction
col("a") / lit(2)                          # division
col("a") % lit(3)                          # modulo

import pyarrow as pa

# Subtract 90 days from a date column
col("ship_date") - lit(pa.scalar((0, 90, 0), type=pa.month_day_nano_interval()))

# Subtract 3 months
col("ship_date") - lit(pa.scalar((3, 0, 0), type=pa.month_day_nano_interval()))

col("a") > 10
col("a") >= 10
col("a") < 10
col("a") <= 10
col("a") == "x"
col("a") != "x"
col("a") == None                           # same as col("a").is_null()
col("a") != None                           # same as col("a").is_not_null()

(col("a") > 1) & (col("b") < 10)   # AND
(col("a") > 1) | (col("b") < 10)   # OR
~(col("a") > 1)                    # NOT

col("a").is_null()
col("a").is_not_null()
col("a").fill_null(lit(0))          # replace NULL with a value (single expression)
F.coalesce(col("a"), col("b"))     # first non-null value
F.nullif(col("a"), lit(0))         # return NULL if a == 0

df.fill_null(0)                     # every column
df.fill_null(0, subset=["a", "b"])  # only these columns

# Simple CASE (matching on a single expression)
status_label = (
    F.case(col("status"))
    .when(lit("A"), lit("Active"))
    .when(lit("I"), lit("Inactive"))
    .otherwise(lit("Unknown"))
)

# Searched CASE (each branch has its own predicate)
severity = (
    F.when(col("value") > 100, lit("high"))
    .when(col("value") > 50, lit("medium"))
    .otherwise(lit("low"))
)

import pyarrow as pa

col("a").cast(pa.float64())
col("a").cast(pa.utf8())
col("a").cast(pa.date32())

col("a").try_cast(pa.int32())        # like cast(), but yields NULL on failure instead of erroring

df.cast({"a": pa.float64(), "b": pa.int32()})

(col("a") + col("b")).alias("total")

col("a").between(1, 10)                                 # 1 <= a <= 10 (bounds auto-wrap)
F.in_list(col("a"), [lit(1), lit(2), lit(3)])           # a IN (1, 2, 3)
F.in_list(col("a"), [lit(1), lit(2)], negated=True)     # a NOT IN (1, 2)

col("struct_col")["field_name"]     # access struct field
col("array_col")[0]                  # access array element (0-indexed)
col("array_col")[1:3]                # array slice (0-indexed)

F.array_transform(col("a"), lambda v: v * 2)    # map: [1,2,3] -> [2,4,6]
F.array_filter(col("a"), lambda v: v > 2)        # filter: [1,2,3] -> [3]
F.array_any_match(col("a"), lambda v: v > 3)     # predicate: any element > 3

F.array_transform(col("a"), F.lambda_(["v"], F.lambda_var("v") * lit(2)))