diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 047b35039..5373057be 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -141,3 +141,45 @@ jobs: run: | cd examples/tpch uv run --no-project pytest _tests.py + + # The top-level examples/*.py are documentation users copy from, so they + # have to keep working. Every one of them is self-contained except the + # five skipped below; run them from the repository root, after the TPC-H + # data exists, against the wheel built above. + - name: Run Python examples + if: matrix.wheel-tag == 'abi3' + run: | + # pandas and polars are neither runtime nor dev dependencies, but the + # import and export examples demonstrate converting to and from them. + uv pip install --python "$PWD/.venv/bin/python" pandas polars + # Skipped, and why each one cannot run here: + # dataframe-parquet.py, sql-parquet.py and sql-to-pandas.py need + # yellow_tripdata_2021-01.parquet, a manual download documented + # in examples/README.md + # ray_pickle_expr.py needs a Ray cluster + # sql-parquet-s3.py needs network access and AWS credentials + skipped=" + dataframe-parquet.py + ray_pickle_expr.py + sql-parquet-s3.py + sql-parquet.py + sql-to-pandas.py + " + failed="" + for example in examples/*.py; do + name=$(basename "$example") + # shellcheck disable=SC2086 # intentional split on whitespace + if printf '%s\n' $skipped | grep -qx "$name"; then + echo "::notice title=Example skipped::$name" + continue + fi + echo "::group::$name" + if ! uv run --python "$PWD/.venv/bin/python" --no-project python "$example"; then + failed="$failed $name" + fi + echo "::endgroup::" + done + if [ -n "$failed" ]; then + echo "::error title=Examples failed::$failed" + exit 1 + fi diff --git a/examples/README.md b/examples/README.md index 7bbb45dcf..4ca940705 100644 --- a/examples/README.md +++ b/examples/README.md @@ -19,7 +19,27 @@ # DataFusion Python Examples -Some examples rely on data which can be downloaded from the following site: +## Running the examples + +Every example is a standalone script; run it from the root of the repository: + +```bash +python examples/create-context.py +``` + +Most of them need nothing but the `datafusion` package and create their own +data. The exceptions are: + +| Example | Needs | +| --- | --- | +| `dataframe-parquet.py`, `sql-parquet.py`, `sql-to-pandas.py` | `yellow_tripdata_2021-01.parquet`, downloaded into the working directory (see below) | +| `import.py`, `export.py`, `sql-to-pandas.py` | `pandas`, `polars` (`sql-to-pandas.py` also needs `matplotlib`) | +| `python-udf-comparisons.py` | the TPC-H dataset in `examples/tpch/data/`, see [`tpch/README.md`](./tpch/README.md) | +| `ray_pickle_expr.py` | `ray` | +| `sql-parquet-s3.py` | network access and AWS credentials in the environment | +| `substrait.py` | the `testing` submodule: `git submodule update --init testing` | + +The NYC taxi data can be downloaded from the following site: - https://www.nyc.gov/site/tlc/about/tlc-trip-record-data.page @@ -27,6 +47,9 @@ Here is a direct link to the file used in the examples: - https://d37ci6vzurychx.cloudfront.net/trip-data/yellow_tripdata_2021-01.parquet +Everything that does not need a manual download or a cluster is run on every +pull request by the `Run Python examples` step in `.github/workflows/test.yml`. + ### Creating a SessionContext - [Creating a SessionContext](./create-context.py) @@ -62,12 +85,6 @@ type and codec boundaries rather than same-library Rust downcasts. - [Serialize query plans using Substrait](./substrait.py) -### Executing SQL against DataFrame Libraries (Experimental) - -- [Executing SQL on Polars](./sql-on-polars.py) -- [Executing SQL on Pandas](./sql-on-pandas.py) -- [Executing SQL on cuDF](./sql-on-cudf.py) - ## TPC-H Examples Within the subdirectory `tpch` there are 22 examples that reproduce queries in diff --git a/examples/csv-read-options.py b/examples/csv-read-options.py index a5952d950..aa0dc923d 100644 --- a/examples/csv-read-options.py +++ b/examples/csv-read-options.py @@ -15,17 +15,49 @@ # specific language governing permissions and limitations # under the License. -"""Example demonstrating CsvReadOptions usage.""" +"""Example demonstrating CsvReadOptions usage. + +The example writes the small CSV files it reads into a temporary directory, so +it is self-contained and can be run from any working directory. +""" + +import gzip +import tempfile +from pathlib import Path from datafusion import CsvReadOptions, SessionContext +# Write the sample data this example reads into a temporary directory, rather +# than checking the files into the repository and reading them by a relative +# path that only resolves from one working directory. +tmp_dir = tempfile.TemporaryDirectory() +data_dir = Path(tmp_dir.name) + +# Comma separated, quoted with `"`, used by most of the examples below. +csv_file = data_dir / "data.csv" +csv_file.write_text('id,name,value\n1,"alice",10\n2,"bob",20\n3,"carol",30\n') + +# Pipe separated and quoted with `'`, to exercise the builder pattern. +pipe_file = data_dir / "data_pipe.csv" +pipe_file.write_text( + "id|name|value\n1|'alice'|10\n2|'bob, the second'|20\n3|'carol'|30\n" +) + +# Gzipped, with a comment line and an `N/A` placeholder, to exercise the +# advanced options. +gzip_file = data_dir / "data.csv.gz" +gzip_file.write_bytes( + gzip.compress(b"# sample data\nid,name,value\n1,alice,10\n2,N/A,20\n3,carol,30\n") +) + # Create a SessionContext ctx = SessionContext() # Example 1: Using CsvReadOptions with default values print("Example 1: Default CsvReadOptions") options = CsvReadOptions() -df = ctx.read_csv("data.csv", options=options) +df = ctx.read_csv(csv_file, options=options) +df.show() # Example 2: Using CsvReadOptions with custom parameters print("\nExample 2: Custom CsvReadOptions") @@ -36,7 +68,8 @@ schema_infer_max_records=1000, file_extension=".csv", ) -df = ctx.read_csv("data.csv", options=options) +df = ctx.read_csv(csv_file, options=options) +df.show() # Example 3: Using the builder pattern (recommended for readability) print("\nExample 3: Builder pattern") @@ -49,7 +82,8 @@ .with_truncated_rows(False) # noqa: FBT003 .with_newlines_in_values(True) # noqa: FBT003 ) -df = ctx.read_csv("data.csv", options=options) +df = ctx.read_csv(pipe_file, options=options) +df.show() # Example 4: Advanced options print("\nExample 4: Advanced options") @@ -64,18 +98,23 @@ .with_file_compression_type("gzip") # Read gzipped CSV .with_file_extension(".gz") ) -df = ctx.read_csv("data.csv.gz", options=options) +df = ctx.read_csv(gzip_file, options=options) +df.show() # Example 5: Register CSV table with options print("\nExample 5: Register CSV table") options = CsvReadOptions().with_has_header(True).with_delimiter(",") # noqa: FBT003 -ctx.register_csv("my_table", "data.csv", options=options) +ctx.register_csv("my_table", csv_file, options=options) df = ctx.sql("SELECT * FROM my_table") +df.show() # Example 6: Backward compatibility (without options) print("\nExample 6: Backward compatibility") # Still works the old way! -df = ctx.read_csv("data.csv", has_header=True, delimiter=",") +df = ctx.read_csv(csv_file, has_header=True, delimiter=",") +df.show() + +tmp_dir.cleanup() print("\nAll examples completed!") print("\nFor all available options, see the CsvReadOptions documentation:") diff --git a/examples/export.py b/examples/export.py index c7a387bcb..d845fe963 100644 --- a/examples/export.py +++ b/examples/export.py @@ -34,19 +34,29 @@ # export to pandas dataframe pandas_df = df.to_pandas() assert pandas_df.shape == (3, 2) +print("pandas DataFrame:") +print(pandas_df) # export to PyArrow table arrow_table = df.to_arrow_table() assert arrow_table.shape == (3, 2) +print("\nPyArrow table:") +print(arrow_table) # export to Polars dataframe polars_df = df.to_polars() assert polars_df.shape == (3, 2) +print("\nPolars DataFrame:") +print(polars_df) # export to Python list of rows pylist = df.to_pylist() assert pylist == [{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}] +print("\nPython list of rows:") +print(pylist) # export to Python dictionary of columns pydict = df.to_pydict() assert pydict == {"a": [1, 2, 3], "b": [4, 5, 6]} +print("\nPython dictionary of columns:") +print(pydict) diff --git a/examples/import.py b/examples/import.py index 7b5ab5082..4e89a8e94 100644 --- a/examples/import.py +++ b/examples/import.py @@ -28,6 +28,8 @@ # represent column values df = ctx.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}) assert type(df) is datafusion.DataFrame +print("from_pydict:") +df.show() # Dataframe: # +---+---+ # | a | b | @@ -40,18 +42,26 @@ # Create a datafusion DataFrame from a Python list of rows df = ctx.from_pylist([{"a": 1, "b": 4}, {"a": 2, "b": 5}, {"a": 3, "b": 6}]) assert type(df) is datafusion.DataFrame +print("from_pylist:") +df.show() # Convert pandas DataFrame to datafusion DataFrame pandas_df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) df = ctx.from_pandas(pandas_df) assert type(df) is datafusion.DataFrame +print("from_pandas:") +df.show() # Convert polars DataFrame to datafusion DataFrame polars_df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) df = ctx.from_polars(polars_df) assert type(df) is datafusion.DataFrame +print("from_polars:") +df.show() # Convert Arrow Table to datafusion DataFrame arrow_table = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]}) df = ctx.from_arrow(arrow_table) assert type(df) is datafusion.DataFrame +print("from_arrow:") +df.show() diff --git a/examples/python-udaf.py b/examples/python-udaf.py index 6655edb0a..5eccb0a58 100644 --- a/examples/python-udaf.py +++ b/examples/python-udaf.py @@ -64,6 +64,8 @@ def evaluate(self) -> pa.Scalar: df = df.aggregate([], [my_udaf(col("a"))]) +df.show() + result = df.collect()[0] assert result.column(0) == pa.array([6.0]) diff --git a/examples/python-udf.py b/examples/python-udf.py index 1c08acd1a..b1b7f8797 100644 --- a/examples/python-udf.py +++ b/examples/python-udf.py @@ -38,6 +38,8 @@ def is_null(array: pa.Array) -> pa.Array: df = df.select(is_null_arr(f.col("a"))) +df.show() + result = df.collect()[0] assert result.column(0) == pa.array([False] * 3) diff --git a/examples/query-pyarrow-data.py b/examples/query-pyarrow-data.py index 9cfe8a62b..bb5a55336 100644 --- a/examples/query-pyarrow-data.py +++ b/examples/query-pyarrow-data.py @@ -35,6 +35,8 @@ col("a") - col("b"), ) +df.show() + # execute and collect the first (and only) batch result = df.collect()[0] diff --git a/examples/sql-to-pandas.py b/examples/sql-to-pandas.py index 34f7bde1b..a65f939b2 100644 --- a/examples/sql-to-pandas.py +++ b/examples/sql-to-pandas.py @@ -34,9 +34,11 @@ # convert to Pandas pandas_df = df.to_pandas() +print(pandas_df) # create a chart fig = pandas_df.plot( kind="bar", title="Trip Count by Number of Passengers" ).get_figure() fig.savefig("chart.png") +print("wrote chart.png") diff --git a/examples/sql-using-python-udaf.py b/examples/sql-using-python-udaf.py index f42bbdc23..b15d812ba 100644 --- a/examples/sql-using-python-udaf.py +++ b/examples/sql-using-python-udaf.py @@ -75,6 +75,7 @@ def evaluate(self) -> pa.Scalar: result_df = ctx.sql( "select a, my_accumulator(b) as b_aggregated from t group by a order by a" ) +result_df.show() # Dataframe: # +---+--------------+ # | a | b_aggregated | diff --git a/examples/sql-using-python-udf.py b/examples/sql-using-python-udf.py index 2f0a0b67d..4edb653c7 100644 --- a/examples/sql-using-python-udf.py +++ b/examples/sql-using-python-udf.py @@ -53,6 +53,7 @@ def is_null(array: pa.Array) -> pa.Array: # Query the DataFrame using SQL result_df = ctx.sql("select a, is_null(b) as b_is_null from t") +result_df.show() # Dataframe: # +---+-----------+ # | a | b_is_null | diff --git a/examples/substrait.py b/examples/substrait.py index fa6f77912..2a525f3a8 100644 --- a/examples/substrait.py +++ b/examples/substrait.py @@ -15,14 +15,21 @@ # specific language governing permissions and limitations # under the License. +from pathlib import Path + from datafusion import SessionContext from datafusion import substrait as ss +# Resolve the data relative to this file so the example can be run from any +# working directory. The file comes from the `testing` git submodule: +# `git submodule update --init testing`. +csv_path = Path(__file__).parent.parent / "testing/data/csv/aggregate_test_100.csv" + # Create a DataFusion context ctx = SessionContext() # Register table with context -ctx.register_csv("aggregate_test_data", "./testing/data/csv/aggregate_test_100.csv") +ctx.register_csv("aggregate_test_data", csv_path) substrait_plan = ss.Serde.serialize_to_plan("SELECT * FROM aggregate_test_data", ctx) # type(substrait_plan) -> @@ -31,6 +38,7 @@ substrait_bytes = substrait_plan.encode() # type(substrait_bytes) -> , at this point the bytes can be distributed to file, network, etc safely # where they could subsequently be deserialized on the receiving end. +print(f"Encoded Substrait plan: {len(substrait_bytes)} bytes") # Alternative serialization approaches # type(substrait_bytes) -> , at this point the bytes can be distributed to file, network, etc safely @@ -44,6 +52,10 @@ # type(df_logical_plan) -> df_logical_plan = ss.Consumer.from_substrait_plan(ctx, substrait_plan) +print("\nLogical plan round-tripped from Substrait:") +print(df_logical_plan.display_indent()) + # Back to Substrait Plan just for demonstration purposes # type(substrait_plan) -> substrait_plan = ss.Producer.to_substrait_plan(df_logical_plan, ctx) +print(f"\nRe-encoded Substrait plan: {len(substrait_plan.encode())} bytes")