Efficient way to create DataFrame with different column types

user_1349428 0
Viewed 182

I need to read data from numeric Postgres table and create DataFrame accordingly.

The default way Pandas is doing it is by using DataFrame.from_records:

df = DataFrame.from_records(data,
                            columns=columns,
                            coerce_float=coerce_float)

When data looks like:

[(0.16275345863180396, 0.16275346), (0.6356328878675244, 0.6356329)...]

And columns looks like:

['a', 'b']

The problem is that the generated DataFrame ignores the original Posgres types: double precision and real.

As I use huge DataFrames and my data is mostly real I'd like to explicitly specify the column types.

So I tried:

df = DataFrame.from_records(np.array(data, dtype=columns),
                            coerce_float=coerce_float)

When data is the same, but columns looks like:

[('a', 'float64'), ('b', 'float32')]

(types are extracted from Postgres as a part of query and converted to Numpy dtypes)

This approach works, but DataFrame construction is 2-3 times slower (for 2M rows DataFrames it takes several seconds), because np.array generation is for some reason very slow. In real life I have 10-200 columns mostly float32.

What is the fastest way to construct DataFrame with specified column types?

pandas dataframe numpy postgresql python
4 Answers

If you know the data columns and its types already, then following format will help to generate data frame with specified datatypes.

    pd.DataFrame(data, columns = columnList, dtype = np.dtype([('type1','type2')]))
user_72469660

You can read the using the standard psycopg2 driver. Here, you can register your own type caster to convert REAL to np.float32 instead of the default python float (the OID of the data type - 700 for REAL - can either be obtained as described in the link or taken from here):

import psycopg2
import numpy as np
import pandas as pd

real_oid = 700
REAL2FLOAT32 = psycopg2.extensions.new_type((real_oid,), 'REAL2FLOAT32', lambda val, cur: np.float32(val))
psycopg2.extensions.register_type(REAL2FLOAT32)

with psycopg2.connect('postgresql://user:pwd@localhost:5432/test') as con:
    with con.cursor() as cur:
        cur.execute('select 0.16275345863180396::double precision, 0.16275346::real')
        # print(cur.description) # to get the OID for real
        rows = cur.fetchall()
        df = pd.DataFrame(rows, columns=['a', 'b'])

Output of df.info():

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 1 entries, 0 to 0
Data columns (total 2 columns):
 #   Column  Non-Null Count  Dtype  
---  ------  --------------  -----  
 0   a       1 non-null      float64
 1   b       1 non-null      float32
dtypes: float32(1), float64(1)
memory usage: 140.0 bytes

If you like, you can integrate this directly into pandas that uses SQLAlchemy in the background like this:

import sqlalchemy

real_oid = 700
REAL2FLOAT32 = psycopg2.extensions.new_type((real_oid,), "REAL2FLOAT32", lambda val, cur: np.float32(val))
psycopg2.extensions.register_type(REAL2FLOAT32)

with sqlalchemy.create_engine('postgresql://user:pwd@localhost:5432/test').connect() as con:
    df = pd.read_sql_query('select 0.16275345863180396::double precision, 0.16275346::real', con)

As for speed, asyncpg claims to be much faster than psycopg2. Here you should be able to use set_type_codec to make your own conversion as shown in this example, but I didn't test it.

user_39443220

IIUC, this is probably helpful:

data = [(0.16275345863180396, 0.16275346), (0.6356328878675244, 0.6356329)]
columns = ['a', 'b']
df = pd.DataFrame(data, columns=columns, dtype='float64').astype({'b':'float32'})
print(df.info()

Output:

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 2 entries, 0 to 1
Data columns (total 2 columns):
 #   Column  Non-Null Count  Dtype  
---  ------  --------------  -----  
 0   a       2 non-null      float64
 1   b       2 non-null      float32
dtypes: float32(1), float64(1)
memory usage: 152.0 bytes
user_74758380

Try connecting to the Postgresql database and read directly to pandas data frame. Not sure if you already tried this way.

import pandas as pd
import psycopg2 as pg
connection= pg.connect("dbname='dbname' user='pguser' host='127.0.0.1' port='15432' password='password'")
df = pd.read_sql('select * from table', connection)
user_143842370
Related