11import math
22import string
33from datetime import datetime , timedelta
4-
5- try :
6- from functools import lru_cache
7- except ImportError :
8- from functools32 import lru_cache
9-
4+ from functools import lru_cache
105from itertools import cycle
116
127import numpy as np
@@ -105,7 +100,7 @@ def get_df_complex_index():
105100 return df
106101
107102
108- def get_dict_of_test_dfs (N = 100 , M = 100 , polars = False , ibis = False ):
103+ def get_dict_of_test_dfs (N = 100 , M = 100 , type = "pandas" ):
109104 NM_values = np .reshape (np .linspace (start = 0.0 , stop = 1.0 , num = N * M ), (N , M ))
110105
111106 test_dfs = {
@@ -266,8 +261,10 @@ def get_dict_of_test_dfs(N=100, M=100, polars=False, ibis=False):
266261 }
267262 ),
268263 }
264+ if type == "pandas" :
265+ return test_dfs
269266
270- if polars :
267+ if type == " polars" :
271268 import polars as pl
272269 import pyarrow as pa
273270
@@ -279,23 +276,42 @@ def get_dict_of_test_dfs(N=100, M=100, polars=False, ibis=False):
279276 pass
280277 return polars_dfs
281278
282- if ibis :
283- import ibis as ib
279+ if type == "ibis_memtable" :
280+ import ibis
284281
285- con = ib .pandas .connect (test_dfs )
286282 ibis_dfs = {}
287- for key in test_dfs :
283+ for key , df in test_dfs .items ():
284+ # Ibis does not support tables with no columns
285+ if not len (df .columns ):
286+ continue
287+ try :
288+ ibis_dfs [key ] = ibis .memtable (df , name = key )
289+ except (TypeError , ibis .common .exceptions .IbisInputError ):
290+ pass
291+
292+ return ibis_dfs
293+
294+ if type == "ibis_connect" :
295+ import ibis
296+
297+ con = ibis .pandas .connect (test_dfs )
298+ ibis_dfs = {}
299+ for key , df in test_dfs .items ():
300+ # Ibis does not support tables with no columns
301+ if not len (df .columns ):
302+ continue
303+
288304 try :
289- ibis_dfs [key ] = con .table (key )
305+ ibis_dfs [f" { key } _connect" ] = con .table (key )
290306 except (TypeError , AttributeError ):
291307 pass
292308
293309 return ibis_dfs
294310
295- return test_dfs
311+ raise NotImplementedError ( type )
296312
297313
298- def get_dict_of_test_series (polars = False ):
314+ def get_dict_of_test_series (type = "pandas" ):
299315 series = {}
300316 for df_name , df in get_dict_of_test_dfs ().items ():
301317 if len (df .columns ) > 6 :
@@ -306,7 +322,10 @@ def get_dict_of_test_series(polars=False):
306322 continue
307323 series ["{}.{}" .format (df_name , col )] = df [col ]
308324
309- if polars :
325+ if type == "pandas" :
326+ return series
327+
328+ if type == "polars" :
310329 import polars as pl
311330 import pyarrow as pa
312331
@@ -325,7 +344,7 @@ def get_dict_of_test_series(polars=False):
325344
326345 return polars_series
327346
328- return series
347+ raise NotImplementedError ( type )
329348
330349
331350@lru_cache ()
0 commit comments