| 68 | ) |
| 69 | @pytest.mark.parametrize("use_head", [True, False]) |
| 70 | def test_single_column(db, use_head): |
| 71 | from sqlalchemy import Column, Integer, MetaData, Table, create_engine |
| 72 | |
| 73 | with tmpfile() as f: |
| 74 | uri = "sqlite:///%s" % f |
| 75 | metadata = MetaData() |
| 76 | engine = create_engine(uri) |
| 77 | table = Table( |
| 78 | "single_column", |
| 79 | metadata, |
| 80 | Column("id", Integer, primary_key=True), |
| 81 | ) |
| 82 | metadata.create_all(engine) |
| 83 | test_data = pd.DataFrame({"id": list(range(50))}).set_index("id") |
| 84 | test_data.to_sql(table.name, uri, index=True, if_exists="replace") |
| 85 | |
| 86 | if use_head: |
| 87 | dask_df = read_sql_table(table.name, uri, index_col="id", npartitions=2) |
| 88 | else: |
| 89 | dask_df = read_sql_table( |
| 90 | table.name, |
| 91 | uri, |
| 92 | head_rows=0, |
| 93 | npartitions=2, |
| 94 | meta=test_data.iloc[:0], |
| 95 | index_col="id", |
| 96 | ) |
| 97 | assert dask_df.index.name == "id" |
| 98 | assert dask_df.npartitions == 2 |
| 99 | pd_dataframe = dask_df.compute() |
| 100 | assert_eq(test_data, pd_dataframe) |
| 101 | |
| 102 | |
| 103 | def test_passing_engine_as_uri_raises_helpful_error(db): |