Skip to content

Commit f1e9c37

Browse files
dedsec-terminaldedsec-terminal
committed
fix: make csv-read-options.py example self-contained with temporary fixtures (#1728)
1 parent 7022baa commit f1e9c37

1 file changed

Lines changed: 89 additions & 68 deletions

File tree

‎examples/csv-read-options.py‎

Lines changed: 89 additions & 68 deletions
Original file line numberDiff line numberDiff line change
@@ -17,80 +17,101 @@
1717

1818
"""Example demonstrating CsvReadOptions usage."""
1919

20+
import gzip
21+
import tempfile
22+
from pathlib import Path
23+
2024
from datafusion import CsvReadOptions, SessionContext
2125

22-
# Create a SessionContext
23-
ctx = SessionContext()
26+
# Create sample data files in a temporary directory to make the example self-contained
27+
with tempfile.TemporaryDirectory() as tmp_dir:
28+
tmp_path = Path(tmp_dir)
29+
csv_file = tmp_path / "data.csv"
30+
gz_file = tmp_path / "data.csv.gz"
31+
32+
sample_csv = "id,name,value\n1,alice,100\n2,bob,200\n3,charlie,null\n"
33+
csv_file.write_text(sample_csv)
34+
35+
with gzip.open(gz_file, "wt") as f:
36+
f.write(sample_csv)
37+
38+
# Create a SessionContext
39+
ctx = SessionContext()
2440

25-
# Example 1: Using CsvReadOptions with default values
26-
print("Example 1: Default CsvReadOptions")
27-
options = CsvReadOptions()
28-
df = ctx.read_csv("data.csv", options=options)
41+
# Example 1: Using CsvReadOptions with default values
42+
print("Example 1: Default CsvReadOptions")
43+
options = CsvReadOptions()
44+
df = ctx.read_csv(str(csv_file), options=options)
45+
df.show()
2946

30-
# Example 2: Using CsvReadOptions with custom parameters
31-
print("\nExample 2: Custom CsvReadOptions")
32-
options = CsvReadOptions(
33-
has_header=True,
34-
delimiter=",",
35-
quote='"',
36-
schema_infer_max_records=1000,
37-
file_extension=".csv",
38-
)
39-
df = ctx.read_csv("data.csv", options=options)
47+
# Example 2: Using CsvReadOptions with custom parameters
48+
print("\nExample 2: Custom CsvReadOptions")
49+
options = CsvReadOptions(
50+
has_header=True,
51+
delimiter=",",
52+
quote='"',
53+
schema_infer_max_records=1000,
54+
file_extension=".csv",
55+
)
56+
df = ctx.read_csv(str(csv_file), options=options)
57+
df.show()
4058

41-
# Example 3: Using the builder pattern (recommended for readability)
42-
print("\nExample 3: Builder pattern")
43-
options = (
44-
CsvReadOptions()
45-
.with_has_header(True) # noqa: FBT003
46-
.with_delimiter("|")
47-
.with_quote("'")
48-
.with_schema_infer_max_records(500)
49-
.with_truncated_rows(False) # noqa: FBT003
50-
.with_newlines_in_values(True) # noqa: FBT003
51-
)
52-
df = ctx.read_csv("data.csv", options=options)
59+
# Example 3: Using the builder pattern (recommended for readability)
60+
print("\nExample 3: Builder pattern")
61+
options = (
62+
CsvReadOptions()
63+
.with_has_header(True) # noqa: FBT003
64+
.with_delimiter("|")
65+
.with_quote("'")
66+
.with_schema_infer_max_records(500)
67+
.with_truncated_rows(False) # noqa: FBT003
68+
.with_newlines_in_values(True) # noqa: FBT003
69+
)
70+
df = ctx.read_csv(str(csv_file), options=options)
5371

54-
# Example 4: Advanced options
55-
print("\nExample 4: Advanced options")
56-
options = (
57-
CsvReadOptions()
58-
.with_has_header(True) # noqa: FBT003
59-
.with_delimiter(",")
60-
.with_comment("#") # Skip lines starting with #
61-
.with_escape("\\") # Escape character
62-
.with_null_regex(r"^(null|NULL|N/A)$") # Treat these as NULL
63-
.with_truncated_rows(True) # noqa: FBT003
64-
.with_file_compression_type("gzip") # Read gzipped CSV
65-
.with_file_extension(".gz")
66-
)
67-
df = ctx.read_csv("data.csv.gz", options=options)
72+
# Example 4: Advanced options
73+
print("\nExample 4: Advanced options")
74+
options = (
75+
CsvReadOptions()
76+
.with_has_header(True) # noqa: FBT003
77+
.with_delimiter(",")
78+
.with_comment("#") # Skip lines starting with #
79+
.with_escape("\\") # Escape character
80+
.with_null_regex(r"^(null|NULL|N/A)$") # Treat these as NULL
81+
.with_truncated_rows(True) # noqa: FBT003
82+
.with_file_compression_type("gzip") # Read gzipped CSV
83+
.with_file_extension(".gz")
84+
)
85+
df = ctx.read_csv(str(gz_file), options=options)
86+
df.show()
6887

69-
# Example 5: Register CSV table with options
70-
print("\nExample 5: Register CSV table")
71-
options = CsvReadOptions().with_has_header(True).with_delimiter(",") # noqa: FBT003
72-
ctx.register_csv("my_table", "data.csv", options=options)
73-
df = ctx.sql("SELECT * FROM my_table")
88+
# Example 5: Register CSV table with options
89+
print("\nExample 5: Register CSV table")
90+
options = CsvReadOptions().with_has_header(True).with_delimiter(",") # noqa: FBT003
91+
ctx.register_csv("my_table", str(csv_file), options=options)
92+
df = ctx.sql("SELECT * FROM my_table")
93+
df.show()
7494

75-
# Example 6: Backward compatibility (without options)
76-
print("\nExample 6: Backward compatibility")
77-
# Still works the old way!
78-
df = ctx.read_csv("data.csv", has_header=True, delimiter=",")
95+
# Example 6: Backward compatibility (without options)
96+
print("\nExample 6: Backward compatibility")
97+
# Still works the old way!
98+
df = ctx.read_csv(str(csv_file), has_header=True, delimiter=",")
99+
df.show()
79100

80-
print("\nAll examples completed!")
81-
print("\nFor all available options, see the CsvReadOptions documentation:")
82-
print(" - has_header: bool")
83-
print(" - delimiter: str")
84-
print(" - quote: str")
85-
print(" - terminator: str | None")
86-
print(" - escape: str | None")
87-
print(" - comment: str | None")
88-
print(" - newlines_in_values: bool")
89-
print(" - schema: pa.Schema | None")
90-
print(" - schema_infer_max_records: int")
91-
print(" - file_extension: str")
92-
print(" - table_partition_cols: list[tuple[str, pa.DataType]]")
93-
print(" - file_compression_type: str")
94-
print(" - file_sort_order: list[list[SortExpr]]")
95-
print(" - null_regex: str | None")
96-
print(" - truncated_rows: bool")
101+
print("\nAll examples completed!")
102+
print("\nFor all available options, see the CsvReadOptions documentation:")
103+
print(" - has_header: bool")
104+
print(" - delimiter: str")
105+
print(" - quote: str")
106+
print(" - terminator: str | None")
107+
print(" - escape: str | None")
108+
print(" - comment: str | None")
109+
print(" - newlines_in_values: bool")
110+
print(" - schema: pa.Schema | None")
111+
print(" - schema_infer_max_records: int")
112+
print(" - file_extension: str")
113+
print(" - table_partition_cols: list[tuple[str, pa.DataType]]")
114+
print(" - file_compression_type: str")
115+
print(" - file_sort_order: list[list[SortExpr]]")
116+
print(" - null_regex: str | None")
117+
print(" - truncated_rows: bool")

0 commit comments

Comments
 (0)