Skip to content

Commit 83032c2

Browse files
Prajwal Narayanaclaude
andcommitted
feat: run examples/*.py in CI and fix the examples that could not run
No workflow executed any top-level `examples/*.py`, so two problems had gone unnoticed: `csv-read-options.py` raised on its first read because the `data.csv` and `data.csv.gz` it reads are not in the repository, and nine scripts printed nothing at all, ending in a bare `assert`. - `csv-read-options.py` now writes its own CSV, pipe-separated CSV and gzipped CSV into a temporary directory, and shows the result of each example. - `export.py`, `import.py`, `python-udaf.py`, `python-udf.py`, `query-pyarrow-data.py`, `sql-to-pandas.py`, `sql-using-python-udaf.py`, `sql-using-python-udf.py` and `substrait.py` print their results, keeping the existing asserts. - `substrait.py` resolves its submodule data relative to `__file__` instead of the working directory. - A `Run Python examples` step runs the examples from the repository root after the TPC-H data is generated, skipping the five that need a manual download, a Ray cluster, or AWS credentials. - `examples/README.md` documents the per-example prerequisites and drops the links to `sql-on-polars.py`, `sql-on-pandas.py` and `sql-on-cudf.py`, which are not in the repository. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
1 parent 7022baa commit 83032c2

12 files changed

Lines changed: 155 additions & 15 deletions

.github/workflows/test.yml

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,3 +141,45 @@ jobs:
141141
run: |
142142
cd examples/tpch
143143
uv run --no-project pytest _tests.py
144+
145+
# The top-level examples/*.py are documentation users copy from, so they
146+
# have to keep working. Every one of them is self-contained except the
147+
# five skipped below; run them from the repository root, after the TPC-H
148+
# data exists, against the wheel built above.
149+
- name: Run Python examples
150+
if: matrix.wheel-tag == 'abi3'
151+
run: |
152+
# pandas and polars are neither runtime nor dev dependencies, but the
153+
# import and export examples demonstrate converting to and from them.
154+
uv pip install --python "$PWD/.venv/bin/python" pandas polars
155+
# Skipped, and why each one cannot run here:
156+
# dataframe-parquet.py, sql-parquet.py and sql-to-pandas.py need
157+
# yellow_tripdata_2021-01.parquet, a manual download documented
158+
# in examples/README.md
159+
# ray_pickle_expr.py needs a Ray cluster
160+
# sql-parquet-s3.py needs network access and AWS credentials
161+
skipped="
162+
dataframe-parquet.py
163+
ray_pickle_expr.py
164+
sql-parquet-s3.py
165+
sql-parquet.py
166+
sql-to-pandas.py
167+
"
168+
failed=""
169+
for example in examples/*.py; do
170+
name=$(basename "$example")
171+
# shellcheck disable=SC2086 # intentional split on whitespace
172+
if printf '%s\n' $skipped | grep -qx "$name"; then
173+
echo "::notice title=Example skipped::$name"
174+
continue
175+
fi
176+
echo "::group::$name"
177+
if ! uv run --python "$PWD/.venv/bin/python" --no-project python "$example"; then
178+
failed="$failed $name"
179+
fi
180+
echo "::endgroup::"
181+
done
182+
if [ -n "$failed" ]; then
183+
echo "::error title=Examples failed::$failed"
184+
exit 1
185+
fi

examples/README.md

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,14 +19,37 @@
1919

2020
# DataFusion Python Examples
2121

22-
Some examples rely on data which can be downloaded from the following site:
22+
## Running the examples
23+
24+
Every example is a standalone script; run it from the root of the repository:
25+
26+
```bash
27+
python examples/create-context.py
28+
```
29+
30+
Most of them need nothing but the `datafusion` package and create their own
31+
data. The exceptions are:
32+
33+
| Example | Needs |
34+
| --- | --- |
35+
| `dataframe-parquet.py`, `sql-parquet.py`, `sql-to-pandas.py` | `yellow_tripdata_2021-01.parquet`, downloaded into the working directory (see below) |
36+
| `import.py`, `export.py`, `sql-to-pandas.py` | `pandas`, `polars` (`sql-to-pandas.py` also needs `matplotlib`) |
37+
| `python-udf-comparisons.py` | the TPC-H dataset in `examples/tpch/data/`, see [`tpch/README.md`](./tpch/README.md) |
38+
| `ray_pickle_expr.py` | `ray` |
39+
| `sql-parquet-s3.py` | network access and AWS credentials in the environment |
40+
| `substrait.py` | the `testing` submodule: `git submodule update --init testing` |
41+
42+
The NYC taxi data can be downloaded from the following site:
2343

2444
- https://www.nyc.gov/site/tlc/about/tlc-trip-record-data.page
2545

2646
Here is a direct link to the file used in the examples:
2747

2848
- https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2021-01.parquet
2949

50+
Everything that does not need a manual download or a cluster is run on every
51+
pull request by the `Run Python examples` step in `.github/workflows/test.yml`.
52+
3053
### Creating a SessionContext
3154

3255
- [Creating a SessionContext](./create-context.py)
@@ -62,12 +85,6 @@ type and codec boundaries rather than same-library Rust downcasts.
6285

6386
- [Serialize query plans using Substrait](./substrait.py)
6487

65-
### Executing SQL against DataFrame Libraries (Experimental)
66-
67-
- [Executing SQL on Polars](./sql-on-polars.py)
68-
- [Executing SQL on Pandas](./sql-on-pandas.py)
69-
- [Executing SQL on cuDF](./sql-on-cudf.py)
70-
7188
## TPC-H Examples
7289

7390
Within the subdirectory `tpch` there are 22 examples that reproduce queries in

examples/csv-read-options.py

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,17 +15,49 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
"""Example demonstrating CsvReadOptions usage."""
18+
"""Example demonstrating CsvReadOptions usage.
19+
20+
The example writes the small CSV files it reads into a temporary directory, so
21+
it is self-contained and can be run from any working directory.
22+
"""
23+
24+
import gzip
25+
import tempfile
26+
from pathlib import Path
1927

2028
from datafusion import CsvReadOptions, SessionContext
2129

30+
# Write the sample data this example reads into a temporary directory, rather
31+
# than checking the files into the repository and reading them by a relative
32+
# path that only resolves from one working directory.
33+
tmp_dir = tempfile.TemporaryDirectory()
34+
data_dir = Path(tmp_dir.name)
35+
36+
# Comma separated, quoted with `"`, used by most of the examples below.
37+
csv_file = data_dir / "data.csv"
38+
csv_file.write_text('id,name,value\n1,"alice",10\n2,"bob",20\n3,"carol",30\n')
39+
40+
# Pipe separated and quoted with `'`, to exercise the builder pattern.
41+
pipe_file = data_dir / "data_pipe.csv"
42+
pipe_file.write_text(
43+
"id|name|value\n1|'alice'|10\n2|'bob, the second'|20\n3|'carol'|30\n"
44+
)
45+
46+
# Gzipped, with a comment line and an `N/A` placeholder, to exercise the
47+
# advanced options.
48+
gzip_file = data_dir / "data.csv.gz"
49+
gzip_file.write_bytes(
50+
gzip.compress(b"# sample data\nid,name,value\n1,alice,10\n2,N/A,20\n3,carol,30\n")
51+
)
52+
2253
# Create a SessionContext
2354
ctx = SessionContext()
2455

2556
# Example 1: Using CsvReadOptions with default values
2657
print("Example 1: Default CsvReadOptions")
2758
options = CsvReadOptions()
28-
df = ctx.read_csv("data.csv", options=options)
59+
df = ctx.read_csv(csv_file, options=options)
60+
df.show()
2961

3062
# Example 2: Using CsvReadOptions with custom parameters
3163
print("\nExample 2: Custom CsvReadOptions")
@@ -36,7 +68,8 @@
3668
schema_infer_max_records=1000,
3769
file_extension=".csv",
3870
)
39-
df = ctx.read_csv("data.csv", options=options)
71+
df = ctx.read_csv(csv_file, options=options)
72+
df.show()
4073

4174
# Example 3: Using the builder pattern (recommended for readability)
4275
print("\nExample 3: Builder pattern")
@@ -49,7 +82,8 @@
4982
.with_truncated_rows(False) # noqa: FBT003
5083
.with_newlines_in_values(True) # noqa: FBT003
5184
)
52-
df = ctx.read_csv("data.csv", options=options)
85+
df = ctx.read_csv(pipe_file, options=options)
86+
df.show()
5387

5488
# Example 4: Advanced options
5589
print("\nExample 4: Advanced options")
@@ -64,18 +98,23 @@
6498
.with_file_compression_type("gzip") # Read gzipped CSV
6599
.with_file_extension(".gz")
66100
)
67-
df = ctx.read_csv("data.csv.gz", options=options)
101+
df = ctx.read_csv(gzip_file, options=options)
102+
df.show()
68103

69104
# Example 5: Register CSV table with options
70105
print("\nExample 5: Register CSV table")
71106
options = CsvReadOptions().with_has_header(True).with_delimiter(",") # noqa: FBT003
72-
ctx.register_csv("my_table", "data.csv", options=options)
107+
ctx.register_csv("my_table", csv_file, options=options)
73108
df = ctx.sql("SELECT * FROM my_table")
109+
df.show()
74110

75111
# Example 6: Backward compatibility (without options)
76112
print("\nExample 6: Backward compatibility")
77113
# Still works the old way!
78-
df = ctx.read_csv("data.csv", has_header=True, delimiter=",")
114+
df = ctx.read_csv(csv_file, has_header=True, delimiter=",")
115+
df.show()
116+
117+
tmp_dir.cleanup()
79118

80119
print("\nAll examples completed!")
81120
print("\nFor all available options, see the CsvReadOptions documentation:")

examples/export.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,19 +34,29 @@
3434
# export to pandas dataframe
3535
pandas_df = df.to_pandas()
3636
assert pandas_df.shape == (3, 2)
37+
print("pandas DataFrame:")
38+
print(pandas_df)
3739

3840
# export to PyArrow table
3941
arrow_table = df.to_arrow_table()
4042
assert arrow_table.shape == (3, 2)
43+
print("\nPyArrow table:")
44+
print(arrow_table)
4145

4246
# export to Polars dataframe
4347
polars_df = df.to_polars()
4448
assert polars_df.shape == (3, 2)
49+
print("\nPolars DataFrame:")
50+
print(polars_df)
4551

4652
# export to Python list of rows
4753
pylist = df.to_pylist()
4854
assert pylist == [{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]
55+
print("\nPython list of rows:")
56+
print(pylist)
4957

5058
# export to Python dictionary of columns
5159
pydict = df.to_pydict()
5260
assert pydict == {"a": [1, 2, 3], "b": [4, 5, 6]}
61+
print("\nPython dictionary of columns:")
62+
print(pydict)

examples/import.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
# represent column values
2929
df = ctx.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
3030
assert type(df) is datafusion.DataFrame
31+
print("from_pydict:")
32+
df.show()
3133
# Dataframe:
3234
# +---+---+
3335
# | a | b |
@@ -40,18 +42,26 @@
4042
# Create a datafusion DataFrame from a Python list of rows
4143
df = ctx.from_pylist([{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}])
4244
assert type(df) is datafusion.DataFrame
45+
print("from_pylist:")
46+
df.show()
4347

4448
# Convert pandas DataFrame to datafusion DataFrame
4549
pandas_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
4650
df = ctx.from_pandas(pandas_df)
4751
assert type(df) is datafusion.DataFrame
52+
print("from_pandas:")
53+
df.show()
4854

4955
# Convert polars DataFrame to datafusion DataFrame
5056
polars_df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
5157
df = ctx.from_polars(polars_df)
5258
assert type(df) is datafusion.DataFrame
59+
print("from_polars:")
60+
df.show()
5361

5462
# Convert Arrow Table to datafusion DataFrame
5563
arrow_table = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
5664
df = ctx.from_arrow(arrow_table)
5765
assert type(df) is datafusion.DataFrame
66+
print("from_arrow:")
67+
df.show()

examples/python-udaf.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ def evaluate(self) -> pa.Scalar:
6464

6565
df = df.aggregate([], [my_udaf(col("a"))])
6666

67+
df.show()
68+
6769
result = df.collect()[0]
6870

6971
assert result.column(0) == pa.array([6.0])

examples/python-udf.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ def is_null(array: pa.Array) -> pa.Array:
3838

3939
df = df.select(is_null_arr(f.col("a")))
4040

41+
df.show()
42+
4143
result = df.collect()[0]
4244

4345
assert result.column(0) == pa.array([False] * 3)

examples/query-pyarrow-data.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@
3535
col("a") - col("b"),
3636
)
3737

38+
df.show()
39+
3840
# execute and collect the first (and only) batch
3941
result = df.collect()[0]
4042

examples/sql-to-pandas.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,11 @@
3434

3535
# convert to Pandas
3636
pandas_df = df.to_pandas()
37+
print(pandas_df)
3738

3839
# create a chart
3940
fig = pandas_df.plot(
4041
kind="bar", title="Trip Count by Number of Passengers"
4142
).get_figure()
4243
fig.savefig("chart.png")
44+
print("wrote chart.png")

examples/sql-using-python-udaf.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@ def evaluate(self) -> pa.Scalar:
7575
result_df = ctx.sql(
7676
"select a, my_accumulator(b) as b_aggregated from t group by a order by a"
7777
)
78+
result_df.show()
7879
# Dataframe:
7980
# +---+--------------+
8081
# | a | b_aggregated |

0 commit comments

Comments
 (0)