feat: run examples/*.py in CI and fix the examples that could not run

.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

examples/README.md

@@ -19,14 +19,37 @@
 
 # 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
 
 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

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

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)

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

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

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)

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]
 

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

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 |

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 |

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) -> <class 'datafusion.substrait.plan'>
@@ -31,6 +38,7 @@
 substrait_bytes = substrait_plan.encode()
 # type(substrait_bytes) -> <class '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) -> <class 'bytes'>, at this point the bytes can be distributed to file, network, etc safely
@@ -44,6 +52,10 @@
 # type(df_logical_plan) -> <class 'substrait.LogicalPlan'>
 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) -> <class 'datafusion.substrait.plan'>
 substrait_plan = ss.Producer.to_substrait_plan(df_logical_plan, ctx)
+print(f"\nRe-encoded Substrait plan: {len(substrait_plan.encode())} bytes")