10
10
from pathlib import Path
11
11
import pyarrow as pa
12
12
import os
13
- import sqlite3
14
13
import sys
15
14
16
15
# needed to resolve ray-project/ray#3744
@@ -47,7 +46,6 @@ def setup_parquet_file(row_size, force=False):
47
46
).to_parquet (TEST_PARQUET_FILENAME )
48
47
49
48
50
- @pytest .fixture
51
49
def create_test_ray_dataframe ():
52
50
df = pd .DataFrame (
53
51
{
@@ -62,7 +60,6 @@ def create_test_ray_dataframe():
62
60
return df
63
61
64
62
65
- @pytest .fixture
66
63
def create_test_pandas_dataframe ():
67
64
df = pandas .DataFrame (
68
65
{
@@ -261,26 +258,41 @@ def teardown_pickle_file():
261
258
262
259
263
260
@pytest .fixture
264
- def setup_sql_file (conn , filename , table , force = False ):
265
- if os .path .exists (filename ) and not force :
266
- pass
267
- else :
268
- df = pandas .DataFrame (
269
- {
270
- "col1" : [0 , 1 , 2 , 3 , 4 , 5 , 6 ],
271
- "col2" : [7 , 8 , 9 , 10 , 11 , 12 , 13 ],
272
- "col3" : [14 , 15 , 16 , 17 , 18 , 19 , 20 ],
273
- "col4" : [21 , 22 , 23 , 24 , 25 , 26 , 27 ],
274
- "col5" : [0 , 0 , 0 , 0 , 0 , 0 , 0 ],
275
- }
276
- )
277
- df .to_sql (table , conn )
278
-
279
-
280
- @pytest .fixture
281
- def teardown_sql_file (filename ):
282
- if os .path .exists (filename ):
283
- os .remove (filename )
261
+ def make_sql_connection ():
262
+ """Sets up sql connections and takes them down after the caller is done.
263
+
264
+ Yields:
265
+ Factory that generates sql connection objects
266
+ """
267
+ filenames = []
268
+
269
+ def _sql_connection (filename , table = "" ):
270
+ # Remove file if exists
271
+ if os .path .exists (filename ):
272
+ os .remove (filename )
273
+ filenames .append (filename )
274
+
275
+ # Create connection and, if needed, table
276
+ conn = "sqlite:///{}" .format (filename )
277
+ if table :
278
+ df = pandas .DataFrame (
279
+ {
280
+ "col1" : [0 , 1 , 2 , 3 , 4 , 5 , 6 ],
281
+ "col2" : [7 , 8 , 9 , 10 , 11 , 12 , 13 ],
282
+ "col3" : [14 , 15 , 16 , 17 , 18 , 19 , 20 ],
283
+ "col4" : [21 , 22 , 23 , 24 , 25 , 26 , 27 ],
284
+ "col5" : [0 , 0 , 0 , 0 , 0 , 0 , 0 ],
285
+ }
286
+ )
287
+ df .to_sql (table , conn )
288
+ return conn
289
+
290
+ yield _sql_connection
291
+
292
+ # Takedown the fixture
293
+ for filename in filenames :
294
+ if os .path .exists (filename ):
295
+ os .remove (filename )
284
296
285
297
286
298
def test_from_parquet ():
@@ -460,21 +472,17 @@ def test_from_pickle():
460
472
teardown_pickle_file ()
461
473
462
474
463
- def test_from_sql ():
475
+ def test_from_sql (make_sql_connection ):
464
476
filename = "test_from_sql.db"
465
- teardown_sql_file (filename )
466
- conn = sqlite3 .connect (filename )
467
477
table = "test_from_sql"
468
- setup_sql_file ( conn , filename , table , True )
478
+ conn = make_sql_connection ( filename , table )
469
479
query = "select * from {0}" .format (table )
470
480
471
481
pandas_df = pandas .read_sql (query , conn )
472
482
modin_df = pd .read_sql (query , conn )
473
483
474
484
assert modin_df_equals_pandas (modin_df , pandas_df )
475
485
476
- teardown_sql_file (filename )
477
-
478
486
479
487
@pytest .mark .skip (reason = "No SAS write methods in Pandas" )
480
488
def test_from_sas ():
@@ -750,20 +758,40 @@ def test_to_pickle():
750
758
teardown_test_file (TEST_PICKLE_DF_FILENAME )
751
759
752
760
753
- def test_to_sql ():
761
+ def test_to_sql_without_index (make_sql_connection ):
762
+ table_name = "tbl_without_index"
754
763
modin_df = create_test_ray_dataframe ()
755
764
pandas_df = create_test_pandas_dataframe ()
756
765
757
- TEST_SQL_DF_FILENAME = "test_df.sql"
758
- TEST_SQL_pandas_FILENAME = "test_pandas.sql"
766
+ # We do not pass the table name so the fixture won't generate a table
767
+ conn = make_sql_connection ("test_to_sql.db" )
768
+ modin_df .to_sql (table_name , conn , index = False )
769
+ df_modin_sql = pandas .read_sql (table_name , con = conn )
770
+
771
+ # We do not pass the table name so the fixture won't generate a table
772
+ conn = make_sql_connection ("test_to_sql_pandas.db" )
773
+ pandas_df .to_sql (table_name , conn , index = False )
774
+ df_pandas_sql = pandas .read_sql (table_name , con = conn )
775
+
776
+ assert df_modin_sql .sort_index ().equals (df_pandas_sql .sort_index ())
777
+
778
+
779
+ def test_to_sql_with_index (make_sql_connection ):
780
+ table_name = "tbl_with_index"
781
+ modin_df = create_test_ray_dataframe ()
782
+ pandas_df = create_test_pandas_dataframe ()
759
783
760
- modin_df .to_pickle (TEST_SQL_DF_FILENAME )
761
- pandas_df .to_pickle (TEST_SQL_pandas_FILENAME )
784
+ # We do not pass the table name so the fixture won't generate a table
785
+ conn = make_sql_connection ("test_to_sql.db" )
786
+ modin_df .to_sql (table_name , conn )
787
+ df_modin_sql = pandas .read_sql (table_name , con = conn , index_col = "index" )
762
788
763
- assert test_files_eq (TEST_SQL_DF_FILENAME , TEST_SQL_pandas_FILENAME )
789
+ # We do not pass the table name so the fixture won't generate a table
790
+ conn = make_sql_connection ("test_to_sql_pandas.db" )
791
+ pandas_df .to_sql (table_name , conn )
792
+ df_pandas_sql = pandas .read_sql (table_name , con = conn , index_col = "index" )
764
793
765
- teardown_test_file (TEST_SQL_DF_FILENAME )
766
- teardown_test_file (TEST_SQL_pandas_FILENAME )
794
+ assert df_modin_sql .sort_index ().equals (df_pandas_sql .sort_index ())
767
795
768
796
769
797
def test_to_stata ():
0 commit comments