|
15 | 15 | # specific language governing permissions and limitations |
16 | 16 | # under the License. |
17 | 17 |
|
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 |
19 | 27 |
|
20 | 28 | from datafusion import CsvReadOptions, SessionContext |
21 | 29 |
|
| 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 | + |
22 | 53 | # Create a SessionContext |
23 | 54 | ctx = SessionContext() |
24 | 55 |
|
25 | 56 | # Example 1: Using CsvReadOptions with default values |
26 | 57 | print("Example 1: Default CsvReadOptions") |
27 | 58 | options = CsvReadOptions() |
28 | | -df = ctx.read_csv("data.csv", options=options) |
| 59 | +df = ctx.read_csv(csv_file, options=options) |
| 60 | +df.show() |
29 | 61 |
|
30 | 62 | # Example 2: Using CsvReadOptions with custom parameters |
31 | 63 | print("\nExample 2: Custom CsvReadOptions") |
|
36 | 68 | schema_infer_max_records=1000, |
37 | 69 | file_extension=".csv", |
38 | 70 | ) |
39 | | -df = ctx.read_csv("data.csv", options=options) |
| 71 | +df = ctx.read_csv(csv_file, options=options) |
| 72 | +df.show() |
40 | 73 |
|
41 | 74 | # Example 3: Using the builder pattern (recommended for readability) |
42 | 75 | print("\nExample 3: Builder pattern") |
|
49 | 82 | .with_truncated_rows(False) # noqa: FBT003 |
50 | 83 | .with_newlines_in_values(True) # noqa: FBT003 |
51 | 84 | ) |
52 | | -df = ctx.read_csv("data.csv", options=options) |
| 85 | +df = ctx.read_csv(pipe_file, options=options) |
| 86 | +df.show() |
53 | 87 |
|
54 | 88 | # Example 4: Advanced options |
55 | 89 | print("\nExample 4: Advanced options") |
|
64 | 98 | .with_file_compression_type("gzip") # Read gzipped CSV |
65 | 99 | .with_file_extension(".gz") |
66 | 100 | ) |
67 | | -df = ctx.read_csv("data.csv.gz", options=options) |
| 101 | +df = ctx.read_csv(gzip_file, options=options) |
| 102 | +df.show() |
68 | 103 |
|
69 | 104 | # Example 5: Register CSV table with options |
70 | 105 | print("\nExample 5: Register CSV table") |
71 | 106 | 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) |
73 | 108 | df = ctx.sql("SELECT * FROM my_table") |
| 109 | +df.show() |
74 | 110 |
|
75 | 111 | # Example 6: Backward compatibility (without options) |
76 | 112 | print("\nExample 6: Backward compatibility") |
77 | 113 | # 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() |
79 | 118 |
|
80 | 119 | print("\nAll examples completed!") |
81 | 120 | print("\nFor all available options, see the CsvReadOptions documentation:") |
|
0 commit comments