Skip to content

Commit

Permalink
simplify test
Browse files Browse the repository at this point in the history
  • Loading branch information
nikil-ravi committed Oct 14, 2024
1 parent 50715e9 commit 3fe8995
Showing 1 changed file with 1 addition and 6 deletions.
7 changes: 1 addition & 6 deletions tests/test_sharded_dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,14 +48,12 @@ def test_basic_parquet_datasource_read_row():
import pyarrow as pa
import pyarrow.parquet as pq

with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as f:
with tempfile.NamedTemporaryFile(suffix=".parquet", delete=True) as f:
# Create a simple dataset
data = {"column1": ["value1", "value2", "value3"], "column2": [10, 20, 30]}
table = pa.Table.from_pydict(data)
pq.write_table(table, f.name)

try:

datasource = ParquetDataSource([os.path.abspath(f.name)])

assert len(datasource.shard_names) == 1, "Expected only one shard"
Expand All @@ -70,6 +68,3 @@ def test_basic_parquet_datasource_read_row():
assert row_data[0]["column2"] == 20
assert row_data[1]["column1"] == "value3"
assert row_data[1]["column2"] == 30

finally:
os.unlink(f.name)

0 comments on commit 3fe8995

Please sign in to comment.