Chalk DataFrame Reference

from chalkdf import DataFrame
# Simple dictionary input
df = DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
# Or use the explicit class method (recommended)
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})

import pyarrow as pa
from chalkdf import DataFrame
schema = pa.schema([("user_id", pa.int64()), ("score", pa.float64())])
df = DataFrame.named_table("users", schema)
# Build a query plan
from chalk.features import _
result_plan = df.filter(_.score > 0.5)
# Inject real data at execution time
import pyarrow as pa
data = pa.table({"user_id": [1, 2, 3], "score": [0.3, 0.8, 0.6]})
result = result_plan.run(tables={"users": data})

import pyarrow as pa
from chalkdf import DataFrame
table = pa.table({"x": [1, 2, 3], "y": ["a", "b", "c"]})
df = DataFrame.from_arrow(table)

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": ["a", "b", "c"]})

import pyarrow as pa
from chalkdf import DataFrame
async def generate_data():
    for i in range(3):
        yield {"x": [i * 10, i * 10 + 1], "y": [i, i]}
schema = pa.schema([("x", pa.int64()), ("y", pa.int64())])
df = DataFrame.from_python_udf(generate_data, schema)
result = df.run()
# Example with PyArrow RecordBatches
async def generate_batches():
    batch1 = pa.RecordBatch.from_pydict({"a": [1, 2], "b": [3, 4]})
    batch2 = pa.RecordBatch.from_pydict({"a": [5, 6], "b": [7, 8]})
    yield batch1
    yield batch2
schema = pa.schema([("a", pa.int64()), ("b", pa.int64())])
df = DataFrame.from_python_udf(generate_batches, schema)
# Example with custom timeout
df = DataFrame.from_python_udf(generate_data, schema, output_timeout=60.0)

from chalkdf import DataFrame
# Scan Parquet files
df = DataFrame.scan(["data/sales_2024.parquet"], name="sales_data")
# Scan CSV with explicit schema
import pyarrow as pa
schema = pa.schema([("id", pa.int64()), ("name", pa.string())])
df = DataFrame.scan(["data/users.csv"], schema=schema)

from chalkdf import DataFrame
from libchalk.chalksql import ChalkSqlCatalog
catalog = ChalkSqlCatalog()
df = DataFrame.from_catalog_table("users", catalog=catalog)

from chalkdf import DataFrame
orders = DataFrame.from_dict({"order_id": [1, 2, 3], "amount": [10.0, 20.0, 5.0]})
result = DataFrame.from_sql(
    "SELECT order_id, amount FROM orders WHERE amount > 8",
    orders=orders,
)

users = DataFrame.from_dict({"id": [1, 2], "name": ["Alice", "Bob"]})
purchases = DataFrame.from_dict({"user_id": [1, 1, 2], "item": ["a", "b", "c"]})
result = DataFrame.from_sql(
    "SELECT u.name, p.item FROM users u JOIN purchases p ON u.id = p.user_id",
    users=users,
    purchases=purchases,
)

import pyarrow as pa
from chalkdf import DataFrame
from chalk.sql import PostgreSQLSource
source = PostgreSQLSource(...)
schema = pa.schema([("user_id", pa.int64()), ("name", pa.string())])
df = DataFrame.from_datasource(source, "SELECT * FROM users", schema)

import pyarrow as pa
from chalkdf import DataFrame
from libchalk.sql import ConnectionPool
from libchalk.sql.snowflake import make_snowflake_connection_factory
factory = make_snowflake_connection_factory(uri="snowflake://user:pass@acct/db")
pool = ConnectionPool(factory, max_pool_size=1)
schema = pa.schema([("user_id", pa.int64()), ("name", pa.string())])
df = DataFrame.scan_from_sql(
    "COPY INTO @stage/out/ FROM (SELECT user_id, name FROM users)",
    pool=pool,
    output_uri_prefix="s3://my-bucket/unload-output/",
    schema=schema,
)

from chalkdf import DataFrame
from chalk.streams import KafkaSource
source = KafkaSource(
    name="my_kafka",
    bootstrap_server="localhost:9092",
    topic="my_topic",
)
# Pull 100 messages, just the raw bytes
df = DataFrame.from_stream_source(source, n=100)
# Pull with full metadata
df = DataFrame.from_stream_source(source, n=100, include_metadata=True)

from chalkdf import DataFrame
from chalkdf.config import CompilationConfig
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
compiled = df.compile(config=CompilationConfig(use_online_hash_join=True))
print(compiled.explain_logical())

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
filtered = df.filter(_.x > 1)
print(filtered.explain_logical())

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
filtered = df.filter(_.x > 1)
print(filtered.explain_physical())

from chalkdf import DataFrame
from chalk.features import _, F
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
# Add a new column using underscore syntax
df2 = df.with_columns({"z": _.x + _.y})
# Add a column using an F function
df3 = df.with_columns({"z_capped": F.least(_.x + _.y, 8)})
# Add a column using .alias()
df4 = df.with_columns((_.x * 2).alias("x_doubled"))
# Both df2, df3, df4 still contain x and y in addition to the new column

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [10, 20, 30]})
df_with_id = df.with_unique_id("row_id")

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"x": [1, 2, 3, 4], "y": [10, 20, 30, 40]})
filtered = df.filter(_.x > 2)

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3, 4, 5]})
# Get rows 1-3 (indices 1, 2, 3)
sliced = df.slice(1, 3)

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
# Reference a column by name to build expressions
col_x = df.col("x")
df_filtered = df.filter(col_x > 1)
# Useful when the column name comes from a variable
target = "y"
df2 = df.with_columns({"doubled": df.col(target) * 2})

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
# Compute a sum from two columns referenced by name
df2 = df.with_columns({"sum": df.col("x") + df.col("y")})

df1 = DataFrame({"x": [1, 2], "y": [10, 20]})
df2 = DataFrame({"x": [3, 4], "y": [30, 40]})
df3 = DataFrame({"x": [5], "y": [50]})
result = df1.union_all(df2, df3)
# result contains all 5 rows from df1, df2, and df3, in any order

df1 = DataFrame({"x": [1, 2], "y": [10, 20]})
df2 = DataFrame({"x": [3, 4], "y": [30, 40]})
result = df1.union(df2)
# result contains all 4 rows from df1 and df2, in any order

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
# Keep only "sum" and "x"; z is dropped
projected = df.project({"sum": _.x + _.y, "x": _.x})

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
selected = df.select("x", "y")

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6], "z": [7, 8, 9]})
df_dropped = df.drop("z")           # raises if "z" missing
df_dropped = df.drop("z", "missing_col", strict=False)  # silently skips

from chalkdf import DataFrame
df = DataFrame.from_dict({"id": [1, 2], "items": [[10, 20], [30]]})
exploded = df.explode("items")

from chalkdf import DataFrame
from libchalk.chalktable import WindowExpr
df = DataFrame.from_dict({
    "idx": [1, 1, 2, 2],
    "v":   [10, 20, 30, 40],
})
# Partition by "idx", sort by "v" ascending, shift "v" by -1 into "v_shifted"
result = df.window(["idx"], ["v"], WindowExpr.shift("v", "v_shifted", -1))
# result schema: idx, v, v_shifted
# v_shifted contains the *next* value of v within each idx partition

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"group": ["A", "A", "B"], "value": [1, 2, 3]})
grouped = df.group_by("group").agg(_.value.sum().alias("total"))

df2 = DataFrame.from_dict({"g1": ["A", "A", "B"], "g2": ["X", "Y", "X"], "val": [1, 2, 3]})
result = df2.group_by("g1", "g2").agg(_.val.sum().alias("sum"))

result = df.group_by(_.group).agg(_.value.mean().alias("avg"))

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"group": ["A", "A", "B"], "value": [1, 2, 3]})
agg_df = df.agg(["group"], _.value.sum().alias("total"))
# Or with a single column:
agg_df = df.agg("group", _.value.sum().alias("total"))

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 1, 2], "y": [10, 20, 30]})
unique = df.distinct_on("x")  # one row per unique x value

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [3, 1, 2], "y": [30, 10, 20]})
# Sort by x ascending
sorted_df = df.order_by("x")
# Sort by x descending, then y ascending
sorted_df = df.order_by(("x", "desc"), "y")

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
df.write_parquet("file:///tmp/output/")  # returns None
result = df.write_parquet("gs://my-bucket/output/", return_table_write_result=True)

from chalkdf import DataFrame
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
renamed = df.rename({"x": "id", "y": "value"})
# Can also use underscore syntax for keys
renamed = df.rename({_.x: "id", _.y: "value"})

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
filtered = df.filter(_.x > 1)
arrow_table = filtered.to_arrow()

from chalkdf import DataFrame
from chalk.features import _
df = DataFrame.from_dict({"x": [1, 2, 3], "y": [4, 5, 6]})
filtered = df.filter(_.x > 1)
materialized = filtered.run()