sqlutilpy 0.26.0__tar.gz → 0.27.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlutilpy
3
- Version: 0.26.0
3
+ Version: 0.27.0
4
4
  Summary: Database query code returning numpy arrays
5
5
  Author-email: "Sergey E. Koposov" <skoposov@ed.ac.uk>
6
6
  License: BSD
@@ -1,12 +1,11 @@
1
1
  """Sqlutilpy module to access SQL databases
2
2
  """
3
- from __future__ import print_function
4
3
  import numpy
5
4
  import numpy as np
6
5
  import psycopg
7
- import threading
8
6
  import collections
9
7
  import warnings
8
+ import concurrent.futures
10
9
 
11
10
  try:
12
11
  import astropy.table as atpy
@@ -19,8 +18,6 @@ except ImportError:
19
18
  # pandas is not installed
20
19
  pandas = None
21
20
 
22
- import queue
23
-
24
21
  _WAIT_SELECT_TIMEOUT = 10
25
22
  STRLEN_DEFAULT = 20
26
23
 
@@ -99,7 +96,7 @@ def getConnection(db=None,
99
96
  return conn
100
97
 
101
98
 
102
- def getCursor(conn, driver=None, preamb=None, notNamed=False):
99
+ def getCursor(conn, driver=None, preamb=None, named=True):
103
100
  """
104
101
  Retrieve the database cursor
105
102
  """
@@ -108,14 +105,14 @@ def getCursor(conn, driver=None, preamb=None, notNamed=False):
108
105
  'psycopg2 driver is not supported anymore using psycopg instead')
109
106
  driver = 'psycopg'
110
107
  if driver == 'psycopg':
111
- cur = conn.cursor()
108
+ cur = conn.cursor(scrollable=False)
112
109
  if preamb is not None:
113
110
  cur.execute(preamb)
114
111
  else:
115
112
  cur.execute('set cursor_tuple_fraction TO 1')
116
113
  # this is required because otherwise PG may decide to execute a
117
114
  # different plan
118
- if notNamed:
115
+ if not named:
119
116
  return cur
120
117
  cur = conn.cursor(name='sqlutilcursor')
121
118
  cur.arraysize = config.arraysize
@@ -132,65 +129,75 @@ def getCursor(conn, driver=None, preamb=None, notNamed=False):
132
129
  return cur
133
130
 
134
131
 
135
- def __fromrecords(recList, dtype=None, intNullVal=None):
132
+ def __fromrecords(recList, dtype=None, intNullVal=None, strNullVal='None'):
136
133
  """
137
134
  This function was taken from np.core.records and updated to
138
- support conversion null integers to intNullVal
135
+ support conversion null integers to intNullVal and strings to strNullVal
139
136
  """
140
137
 
141
- shape = None
142
138
  descr = np.dtype((np.record, dtype))
143
- try:
144
- retval = np.array(recList, dtype=descr)
145
- except TypeError: # list of lists instead of list of tuples
146
- shape = (len(recList), )
147
- _array = np.recarray(shape, descr)
139
+
140
+ # Check if we must avoid fast path for strings to prevent
141
+ # None -> 'None' conversion
142
+ # Note: numpy's default conversion for None to string is 'None', so if
143
+ # strNullVal is 'None', we can still use the fast path.
144
+ force_object_strings = False
145
+ if strNullVal is not None and strNullVal != 'None':
146
+ if any(dtype.fields[n][0].kind in ('S', 'U') for n in dtype.names):
147
+ force_object_strings = True
148
+
149
+ if not force_object_strings:
148
150
  try:
149
- for k in range(_array.size):
150
- _array[k] = tuple(recList[k])
151
- except TypeError:
152
- convs = []
153
- ncols = len(dtype.fields)
154
- for _k in dtype.names:
155
- _v = dtype.fields[_k]
156
- if _v[0] in [np.int16, np.int32, np.int64]:
157
- convs.append(lambda x: intNullVal if x is None else x)
158
- else:
159
- convs.append(lambda x: x)
160
- convs = tuple(convs)
161
-
162
- def convF(x):
163
- return [convs[_](x[_]) for _ in range(ncols)]
164
-
165
- for k in range(k, _array.size):
166
- try:
167
- _array[k] = tuple(recList[k])
168
- except TypeError:
169
- _array[k] = tuple(convF(recList[k]))
170
- return _array
171
- else:
172
- if shape is not None and retval.shape != shape:
173
- retval.shape = shape
151
+ retval = np.array(recList, dtype=descr)
152
+ return retval.view(numpy.recarray)
153
+ except (TypeError, ValueError):
154
+ # Failed (likely due to None in int col). Fall through to object
155
+ # path.
156
+ pass
174
157
 
175
- res = retval.view(numpy.recarray)
158
+ # Vectorized fallback path using object arrays
159
+ names = dtype.names
160
+ new_formats = []
161
+ converters = {}
176
162
 
177
- return res
163
+ for name in names:
164
+ dt = dtype.fields[name][0]
165
+ if dt.kind in ('i', 'u'):
166
+ new_formats.append(object)
167
+ converters[name] = intNullVal
168
+ elif dt.kind in ('S', 'U') and strNullVal is not None:
169
+ new_formats.append(object)
170
+ converters[name] = strNullVal
171
+ else:
172
+ new_formats.append(dt)
178
173
 
174
+ temp_dtype = np.dtype({'names': names, 'formats': new_formats})
179
175
 
180
- def __converter(qIn, qOut, endEvent, dtype, intNullVal):
181
- """ Convert the input stream of tuples into numpy arrays """
182
- while (not endEvent.is_set()):
183
- try:
184
- tups = qIn.get(True, 0.1)
185
- except queue.Empty:
186
- continue
187
- try:
188
- res = __fromrecords(tups, dtype=dtype, intNullVal=intNullVal)
189
- except Exception:
190
- print('Failed to convert input data into array')
191
- endEvent.set()
192
- raise
193
- qOut.put(res)
176
+ if len(recList) == 0:
177
+ return np.recarray((0, ), dtype=descr)
178
+
179
+ try:
180
+ arr = np.array(recList, dtype=temp_dtype)
181
+ except (TypeError, ValueError):
182
+ # Handle cases where np.array fails (e.g. ragged lists)
183
+ # by manually filling the object array
184
+ shape = (len(recList), )
185
+ arr = np.recarray(shape, temp_dtype)
186
+ for k in range(len(recList)):
187
+ arr[k] = tuple(recList[k])
188
+
189
+ # Apply conversions
190
+ for name, val in converters.items():
191
+ col = arr[name]
192
+ # Check for None efficiently
193
+ mask = (col == None)
194
+ if np.any(mask):
195
+ if val is not None:
196
+ col[mask] = val
197
+
198
+ # Cast to final dtype
199
+ retval = arr.astype(descr)
200
+ return retval.view(numpy.recarray)
194
201
 
195
202
 
196
203
  def __getDType(row, typeCodes, strLength):
@@ -247,6 +254,112 @@ def __getDType(row, typeCodes, strLength):
247
254
  return dtype
248
255
 
249
256
 
257
+ def __get_representative_row(batch, type_codes):
258
+ """
259
+ Scans the batch to find the maximum length of string columns and returns
260
+ a representative row (based on the first row) with those maximum lengths.
261
+ This ensures that __getDType determines the correct string length for
262
+ the array.
263
+ """
264
+ first_row = list(batch[0])
265
+ strTypes = [25, 1042, 1043]
266
+
267
+ for i, type_code in enumerate(type_codes):
268
+ if type_code in strTypes:
269
+ max_len = 0
270
+ found = False
271
+ for row in batch:
272
+ val = row[i]
273
+ if val is not None:
274
+ max_len = max(max_len, len(val))
275
+ found = True
276
+ if found:
277
+ # Update representative row with max length string
278
+ # __getDType uses max(strLength, len(val))
279
+ if first_row[i] is None or len(first_row[i]) < max_len:
280
+ first_row[i] = 'x' * max_len
281
+ return first_row
282
+
283
+
284
+ def __fix_numpy_record_array(res, intNullVal, strNullVal):
285
+ if res.size == 0:
286
+ return res
287
+
288
+ names = res.dtype.names
289
+ new_arrays = []
290
+
291
+ for name in names:
292
+ col = res[name]
293
+ dtype = col.dtype
294
+
295
+ if dtype.kind == 'O':
296
+ # Check for None (vectorized)
297
+ # We treat everything as candidates if it's Object type
298
+
299
+ # Find first non-None value to guess type
300
+ valid_mask = (col != None)
301
+ if not np.any(valid_mask):
302
+ # All None.
303
+ # We can't determine type. Return as is (Object filled with
304
+ # None)
305
+ new_arrays.append(col)
306
+ continue
307
+
308
+ # Take a sample (first valid)
309
+ first_val = col[valid_mask][0]
310
+
311
+ # INTEGER
312
+ if isinstance(
313
+ first_val,
314
+ (int, np.integer)) and not isinstance(first_val,
315
+ (bool, np.bool_)):
316
+ # Verify all valid are integers
317
+ valid_vals = col[valid_mask]
318
+ is_valid = all(
319
+ isinstance(x, (int, np.integer))
320
+ and not isinstance(x, (bool, np.bool_))
321
+ for x in valid_vals)
322
+ if is_valid:
323
+ # Convert
324
+ temp_col = np.copy(col)
325
+ temp_col[~valid_mask] = intNullVal
326
+ try:
327
+ new_col = temp_col.astype(int)
328
+ new_arrays.append(new_col)
329
+ continue
330
+ except Exception:
331
+ pass
332
+
333
+ # FLOAT
334
+ if isinstance(first_val, (float, np.floating)):
335
+ if all(
336
+ isinstance(x, (float, np.floating, int, np.integer))
337
+ for x in col[valid_mask]):
338
+ temp_col = np.copy(col)
339
+ temp_col[~valid_mask] = np.nan
340
+ try:
341
+ new_col = temp_col.astype(float)
342
+ new_arrays.append(new_col)
343
+ continue
344
+ except Exception:
345
+ pass
346
+
347
+ # STRING
348
+ if isinstance(first_val, str):
349
+ if all(isinstance(x, str) for x in col[valid_mask]):
350
+ # Replicate PG behavior: None -> strNullVal
351
+ temp_col = np.copy(col)
352
+ temp_col[~valid_mask] = strNullVal
353
+ # Let numpy infer string length
354
+ new_col = np.array(list(temp_col))
355
+ new_arrays.append(new_col)
356
+ continue
357
+
358
+ new_arrays.append(col)
359
+
360
+ return numpy.rec.fromarrays(new_arrays, names=names)
361
+
362
+
250
363
  def get(query,
251
364
  params=None,
252
365
  db="wsdb",
@@ -259,9 +372,11 @@ def get(query,
259
372
  port=None,
260
373
  strLength=STRLEN_DEFAULT,
261
374
  timeout=None,
262
- notNamed=False,
375
+ batched=True,
263
376
  asDict=False,
264
- intNullVal=-9999):
377
+ intNullVal=-9999,
378
+ strNullVal='None',
379
+ nthreads=1):
265
380
  '''
266
381
  Executes the sql query and returns the tuple or dictionary
267
382
  with the numpy arrays.
@@ -284,6 +399,9 @@ def get(query,
284
399
  intNullVal : integer, optional
285
400
  All the integer columns with nulls will have null replaced by
286
401
  this value
402
+ strNullVal : string, optional
403
+ All the string columns with nulls will have null replaced by
404
+ this value
287
405
  db : string
288
406
  The name of the database
289
407
  driver : string, optional
@@ -298,7 +416,16 @@ def get(query,
298
416
  Port of the database
299
417
  preamb : string
300
418
  SQL code to be executed before the query
301
-
419
+ batched : bool, optional
420
+ Whether to use named or not-named cursor with PostgreSQL
421
+ The default setting of True leads to retrieval of results in batches
422
+ and on the fly conversion of
423
+ retrieved results into numpy, thus it will not use more memory than
424
+ needed to store the results of the query. The
425
+ The batched=False will use at least twice the amount of memory
426
+ needed to store the results, but with the benefit of faster query
427
+ execution, because these queries use PostgreSQL parallelism and thus
428
+ can be ~ a factor of few faster
302
429
  Returns
303
430
  -------
304
431
  ret : Tuple or dictionary
@@ -326,119 +453,114 @@ def get(query,
326
453
  port=port,
327
454
  timeout=timeout)
328
455
  try:
329
- cur = getCursor(conn, driver=driver, preamb=preamb, notNamed=notNamed)
456
+ if driver == 'psycopg2':
457
+ warnings.warn('psycopg2 driver is not supported anymore. '
458
+ 'We using psycopg instead')
459
+ driver = 'psycopg'
460
+ if batched:
461
+ # This is related to postgresql server cursors vs
462
+ # application cursors
463
+ namedCursor = True
464
+ else:
465
+ namedCursor = False
466
+ cur = getCursor(conn, driver=driver, preamb=preamb, named=namedCursor)
330
467
 
331
468
  if params is None:
332
469
  res = cur.execute(query)
333
470
  else:
334
471
  res = cur.execute(query, params)
335
472
 
336
- qIn = queue.Queue(1)
337
- qOut = queue.Queue()
338
- endEvent = threading.Event()
339
- nrec = 0 # keeps the number of arrays sent to the other thread
340
- # minus number received
341
- reslist = []
342
- proc = None
343
473
  colNames = []
344
- if driver == 'psycopg2':
345
- warnings.warn('psycopg2 driver is not supported anymore. '
346
- 'We using psycopg instead')
347
- driver = 'psycopg'
348
474
  if driver == 'psycopg':
349
- try:
350
- while True:
351
- # Iterating over the cursor, retrieving batches of results
352
- # and then sending them for conversion
353
- tups = cur.fetchmany()
354
- desc = cur.description
355
-
356
- # If the is just the start we need to launch the
357
- # thread doing the conversion
358
- no_results = tups == []
359
- if nrec == 0:
360
- typeCodes = [_tmp.type_code for _tmp in desc]
361
- colNames = [_tmp.name for _tmp in cur.description]
362
-
363
- # No more data
364
- if no_results:
365
- dtype = __getDType([None] * len(typeCodes), typeCodes,
366
- strLength)
367
- break
368
-
369
- dtype = __getDType(tups[0], typeCodes, strLength)
370
-
371
- # Send the new batch for conversion
372
- qIn.put(tups)
373
-
374
- # If the is just the start we need to launch the
375
- # thread doing the conversion
376
- if nrec == 0:
377
- proc = threading.Thread(target=__converter,
378
- args=(qIn, qOut, endEvent,
379
- dtype, intNullVal))
380
- proc.start()
381
-
382
- # nrec is the number of batches in conversion currently
383
- nrec += 1
384
-
385
- # Try to retrieve one processed batch without waiting
386
- # on it
387
- try:
388
- reslist.append(qOut.get(False))
389
- nrec -= 1
390
- except queue.Empty:
391
- pass
392
-
393
- # Now we are done fetching the data from the DB, we
394
- # just need to assemble the converted results
395
- while (nrec != 0):
396
- try:
397
- reslist.append(qOut.get(True, 0.1))
398
- nrec -= 1
399
- except queue.Empty:
400
- # continue looping unless the endEvent was set
401
- # which should happen in the case of the crash
402
- # of the converter thread
403
- if endEvent.is_set():
404
- raise Exception('Child thread failed')
405
- endEvent.set()
406
- except BaseException:
407
- endEvent.set()
408
- if proc is not None:
409
- # notice that here the timeout is larger than the timeout
410
- proc.join(0.2)
411
- # in the converter process
412
- if proc.is_alive():
413
- # could not kill
414
- pass
415
- raise
416
- if proc is not None:
417
- proc.join()
418
- if reslist == []:
419
- nCols = len(desc)
420
- res = numpy.array([],
421
- dtype=numpy.dtype([('a%d' % i, 'f')
422
- for i in range(nCols)]))
475
+ # Fetch first batch to determine dtype
476
+ if batched:
477
+ first_batch = cur.fetchmany(config.arraysize)
478
+ else:
479
+ first_batch = cur.fetchall()
480
+ desc = cur.description
481
+
482
+ # Check if we have description
483
+ # (for non-select queries it might be None)
484
+ if desc:
485
+ type_codes = [_tmp.type_code for _tmp in desc]
486
+ colNames = [_tmp.name for _tmp in desc]
487
+ else:
488
+ type_codes = []
489
+ colNames = []
490
+
491
+ if not first_batch:
492
+ # Empty result set or no result (e.g. insert)
493
+ if type_codes:
494
+ dtype = __getDType([None] * len(type_codes), type_codes,
495
+ strLength)
496
+ res = numpy.array([], dtype=dtype)
497
+ else:
498
+ # Fallback or empty
499
+ dtype = None
500
+ res = []
423
501
  else:
424
- res = numpy.concatenate(reslist)
502
+
503
+ def process_batch(batch):
504
+ # Determine dtype from the first row of the current batch
505
+ # This allows adapting to string lengths in later batches
506
+ # if the first batch had Nulls
507
+ if isinstance(batch, tuple):
508
+ batch = list(batch)
509
+
510
+ first_row = __get_representative_row(batch, type_codes)
511
+ dtype = __getDType(first_row, type_codes, strLength)
512
+ return __fromrecords(batch,
513
+ dtype=dtype,
514
+ intNullVal=intNullVal,
515
+ strNullVal=strNullVal)
516
+
517
+ def batch_iter(first):
518
+ yield first
519
+ while True:
520
+ batch = cur.fetchmany(config.arraysize)
521
+ if not batch:
522
+ break
523
+ yield batch
524
+
525
+ if batched:
526
+ with concurrent.futures.ThreadPoolExecutor(
527
+ max_workers=nthreads) as executor:
528
+ results_list = list(
529
+ executor.map(process_batch,
530
+ batch_iter(first_batch)))
531
+ else:
532
+ results_list = [process_batch(first_batch)]
533
+ # Check if results_list is valid
534
+ if results_list:
535
+ res = numpy.concatenate(results_list)
536
+ else:
537
+ # Should not be reached given logic above, but fallback
538
+ dtype = __getDType(first_batch[0], type_codes, strLength)
539
+ res = numpy.array([], dtype=dtype)
425
540
 
426
541
  elif driver == 'sqlite3':
427
542
  tups = cur.fetchall()
428
543
  colNames = [_tmp[0] for _tmp in cur.description]
429
544
  if len(tups) > 0:
430
545
  res = numpy.rec.array(tups)
546
+ res = __fix_numpy_record_array(res, intNullVal, strNullVal)
431
547
  else:
432
- return [[]] * len(cur.description)
548
+ res = [numpy.array([]) for _ in cur.description]
433
549
  elif driver == 'duckdb':
434
550
  tups = cur.fetchall()
435
551
  colNames = [_tmp[0] for _tmp in cur.description]
436
552
  if len(tups) > 0:
437
553
  res = numpy.rec.array(tups)
554
+ res = __fix_numpy_record_array(res, intNullVal, strNullVal)
438
555
  else:
439
- return [[]] * len(cur.description)
556
+ res = [numpy.array([]) for _ in cur.description]
440
557
 
441
- res = [res[tmp] for tmp in res.dtype.names]
558
+ if isinstance(res, numpy.ndarray) and res.dtype.names:
559
+ res = [res[tmp] for tmp in res.dtype.names]
560
+ elif isinstance(res, list) and not res:
561
+ # Handle empty result logic from original code if needed,
562
+ # but keeping consistent with original returning array list
563
+ pass
442
564
 
443
565
  except BaseException:
444
566
  failure_cleanup(conn, connSupplied)
@@ -517,7 +639,7 @@ def execute(query,
517
639
  host=host,
518
640
  timeout=timeout)
519
641
  try:
520
- cur = getCursor(conn, driver=driver, preamb=preamb, notNamed=True)
642
+ cur = getCursor(conn, driver=driver, preamb=preamb, named=False)
521
643
  if params is not None:
522
644
  cur.execute(query, params)
523
645
  else:
@@ -733,7 +855,7 @@ table/pandas/dictionary or provide a separate list of arrays and their names')
733
855
  fixed_names.append(fixed_name)
734
856
  names = fixed_names
735
857
  try:
736
- cur = getCursor(conn, driver=driver, preamb=preamb, notNamed=True)
858
+ cur = getCursor(conn, driver=driver, preamb=preamb, named=False)
737
859
  if createTable:
738
860
  query1 = __create_schema(tableName, arrays, names, temp=temp)
739
861
  cur.execute(query1)
@@ -777,6 +899,7 @@ def local_join(query,
777
899
  timeout=None,
778
900
  strLength=STRLEN_DEFAULT,
779
901
  intNullVal=-9999,
902
+ strNullVal='None',
780
903
  asDict=False):
781
904
  """
782
905
  Join your local data in python with the data in the database
@@ -829,7 +952,8 @@ def local_join(query,
829
952
  preamb=preamb,
830
953
  strLength=strLength,
831
954
  asDict=asDict,
832
- intNullVal=intNullVal)
955
+ intNullVal=intNullVal,
956
+ strNullVal=strNullVal)
833
957
  except BaseException:
834
958
  failure_cleanup(conn, connSupplied)
835
959
  raise
@@ -0,0 +1 @@
1
+ __version__ = '0.27.0'
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: sqlutilpy
3
- Version: 0.26.0
3
+ Version: 0.27.0
4
4
  Summary: Database query code returning numpy arrays
5
5
  Author-email: "Sergey E. Koposov" <skoposov@ed.ac.uk>
6
6
  License: BSD
@@ -11,6 +11,8 @@ src/sqlutilpy.egg-info/dependency_links.txt
11
11
  src/sqlutilpy.egg-info/requires.txt
12
12
  src/sqlutilpy.egg-info/top_level.txt
13
13
  tests/killer.py
14
+ tests/test_consistency.py
14
15
  tests/test_duckdb.py
16
+ tests/test_fix_numpy.py
15
17
  tests/test_pg.py
16
18
  tests/test_sqlite.py