| 134 | |
| 135 | |
| 136 | def test_write_sql(temp_database: str): |
| 137 | connection = sqlite3.connect(temp_database) |
| 138 | connection.cursor().execute("CREATE TABLE test(string, number)") |
| 139 | dataset = ray.data.from_items( |
| 140 | [{"string": "spam", "number": 0}, {"string": "ham", "number": 1}] |
| 141 | ) |
| 142 | |
| 143 | dataset.write_sql( |
| 144 | "INSERT INTO test VALUES(?, ?)", lambda: sqlite3.connect(temp_database) |
| 145 | ) |
| 146 | |
| 147 | result = connection.cursor().execute("SELECT * FROM test ORDER BY number") |
| 148 | assert result.fetchall() == [("spam", 0), ("ham", 1)] |
| 149 | |
| 150 | |
| 151 | @pytest.mark.parametrize("num_blocks", (1, 20)) |