sqlutilpy 0.26.0__tar.gz → 0.27.0__tar.gz

{sqlutilpy-0.26.0/src/sqlutilpy.egg-info → sqlutilpy-0.27.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlutilpy
-Version: 0.26.0
+Version: 0.27.0
 Summary: Database query code returning numpy arrays
 Author-email: "Sergey E. Koposov" <skoposov@ed.ac.uk>
 License: BSD

{sqlutilpy-0.26.0 → sqlutilpy-0.27.0}/src/sqlutilpy/sqlutil.py

@@ -1,12 +1,11 @@
 """Sqlutilpy module to access SQL databases
 """
-from __future__ import print_function
 import numpy
 import numpy as np
 import psycopg
-import threading
 import collections
 import warnings
+import concurrent.futures
 
 try:
     import astropy.table as atpy
@@ -19,8 +18,6 @@ except ImportError:
     # pandas is not installed
     pandas = None
 
-import queue
-
 _WAIT_SELECT_TIMEOUT = 10
 STRLEN_DEFAULT = 20
 
@@ -99,7 +96,7 @@ def getConnection(db=None,
     return conn
 
 
-def getCursor(conn, driver=None, preamb=None, notNamed=False):
+def getCursor(conn, driver=None, preamb=None, named=True):
     """
     Retrieve the database cursor
     """
@@ -108,14 +105,14 @@ def getCursor(conn, driver=None, preamb=None, notNamed=False):
             'psycopg2 driver is not supported anymore using psycopg instead')
         driver = 'psycopg'
     if driver == 'psycopg':
-        cur = conn.cursor()
+        cur = conn.cursor(scrollable=False)
         if preamb is not None:
             cur.execute(preamb)
         else:
             cur.execute('set cursor_tuple_fraction TO 1')
             # this is required because otherwise PG may decide to execute a
             # different plan
-        if notNamed:
+        if not named:
             return cur
         cur = conn.cursor(name='sqlutilcursor')
         cur.arraysize = config.arraysize
@@ -132,65 +129,75 @@ def getCursor(conn, driver=None, preamb=None, notNamed=False):
     return cur
 
 
-def __fromrecords(recList, dtype=None, intNullVal=None):
+def __fromrecords(recList, dtype=None, intNullVal=None, strNullVal='None'):
     """
     This function was taken from np.core.records and updated to
-    support conversion null integers to intNullVal
+    support conversion null integers to intNullVal and strings to strNullVal
     """
 
-    shape = None
     descr = np.dtype((np.record, dtype))
-    try:
-        retval = np.array(recList, dtype=descr)
-    except TypeError:  # list of lists instead of list of tuples
-        shape = (len(recList), )
-        _array = np.recarray(shape, descr)
+
+    # Check if we must avoid fast path for strings to prevent
+    # None -> 'None' conversion
+    # Note: numpy's default conversion for None to string is 'None', so if
+    # strNullVal is 'None', we can still use the fast path.
+    force_object_strings = False
+    if strNullVal is not None and strNullVal != 'None':
+        if any(dtype.fields[n][0].kind in ('S', 'U') for n in dtype.names):
+            force_object_strings = True
+
+    if not force_object_strings:
         try:
-            for k in range(_array.size):
-                _array[k] = tuple(recList[k])
-        except TypeError:
-            convs = []
-            ncols = len(dtype.fields)
-            for _k in dtype.names:
-                _v = dtype.fields[_k]
-                if _v[0] in [np.int16, np.int32, np.int64]:
-                    convs.append(lambda x: intNullVal if x is None else x)
-                else:
-                    convs.append(lambda x: x)
-            convs = tuple(convs)
-
-            def convF(x):
-                return [convs[_](x[_]) for _ in range(ncols)]
-
-            for k in range(k, _array.size):
-                try:
-                    _array[k] = tuple(recList[k])
-                except TypeError:
-                    _array[k] = tuple(convF(recList[k]))
-        return _array
-    else:
-        if shape is not None and retval.shape != shape:
-            retval.shape = shape
+            retval = np.array(recList, dtype=descr)
+            return retval.view(numpy.recarray)
+        except (TypeError, ValueError):
+            # Failed (likely due to None in int col). Fall through to object
+            # path.
+            pass
 
-    res = retval.view(numpy.recarray)
+    # Vectorized fallback path using object arrays
+    names = dtype.names
+    new_formats = []
+    converters = {}
 
-    return res
+    for name in names:
+        dt = dtype.fields[name][0]
+        if dt.kind in ('i', 'u'):
+            new_formats.append(object)
+            converters[name] = intNullVal
+        elif dt.kind in ('S', 'U') and strNullVal is not None:
+            new_formats.append(object)
+            converters[name] = strNullVal
+        else:
+            new_formats.append(dt)
 
+    temp_dtype = np.dtype({'names': names, 'formats': new_formats})
 
-def __converter(qIn, qOut, endEvent, dtype, intNullVal):
-    """ Convert the input stream of tuples into numpy arrays """
-    while (not endEvent.is_set()):
-        try:
-            tups = qIn.get(True, 0.1)
-        except queue.Empty:
-            continue
-        try:
-            res = __fromrecords(tups, dtype=dtype, intNullVal=intNullVal)
-        except Exception:
-            print('Failed to convert input data into array')
-            endEvent.set()
-            raise
-        qOut.put(res)
+    if len(recList) == 0:
+        return np.recarray((0, ), dtype=descr)
+
+    try:
+        arr = np.array(recList, dtype=temp_dtype)
+    except (TypeError, ValueError):
+        # Handle cases where np.array fails (e.g. ragged lists)
+        # by manually filling the object array
+        shape = (len(recList), )
+        arr = np.recarray(shape, temp_dtype)
+        for k in range(len(recList)):
+            arr[k] = tuple(recList[k])
+
+    # Apply conversions
+    for name, val in converters.items():
+        col = arr[name]
+        # Check for None efficiently
+        mask = (col == None)
+        if np.any(mask):
+            if val is not None:
+                col[mask] = val
+
+    # Cast to final dtype
+    retval = arr.astype(descr)
+    return retval.view(numpy.recarray)
 
 
 def __getDType(row, typeCodes, strLength):
@@ -247,6 +254,112 @@ def __getDType(row, typeCodes, strLength):
     return dtype
 
 
+def __get_representative_row(batch, type_codes):
+    """
+    Scans the batch to find the maximum length of string columns and returns
+    a representative row (based on the first row) with those maximum lengths.
+    This ensures that __getDType determines the correct string length for
+    the array.
+    """
+    first_row = list(batch[0])
+    strTypes = [25, 1042, 1043]
+
+    for i, type_code in enumerate(type_codes):
+        if type_code in strTypes:
+            max_len = 0
+            found = False
+            for row in batch:
+                val = row[i]
+                if val is not None:
+                    max_len = max(max_len, len(val))
+                    found = True
+            if found:
+                # Update representative row with max length string
+                # __getDType uses max(strLength, len(val))
+                if first_row[i] is None or len(first_row[i]) < max_len:
+                    first_row[i] = 'x' * max_len
+    return first_row
+
+
+def __fix_numpy_record_array(res, intNullVal, strNullVal):
+    if res.size == 0:
+        return res
+
+    names = res.dtype.names
+    new_arrays = []
+
+    for name in names:
+        col = res[name]
+        dtype = col.dtype
+
+        if dtype.kind == 'O':
+            # Check for None (vectorized)
+            # We treat everything as candidates if it's Object type
+
+            # Find first non-None value to guess type
+            valid_mask = (col != None)
+            if not np.any(valid_mask):
+                # All None.
+                # We can't determine type. Return as is (Object filled with
+                # None)
+                new_arrays.append(col)
+                continue
+
+            # Take a sample (first valid)
+            first_val = col[valid_mask][0]
+
+            # INTEGER
+            if isinstance(
+                    first_val,
+                (int, np.integer)) and not isinstance(first_val,
+                                                      (bool, np.bool_)):
+                # Verify all valid are integers
+                valid_vals = col[valid_mask]
+                is_valid = all(
+                    isinstance(x, (int, np.integer))
+                    and not isinstance(x, (bool, np.bool_))
+                    for x in valid_vals)
+                if is_valid:
+                    # Convert
+                    temp_col = np.copy(col)
+                    temp_col[~valid_mask] = intNullVal
+                    try:
+                        new_col = temp_col.astype(int)
+                        new_arrays.append(new_col)
+                        continue
+                    except Exception:
+                        pass
+
+            # FLOAT
+            if isinstance(first_val, (float, np.floating)):
+                if all(
+                        isinstance(x, (float, np.floating, int, np.integer))
+                        for x in col[valid_mask]):
+                    temp_col = np.copy(col)
+                    temp_col[~valid_mask] = np.nan
+                    try:
+                        new_col = temp_col.astype(float)
+                        new_arrays.append(new_col)
+                        continue
+                    except Exception:
+                        pass
+
+            # STRING
+            if isinstance(first_val, str):
+                if all(isinstance(x, str) for x in col[valid_mask]):
+                    # Replicate PG behavior: None -> strNullVal
+                    temp_col = np.copy(col)
+                    temp_col[~valid_mask] = strNullVal
+                    # Let numpy infer string length
+                    new_col = np.array(list(temp_col))
+                    new_arrays.append(new_col)
+                    continue
+
+        new_arrays.append(col)
+
+    return numpy.rec.fromarrays(new_arrays, names=names)
+
+
 def get(query,
         params=None,
         db="wsdb",
@@ -259,9 +372,11 @@ def get(query,
         port=None,
         strLength=STRLEN_DEFAULT,
         timeout=None,
-        notNamed=False,
+        batched=True,
         asDict=False,
-        intNullVal=-9999):
+        intNullVal=-9999,
+        strNullVal='None',
+        nthreads=1):
     '''
     Executes the sql query and returns the tuple or dictionary
     with the numpy arrays.
@@ -284,6 +399,9 @@ def get(query,
     intNullVal : integer, optional
         All the integer columns with nulls will have null replaced by
         this value
+    strNullVal : string, optional
+        All the string columns with nulls will have null replaced by
+        this value
     db : string
         The name of the database
     driver : string, optional
@@ -298,7 +416,16 @@ def get(query,
         Port of the database
     preamb : string
         SQL code to be executed before the query
-
+    batched : bool, optional
+        Whether to use named or not-named cursor with PostgreSQL
+        The default setting of True leads to retrieval of results in batches
+        and on the fly conversion of
+        retrieved results into numpy, thus it will not use more memory than
+        needed to store the results of the query. The
+        The batched=False will use at least twice the amount of memory
+        needed to store the results, but with the benefit of faster query
+        execution, because these queries use PostgreSQL parallelism and thus
+        can be ~ a factor of few faster
     Returns
     -------
     ret : Tuple or dictionary
@@ -326,119 +453,114 @@ def get(query,
                              port=port,
                              timeout=timeout)
     try:
-        cur = getCursor(conn, driver=driver, preamb=preamb, notNamed=notNamed)
+        if driver == 'psycopg2':
+            warnings.warn('psycopg2 driver is not supported anymore. '
+                          'We using psycopg instead')
+            driver = 'psycopg'
+        if batched:
+            # This is related to postgresql server cursors vs
+            # application cursors
+            namedCursor = True
+        else:
+            namedCursor = False
+        cur = getCursor(conn, driver=driver, preamb=preamb, named=namedCursor)
 
         if params is None:
             res = cur.execute(query)
         else:
             res = cur.execute(query, params)
 
-        qIn = queue.Queue(1)
-        qOut = queue.Queue()
-        endEvent = threading.Event()
-        nrec = 0  # keeps the number of arrays sent to the other thread
-        # minus number received
-        reslist = []
-        proc = None
         colNames = []
-        if driver == 'psycopg2':
-            warnings.warn('psycopg2 driver is not supported anymore. '
-                          'We using psycopg instead')
-            driver = 'psycopg'
         if driver == 'psycopg':
-            try:
-                while True:
-                    # Iterating over the cursor, retrieving batches of results
-                    # and then sending them for conversion
-                    tups = cur.fetchmany()
-                    desc = cur.description
-
-                    # If the is just the start we need to launch the
-                    # thread doing the conversion
-                    no_results = tups == []
-                    if nrec == 0:
-                        typeCodes = [_tmp.type_code for _tmp in desc]
-                        colNames = [_tmp.name for _tmp in cur.description]
-
-                    # No more data
-                    if no_results:
-                        dtype = __getDType([None] * len(typeCodes), typeCodes,
-                                           strLength)
-                        break
-
-                    dtype = __getDType(tups[0], typeCodes, strLength)
-
-                    # Send the new batch for conversion
-                    qIn.put(tups)
-
-                    # If the is just the start we need to launch the
-                    # thread doing the conversion
-                    if nrec == 0:
-                        proc = threading.Thread(target=__converter,
-                                                args=(qIn, qOut, endEvent,
-                                                      dtype, intNullVal))
-                        proc.start()
-
-                    # nrec is the number of batches in conversion currently
-                    nrec += 1
-
-                    # Try to retrieve one processed batch without waiting
-                    # on it
-                    try:
-                        reslist.append(qOut.get(False))
-                        nrec -= 1
-                    except queue.Empty:
-                        pass
-
-                # Now we are done fetching the data from the DB, we
-                # just need to assemble the converted results
-                while (nrec != 0):
-                    try:
-                        reslist.append(qOut.get(True, 0.1))
-                        nrec -= 1
-                    except queue.Empty:
-                        # continue looping unless the endEvent was set
-                        # which should happen in the case of the crash
-                        # of the converter thread
-                        if endEvent.is_set():
-                            raise Exception('Child thread failed')
-                endEvent.set()
-            except BaseException:
-                endEvent.set()
-                if proc is not None:
-                    # notice that here the timeout is larger than the timeout
-                    proc.join(0.2)
-                    # in the converter process
-                    if proc.is_alive():
-                        # could not kill
-                        pass
-                raise
-            if proc is not None:
-                proc.join()
-            if reslist == []:
-                nCols = len(desc)
-                res = numpy.array([],
-                                  dtype=numpy.dtype([('a%d' % i, 'f')
-                                                     for i in range(nCols)]))
+            # Fetch first batch to determine dtype
+            if batched:
+                first_batch = cur.fetchmany(config.arraysize)
+            else:
+                first_batch = cur.fetchall()
+            desc = cur.description
+
+            # Check if we have description
+            # (for non-select queries it might be None)
+            if desc:
+                type_codes = [_tmp.type_code for _tmp in desc]
+                colNames = [_tmp.name for _tmp in desc]
+            else:
+                type_codes = []
+                colNames = []
+
+            if not first_batch:
+                # Empty result set or no result (e.g. insert)
+                if type_codes:
+                    dtype = __getDType([None] * len(type_codes), type_codes,
+                                       strLength)
+                    res = numpy.array([], dtype=dtype)
+                else:
+                    # Fallback or empty
+                    dtype = None
+                    res = []
             else:
-                res = numpy.concatenate(reslist)
+
+                def process_batch(batch):
+                    # Determine dtype from the first row of the current batch
+                    # This allows adapting to string lengths in later batches
+                    # if the first batch had Nulls
+                    if isinstance(batch, tuple):
+                        batch = list(batch)
+
+                    first_row = __get_representative_row(batch, type_codes)
+                    dtype = __getDType(first_row, type_codes, strLength)
+                    return __fromrecords(batch,
+                                         dtype=dtype,
+                                         intNullVal=intNullVal,
+                                         strNullVal=strNullVal)
+
+                def batch_iter(first):
+                    yield first
+                    while True:
+                        batch = cur.fetchmany(config.arraysize)
+                        if not batch:
+                            break
+                        yield batch
+
+                if batched:
+                    with concurrent.futures.ThreadPoolExecutor(
+                            max_workers=nthreads) as executor:
+                        results_list = list(
+                            executor.map(process_batch,
+                                         batch_iter(first_batch)))
+                else:
+                    results_list = [process_batch(first_batch)]
+                # Check if results_list is valid
+                if results_list:
+                    res = numpy.concatenate(results_list)
+                else:
+                    # Should not be reached given logic above, but fallback
+                    dtype = __getDType(first_batch[0], type_codes, strLength)
+                    res = numpy.array([], dtype=dtype)
 
         elif driver == 'sqlite3':
             tups = cur.fetchall()
             colNames = [_tmp[0] for _tmp in cur.description]
             if len(tups) > 0:
                 res = numpy.rec.array(tups)
+                res = __fix_numpy_record_array(res, intNullVal, strNullVal)
             else:
-                return [[]] * len(cur.description)
+                res = [numpy.array([]) for _ in cur.description]
         elif driver == 'duckdb':
             tups = cur.fetchall()
             colNames = [_tmp[0] for _tmp in cur.description]
             if len(tups) > 0:
                 res = numpy.rec.array(tups)
+                res = __fix_numpy_record_array(res, intNullVal, strNullVal)
             else:
-                return [[]] * len(cur.description)
+                res = [numpy.array([]) for _ in cur.description]
 
-        res = [res[tmp] for tmp in res.dtype.names]
+        if isinstance(res, numpy.ndarray) and res.dtype.names:
+            res = [res[tmp] for tmp in res.dtype.names]
+        elif isinstance(res, list) and not res:
+            # Handle empty result logic from original code if needed,
+            # but keeping consistent with original returning array list
+            pass
 
     except BaseException:
         failure_cleanup(conn, connSupplied)
@@ -517,7 +639,7 @@ def execute(query,
                              host=host,
                              timeout=timeout)
     try:
-        cur = getCursor(conn, driver=driver, preamb=preamb, notNamed=True)
+        cur = getCursor(conn, driver=driver, preamb=preamb, named=False)
         if params is not None:
             cur.execute(query, params)
         else:
@@ -733,7 +855,7 @@ table/pandas/dictionary or provide a separate list of arrays and their names')
         fixed_names.append(fixed_name)
     names = fixed_names
     try:
-        cur = getCursor(conn, driver=driver, preamb=preamb, notNamed=True)
+        cur = getCursor(conn, driver=driver, preamb=preamb, named=False)
         if createTable:
             query1 = __create_schema(tableName, arrays, names, temp=temp)
             cur.execute(query1)
@@ -777,6 +899,7 @@ def local_join(query,
                timeout=None,
                strLength=STRLEN_DEFAULT,
                intNullVal=-9999,
+               strNullVal='None',
                asDict=False):
     """
     Join your local data in python with the data in the database
@@ -829,7 +952,8 @@ def local_join(query,
                   preamb=preamb,
                   strLength=strLength,
                   asDict=asDict,
-                  intNullVal=intNullVal)
+                  intNullVal=intNullVal,
+                  strNullVal=strNullVal)
     except BaseException:
         failure_cleanup(conn, connSupplied)
         raise

sqlutilpy-0.27.0/src/sqlutilpy/version.py

@@ -0,0 +1 @@
+__version__ = '0.27.0'

{sqlutilpy-0.26.0 → sqlutilpy-0.27.0/src/sqlutilpy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sqlutilpy
-Version: 0.26.0
+Version: 0.27.0
 Summary: Database query code returning numpy arrays
 Author-email: "Sergey E. Koposov" <skoposov@ed.ac.uk>
 License: BSD

{sqlutilpy-0.26.0 → sqlutilpy-0.27.0}/src/sqlutilpy.egg-info/SOURCES.txt

@@ -11,6 +11,8 @@ src/sqlutilpy.egg-info/dependency_links.txt
 src/sqlutilpy.egg-info/requires.txt
 src/sqlutilpy.egg-info/top_level.txt
 tests/killer.py
+tests/test_consistency.py
 tests/test_duckdb.py
+tests/test_fix_numpy.py
 tests/test_pg.py
 tests/test_sqlite.py