From dd956116b73416de454e6f5b0fbeb6ddcba1f0c4 Mon Sep 17 00:00:00 2001 From: llllllllll Date: Wed, 5 Aug 2015 14:31:54 -0400 Subject: [PATCH] ENH: sqlalchemy --- etc/requirements.txt | 2 + zipline/assets/asset_writer.py | 358 +++++++++++++------------------ zipline/assets/assets.py | 375 +++++++++++++++++---------------- 3 files changed, 348 insertions(+), 387 deletions(-) diff --git a/etc/requirements.txt b/etc/requirements.txt index 34c99b95..43a504e9 100644 --- a/etc/requirements.txt +++ b/etc/requirements.txt @@ -45,3 +45,5 @@ bcolz==0.10.0 # Command line interface helper click==4.0.0 + +toolz==0.7.2 diff --git a/zipline/assets/asset_writer.py b/zipline/assets/asset_writer.py index 2dbc9c8c..b8190cdb 100644 --- a/zipline/assets/asset_writer.py +++ b/zipline/assets/asset_writer.py @@ -3,10 +3,11 @@ from abc import ( abstractmethod, ) -import pandas as pd import numpy as np +import pandas as pd from pandas.tseries.tools import normalize_date from six import with_metaclass, string_types +import sqlalchemy as sa from zipline.errors import ( ConsumeAssetMetaDataError, @@ -17,7 +18,7 @@ from zipline.assets import ( Asset, Equity, Future ) -ASSET_FIELDS = [ +ASSET_FIELDS = frozenset({ 'sid', 'asset_type', 'symbol', @@ -34,43 +35,44 @@ ASSET_FIELDS = [ 'file_name', # Used as symbol 'company_name', # Used as asset_name 'start_date_nano', # Used as start_date - 'end_date_nano' # Used as end_date -] + 'end_date_nano', # Used as end_date +}) # Expected fields for an Asset's metadata -ASSET_TABLE_FIELDS = [ +ASSET_TABLE_FIELDS = frozenset({ 'sid', 'symbol', 'asset_name', 'start_date', 'end_date', 'first_traded', - 'exchange' -] + 'exchange', +}) + # Expected fields for an Asset's metadata -FUTURE_TABLE_FIELDS = ASSET_TABLE_FIELDS + [ +FUTURE_TABLE_FIELDS = ASSET_TABLE_FIELDS | { 'root_symbol_id', 'notice_date', 'expiration_date', 'contract_multiplier', -] +} EQUITY_TABLE_FIELDS = ASSET_TABLE_FIELDS -EXCHANGE_TABLE_FIELDS = [ +EXCHANGE_TABLE_FIELDS = frozenset({ 'exchange_id', 'exchange', 'timezone' -] +}) -ROOT_SYMBOL_TABLE_FIELDS = [ +ROOT_SYMBOL_TABLE_FIELDS = ({ 'root_symbol_id', 'root_symbol', 'sector', 'description', 'exchange_id' -] +}) class AssetDBWriter(with_metaclass(ABCMeta)): @@ -81,237 +83,173 @@ class AssetDBWriter(with_metaclass(ABCMeta)): Methods ------- - write_all(db_conn, fuzzy_char=None, allow_sid_assignment=True, + write_all(engine, fuzzy_char=None, allow_sid_assignment=True, constraints=False) Write the data supplied at initialization to the database. - init_db(db_conn, constraints=False) + init_db(engine, constraints=False) Create the SQLite tables (called by write_all). load_data() Returns data in standard format. """ + def __init__(self): + self.sql_metadata = None - def write_all(self, db_conn, fuzzy_char=None, allow_sid_assignment=True, + def write_all(self, + engine, + fuzzy_char=None, constraints=False): """ Write pre-supplied data to SQLite. Parameters ---------- - db_conn: sqlite3.Connection - A connection to a SQLite database. - fuzzy_char: string + engine : Engine + An engine to a SQL database. + fuzzy_char : str, optional A string for use in fuzzy matching. - allow_sid_assignment: boolean + allow_sid_assignment: bool, optional If True then the class can assign sids where necessary. - constraints: boolean + constraints : bool, optional If True, create SQL ForeignKey and Index constraints. """ - self.fuzzy_char = fuzzy_char - self.allow_sid_assignment = allow_sid_assignment - if allow_sid_assignment: - ts = normalize_date(pd.Timestamp('now', tz='UTC')) - # Store as seconds since UNIX Epoch for compatibility - # with SQL. - self.end_date_to_assign = (ts.value // 10 ** 9) - - # Store a nested-dict of all metadata for - # reference when building Assets - self.metadata_cache = {} - # Create SQL tables - self.init_db(db_conn, constraints) - + self.init_db(engine, constraints) # Get the data to add to SQL equities, futures, exchanges, root_symbols = self.load_data() + with engine.begin() as txn: + self._write_exchanges(exchanges, txn) + self._write_root_symbols(root_symbols, txn) + self._write_futures(futures, txn) + self._write_equities(equities, txn) - c = db_conn.cursor() - # Write to the SQL tables, using the raw SQL driver instead - # of the pandas.DataFrame.to_sql method, as the former - # allows us to create an SQL transaction - c.execute('BEGIN') - # Everything between here and the db_conn.commit() - # will be part of of the same SQL transaction. - self._write_exchanges(exchanges, db_conn) - self._write_root_symbols(root_symbols, db_conn) - self._write_futures(futures, db_conn) - self._write_equities(equities, db_conn) - db_conn.commit() + def _write_exchanges(self, exchanges, bind=None): + self.futures_exchanges.insert().values( + exchanges.to_records(), + ).execute(bind=bind) - def _write_exchanges(self, exchanges, db_conn): + def _write_root_symbols(self, root_symbols, bind=None): + self.futures_root_symbols.insert().values( + root_symbols.to_records(), + ).execute(bind=bind) - data = [tuple(x) for x in exchanges.to_records()] - - c = db_conn.cursor() - # The OR IGNORE syntax means we do not insert data - # which would violate an SQL constraint. - c.executemany(""" - INSERT OR IGNORE INTO futures_exchanges - ('exchange_id', 'exchange', 'timezone') - VALUES (?, ?, ?) - """, data) - - def _write_root_symbols(self, root_symbols, db_conn): - - data = [tuple(x) for x in root_symbols.to_records()] - - c = db_conn.cursor() - c.executemany(""" - INSERT OR IGNORE INTO futures_root_symbols - ('root_symbol_id', 'root_symbol', 'sector', - 'description', 'exchange_id') - VALUES(?, ?, ?, ?, ?) - """, data) - - def _write_futures(self, futures, db_conn): - - # Retrieve the data to add to the futures_contracts table - data = [tuple(x) for x in futures.to_records()] - - # Retrieve the data to add to the asset_router table - sids = [x[0] for x in data] - futs = ['future'] * len(sids) - to_insert = zip(sids, futs) - - c = db_conn.cursor() - c.executemany(""" - INSERT OR IGNORE INTO futures_contracts - ('sid', 'symbol', 'root_symbol', 'asset_name', - 'start_date', 'end_date', 'first_traded', 'exchange', - 'notice_date', 'expiration_date', 'contract_multiplier') - VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) - """, data) - - c.executemany(""" - INSERT INTO asset_router - ('sid', 'asset_type') - VALUES(?,?) - """, to_insert) - - def _write_equities(self, equities, db_conn): + def _write_futures(self, futures, bind=None): + recs = futures.to_records() + self.futures_contracts.insert().values(recs).execute(bind=bind) + self.asset_router.insert().values([ + (rec['sid'], 'future') for rec in recs + ]).execute(bind=bind) + def _write_equities(self, equities, fuzzy_char, bind=None): # Apply fuzzy matching. - if self.fuzzy_char: - equities['fuzzy'] = equities['symbol'].str.\ - replace(self.fuzzy_char, '') + if fuzzy_char: + equities['fuzzy'] = equities['symbol'].str.replace(fuzzy_char, '') - # Retrieve the data to add to the equitites table - data = [tuple(x) for x in equities.to_records()] + recs = equities.to_records() + self.equities.insert().values(recs).execute(bind=bind) + self.asset_router.insert().values([ + (rec['sid'], 'equity') for rec in recs + ]).execute(bind=bind) - # Retrieve the data to add to the asset_router table - sids = [x[0] for x in data] - eqs = ['equity'] * len(sids) - to_insert = zip(sids, eqs) - - c = db_conn.cursor() - c.executemany(""" - INSERT OR IGNORE INTO equities - ('sid', 'symbol', 'asset_name', 'start_date', - 'end_date', 'first_traded', 'exchange', 'fuzzy') - VALUES(?, ?, ?, ?, ?, ?, ?, ?) - """, data) - - c.executemany(""" - INSERT INTO asset_router - ('sid', 'asset_type') - VALUES(?,?) - """, to_insert) - - def init_db(self, - db_conn, - constraints=False): + def init_db(self, engine, constraints=False): """Connect to database and create tables. Parameters ---------- - db_conn: sqlite3.Connection - A connection to a SQLite database. - constraints: boolean + engine : Engine + An engine to a SQL database. + constraints : bool If True, create SQL ForeignKey and Index constraints. """ - - c = db_conn.cursor() - - c.execute(""" - CREATE TABLE IF NOT EXISTS equities ( - sid INTEGER NOT NULL, - symbol TEXT, - asset_name TEXT, - start_date INTEGER DEFAULT 0, - end_date INTEGER, - first_traded INTEGER, - exchange TEXT, - fuzzy TEXT - )""") - - c.execute(""" - CREATE TABLE IF NOT EXISTS futures_exchanges ( - exchange_id INTEGER NOT NULL, - exchange TEXT, - timezone TEXT - )""") - - c.execute(""" - CREATE TABLE IF NOT EXISTS futures_root_symbols ( - root_symbol_id INTEGER NOT NULL, - root_symbol TEXT, - sector TEXT, - description TEXT, - exchange_id INTEGER{fk} - )""".format(fk=", FOREIGN KEY(exchange_id) REFERENCES " - "futures_exchanges(exchange_id)" - if constraints else "")) - - c.execute(""" - CREATE TABLE IF NOT EXISTS futures_contracts ( - sid INTEGER NOT NULL, - symbol TEXT, - root_symbol_id INTEGER, - root_symbol TEXT, - asset_name TEXT, - start_date INTEGER DEFAULT 0, - end_date INTEGER, - first_traded INTEGER, - exchange_id INTEGER, - exchange TEXT, - notice_date INTEGER, - expiration_date INTEGER, - contract_multiplier REAL{fk} - )""".format(fk=", FOREIGN KEY(exchange_id) REFERENCES " - "futures_exchanges(exchange_id), " - "FOREIGN KEY(root_symbol_id) REFERENCES " - "futures_root_symbols(root_symbol_id)" - if constraints else "")) - - c.execute(""" - CREATE TABLE IF NOT EXISTS asset_router - (sid integer, - asset_type text - )""") - - # Note: Would be optimal to use INTEGER PRIMARY KEY here, but SQLite - # tables cannot be modified after creation. Using CREATE UNIQUE INDEX - # only marginally less performant. - if constraints: - - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_equities_sid ' - 'ON equities(sid)') - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_equities_symbol ' - 'ON equities(symbol)') - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_equities_fuzzy ' - 'ON equities(fuzzy)') - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_futures_exchanges_en ' # noqa - 'ON futures_exchanges(exchange_id)') - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_futures_contracts_sid ' # noqa - 'ON futures_contracts(sid)') - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_futures_root_symbols_id ' # noqa - 'ON futures_root_symbols(root_symbol_id)') - c.execute('CREATE UNIQUE INDEX IF NOT EXISTS ix_asset_router_sid ' - 'ON asset_router(sid)') - - # Note: Also need a max_date table. - - db_conn.commit() + self.sql_metadata = metadata = sa.MetaData(bind=engine) + self.equities = sa.Table( + 'equities', + metadata, + sa.Column( + 'sid', + sa.Integer, + nullable=False, # in case constraints is False + primary_key=constraints, + ), + sa.Column('symbol', sa.Text), + sa.Column('asset_name', sa.Text), + sa.Column('start_date', sa.Integer, default=0), + sa.Column('end_date', sa.Integer), + sa.Column('first_traded', sa.Integer), + sa.Column('exchange', sa.Text), + sa.Column('fuzzy', sa.Text), + ) + self.futures_exchanges = sa.Table( + 'futures_exchanges', + metadata, + sa.Column( + 'exchange_id', + sa.Integer, + nullable=False, # in case constraints is False + primary_key=constraints, + ), + sa.Column('exchange', sa.Text), + sa.Column('timezone', sa.Text), + ) + self.futures_root_symbols = sa.Table( + 'futures_root_symbols', + metadata, + sa.Column( + 'root_symbol_id', + sa.Integer, + nullable=False, # in case constraints is False + primary_key=constraints, + ), + sa.Column('root_symbol', sa.Text), + sa.Column('sector', sa.Text), + sa.Column('description', sa.Text), + sa.Column( + 'exchange_id', + sa.Integer, + *((sa.ForeignKey(self.futures_exchanges.c.exchange_id),) + if constraints else ()) + ), + ) + self.futures_contracts = sa.Table( + 'futures_contracts', + metadata, + sa.Column( + 'sid', + sa.Integer, + nullable=False, # in case constraints is False + primary_key=constraints, + ), + sa.Column('symbol', sa.Text), + sa.Column( + 'root_symbol_id', + sa.Integer, + *((sa.ForeignKey(self.futures_root_symbols.c.root_symbol_id),) + if constraints else ()) + ), + sa.Column('root_symbol', sa.Text), + sa.Column('asset_name', sa.Text), + sa.Column('start_date', sa.Integer, default=0), + sa.Column('end_date', sa.Integer), + sa.Column('first_traded', sa.Integer), + sa.Column( + 'exchange_id', + sa.Integer, + *((sa.ForeignKey(self.futures_exchanges.c.exchange_id),) + if constraints else ()) + ), + sa.column('exchange', sa.Text), + sa.Column('notice_date', sa.Integer), + sa.Column('expiration_date', sa.Integer), + sa.Column('contract_multiplier', sa.Float), + ) + self.asset_router = sa.Table( + 'asset_router', + metadata, + sa.Column('sid', sa.Integer, primary_key=constraints), + sa.Column('asset_type', sa.Text), + ) + metadata.create_all(checkfirst=True) + return metadata @staticmethod def dict_subset(dict_, subset): diff --git a/zipline/assets/assets.py b/zipline/assets/assets.py index 05c8bcb4..64c684d9 100644 --- a/zipline/assets/assets.py +++ b/zipline/assets/assets.py @@ -13,15 +13,18 @@ # limitations under the License. from abc import ABCMeta +from functools import partial, itemgetter from numbers import Integral -from sqlite3 import Row +from operator import getitem import warnings -import numpy as np from logbook import Logger +import numpy as np import pandas as pd from pandas.tseries.tools import normalize_date from six import with_metaclass, string_types +import sqlalchemy as sa +from toolz import compose from zipline.errors import ( MultipleSymbolsFound, @@ -31,34 +34,74 @@ from zipline.errors import ( MapAssetIdentifierIndexError, ) from zipline.assets import ( - Asset, Equity, Future + Asset, Equity, Future, ) from zipline.assets.asset_writer import ( FUTURE_TABLE_FIELDS, - EQUITY_TABLE_FIELDS + EQUITY_TABLE_FIELDS, ) log = Logger('assets.py') -# Create the query once from the fields, so that the join is not done -# repeatedly. -FUTURE_BY_SID_QUERY = 'select {0} from futures_contracts where sid=?'.format( - ", ".join(FUTURE_TABLE_FIELDS)) - -EQUITY_BY_SID_QUERY = 'select {0} from equities where sid=?'.format( - ", ".join(EQUITY_TABLE_FIELDS)) - class AssetFinder(object): - def __init__(self, conn, allow_sid_assignment=True, - fuzzy_char=None): + def __init__(self, engine, allow_sid_assignment=True, fuzzy_char=None): self.fuzzy_char = fuzzy_char self.allow_sid_assignment = allow_sid_assignment - self.conn = conn + self.engine = engine + metadata = sa.Metadata(bind=engine) + self.equities = equities = sa.Table( + 'equities', + metadata, + autoload_with=engine, + ) + self.futures_exchanges = sa.Table( + 'futures_exchanges', + metadata, + autoload_with=engine, + ) + self.futures_root_symbols = sa.Table( + 'futures_root_symbols', + metadata, + autoload_with=engine, + ) + self.futures_contracts = futures_contracts = sa.Table( + 'futures_contracts', + metadata, + autoload_with=engine, + ) + self.asset_router = sa.Table( + 'asset_router', + metadata, + autoload_with=engine, + ) + # Create the equity and future queries once. + _equity_sid = equities.c.sid + _equity_by_sid = sa.select( + tuple(map(partial(getitem, equities.c), EQUITY_TABLE_FIELDS)), + ) + + def select_equity_by_sid(sid): + return _equity_by_sid.where(_equity_sid == int(sid)) + + self.select_equity_by_sid = select_equity_by_sid + + _future_sid = futures_contracts.c.sid + _future_by_sid = sa.select( + tuple(map( + partial(getitem, futures_contracts.c), + FUTURE_TABLE_FIELDS, + )), + ) + + def select_future_by_sid(sid): + return _future_by_sid.where(_future_sid == int(sid)) + + self.select_future_by_sid = select_future_by_sid # Cache for lookup of assets by sid, the objects in the asset lookp may # be shared with the results from equity and future lookup caches. # @@ -86,18 +129,12 @@ class AssetFinder(object): except KeyError: pass - c = self.conn.cursor() - # Python 3 compatibility required forcing to int for sid = 0. - t = (int(sid),) - query = 'SELECT asset_type FROM asset_router WHERE sid=:sid' - c.execute(query, t) - data = c.fetchone() - if data is None: - return - - asset_type = data[0] - self._asset_type_cache[sid] = asset_type + asset_type = sa.select((self.asset_router.c.asset_type,)).where( + self.asset_router.c.sid == int(sid), + ).scalar().execute() + if asset_type is not None: + self._asset_type_cache[sid] = asset_type return asset_type def retrieve_asset(self, sid, default_none=False): @@ -139,11 +176,7 @@ class AssetFinder(object): except KeyError: pass - c = self.conn.cursor() - c.row_factory = Row - t = (int(sid),) - c.execute(EQUITY_BY_SID_QUERY, t) - data = dict(c.fetchone()) + data = self.select_equity_by_sid(sid).execute().fetchone() if data: if data['start_date']: data['start_date'] = pd.Timestamp(data['start_date'], tz='UTC') @@ -171,11 +204,7 @@ class AssetFinder(object): except KeyError: pass - c = self.conn.cursor() - t = (int(sid),) - c.row_factory = Row - c.execute(FUTURE_BY_SID_QUERY, t) - data = dict(c.fetchone()) + data = self.select_future_by_sid(sid).execute().fetchone() if data: if data['start_date']: data['start_date'] = pd.Timestamp(data['start_date'], tz='UTC') @@ -209,79 +238,69 @@ class AssetFinder(object): If multiple Assets are found and as_of_date is not set, raises MultipleSymbolsFound. - If no Asset was active at as_of_date, and allow_expired is False - raises SymbolNotFound. + If no Asset was active at as_of_date raises SymbolNotFound. """ if as_of_date is not None: as_of_date = pd.Timestamp(normalize_date(as_of_date)) - c = self.conn.cursor() - + equities_cols = self.equities.c if as_of_date: - # If one SID exists for symbol, return that symbol - t = (symbol, as_of_date.value, as_of_date.value) - query = ("SELECT sid FROM equities " - "WHERE symbol=? " - "AND start_date<=? " - "AND end_date>=?") - c.execute(query, t) - candidates = c.fetchall() + ad_value = as_of_date.value + # If one SID exists for symbol, return that symbol + candidates = sa.select((equities_cols.sid,)).where( + (equities_cols.symbol == symbol) & + (equities_cols.start_date <= ad_value) & + (equities_cols.end_date >= ad_value), + ).execute().fetchall() if len(candidates) == 1: - return self._retrieve_equity(candidates[0][0]) + return self._retrieve_equity(candidates[0]['sid']) # If no SID exists for symbol, return SID with the # highest-but-not-over end_date - if len(candidates) == 0: - t = (symbol, as_of_date.value) - query = ("SELECT sid FROM equities " - "WHERE symbol=? " - "AND start_date<=? " - "ORDER BY end_date DESC " - "LIMIT 1") - c.execute(query, t) - data = c.fetchone() - - if data: - return self._retrieve_equity(data[0]) + elif not candidates: + sid = sa.select((equities_cols.sid,)).where( + (equities_cols.symbol == symbol) & + (equities_cols.start_date <= ad_value), + ).order_by( + equities_cols.end_date.desc(), + ).scalar().execute() + if sid: + return self._retrieve_equity(sid) # If multiple SIDs exist for symbol, return latest start_date with # end_date as a tie-breaker - if len(candidates) > 1: - t = (symbol, as_of_date.value) - query = ("SELECT sid FROM equities " - "WHERE symbol=? " + - "AND start_date<=? " + - "ORDER BY start_date DESC, end_date DESC " + - "LIMIT 1") - c.execute(query, t) - data = c.fetchone() - - if data: - return self._retrieve_equity(data[0]) + elif len(candidates) > 1: + sid = sa.select((equities_cols.sid,)).where( + (equities_cols.symbol == symbol) & + (equities_cols.start_date <= ad_value), + ).order_by( + equities_cols.start_date.desc(), + equities_cols.end_date.desc(), + ).scalar().execute() + if sid: + return self._retrieve_equity(sid) raise SymbolNotFound(symbol=symbol) else: - t = (symbol,) - query = ("SELECT sid FROM equities WHERE symbol=?") - c.execute(query, t) - data = c.fetchall() - - if len(data) == 1: - return self._retrieve_equity(data[0][0]) - elif not data: + sids = sa.select((equities_cols.sid,)).where( + equities_cols.symbol == sid, + ).execute().fetchall() + if len(sids) == 1: + return self._retrieve_equity(sids[0]['sid']) + elif not sids: raise SymbolNotFound(symbol=symbol) else: - options = [] - for row in data: - sid = row[0] - asset = self._retrieve_equity(sid) - options.append(asset) - raise MultipleSymbolsFound(symbol=symbol, - options=options) + raise MultipleSymbolsFound( + symbol=symbol, + options=list(map( + compose(self._retrieve_equity, itemgetter('sid')), + sids, + )) + ) - def lookup_symbol(self, symbol, as_of_date, fuzzy=False): + def lookup_symbol(self, symbol, as_of_date, fuzzy=None): """ If a fuzzy string is provided, then we try various symbols based on the provided symbol. This is to facilitate mapping from a broker's @@ -291,41 +310,39 @@ class AssetFinder(object): so we can find a match by inserting an underscore. """ symbol = symbol.upper() - as_of_date = normalize_date(as_of_date) + ad_value = normalize_date(as_of_date).value - if not fuzzy: + if fuzzy is None: try: return self.lookup_symbol_resolve_multiple(symbol, as_of_date) except SymbolNotFound: return None - else: - c = self.conn.cursor() - fuzzy = symbol.replace(self.fuzzy_char, '') - t = (fuzzy, as_of_date.value, as_of_date.value) - query = ("SELECT sid FROM equities " - "WHERE fuzzy=? " + - "AND start_date<=? " + - "AND end_date>=?") - c.execute(query, t) - candidates = c.fetchall() - # If one SID exists for symbol, return that symbol - if len(candidates) == 1: - return self._retrieve_equity(candidates[0][0]) + equities_cols = self.equities.c + candidates = sa.select((equities_cols.sid,)).where( + (equities_cols.fuzzy == fuzzy) & + (equities_cols.start_date <= ad_value) & + (equities_cols.end_date >= ad_value), + ).execute().fetchall() - # If multiple SIDs exist for symbol, return latest start_date with - # end_date as a tie-breaker - if len(candidates) > 1: - t = (symbol, as_of_date.value) - query = ("SELECT sid FROM equities " - "WHERE symbol=? " + - "AND start_date<=? " + - "ORDER BY start_date desc, end_date desc" + - "LIMIT 1") - c.execute(query, t) - data = c.fetchone() - if data: - return self._retrieve_equity(data[0]) + # If one SID exists for symbol, return that symbol + if len(candidates) == 1: + return self._retrieve_equity(candidates[0]['sid']) + + # If multiple SIDs exist for symbol, return latest start_date with + # end_date as a tie-breaker + elif candidates: + sid = sa.select((equities_cols.sid,)).where( + (equities_cols.symbol == symbol) & + (equities_cols.start_date <= ad_value), + ).order_by( + equities_cols.start_date.desc(), + equities_cols.end_date.desc(), + ).scalar().execute() + if sid: + return self._retrieve_equity(sid) + + raise SymbolNotFound(symbol=symbol) def lookup_future_chain(self, root_symbol, as_of_date, knowledge_date): """ Return the futures chain for a given root symbol. @@ -359,55 +376,54 @@ class AssetFinder(object): Raised when a future chain could not be found for the given root symbol. """ - c = self.conn.cursor() + + fc_cols = self.futures_contracts.c + if as_of_date is pd.NaT: # If the as_of_date is NaT, get all contracts for this # root symbol. - t = {'root_symbol': root_symbol} - c.execute(""" - select sid from futures - where root_symbol=:root_symbol - order by notice_date asc - """, t) + sids = list(map( + itemgetter('sid'), + sa.select((fc_cols.sid,)).where( + (fc_cols.root_symbol == root_symbol), + ).order_by( + fc_cols.notice_date.asc(), + ).execute().fetchall())) else: + as_of_date = as_of_date.value if knowledge_date is pd.NaT: # If knowledge_date is NaT, default to using as_of_date - t = {'root_symbol': root_symbol, - 'as_of_date': as_of_date.value, - 'knowledge_date': as_of_date.value} + knowledge_date = as_of_date.value else: - t = {'root_symbol': root_symbol, - 'as_of_date': as_of_date.value, - 'knowledge_date': knowledge_date.value} + knowledge_date = knowledge_date.value + + sids = list(map( + itemgetter('sid'), + sa.select((fc_cols.sid,)).where( + (fc_cols.root_symbol == root_symbol) & + (fc_cols.notice_date >= as_of_date) & + (fc_cols.start_date <= knowledge_date), + ).order_by( + fc_cols.notice_date.asc(), + ).execute().fetchall() + )) - c.execute(""" - select sid from futures - where root_symbol=:root_symbol - and :as_of_date < notice_date - and start_date <= :knowledge_date - order by notice_date asc - """, t) - sids = [r[0] for r in c.fetchall()] if not sids: # Check if root symbol exists. - c.execute(""" - SELECT COUNT(sid) FROM futures_contracts - WHERE root_symbol=:root_symbol - """, t) - count = c.fetchone()[0] + count = sa.select((sa.func.count(fc_cols.sid),)).where( + fc_cols.root_symbol == root_symbol, + ).scalar().execute() if count == 0: raise RootSymbolNotFound(root_symbol=root_symbol) - else: - # If symbol exists, return empty future chain. - return [] - return [self._retrieve_futures_contract(sid) for sid in sids] + + return map(self._retrieve_futures_contract, sids) @property def sids(self): - c = self.conn.cursor() - query = 'SELECT sid FROM asset_router' - c.execute(query) - return [r[0] for r in c.fetchall()] + return tuple(map( + itemgetter('sid'), + sa.select(self.asset_router.c.sid,).execute().fetchall(), + )) def _lookup_generic_scalar(self, asset_convertible, @@ -550,33 +566,38 @@ class AssetFinder(object): def _compute_asset_lifetimes(self): """ Compute and cache a recarry of asset lifetimes. - - FUTURE OPTIMIZATION: We're looping over a big array, which means this - probably should be in C/Cython. """ - with self.conn as transaction: - results = transaction.execute( - 'SELECT sid, start_date, end_date from equities' - ).fetchall() - - lifetimes = np.recarray( - shape=(len(results),), - dtype=[('sid', 'i8'), ('start', 'i8'), ('end', 'i8')], - ) - - # TODO: This is **WAY** slower than it could be because we have to - # check for None everywhere. If we represented "no start date" as - # 0, and "no end date" as MAX_INT in our metadata, this would be - # significantly faster. - NO_START = 0 - NO_END = np.iinfo(int).max - for idx, (sid, start, end) in enumerate(results): - lifetimes[idx] = ( - sid, - start if start is not None else NO_START, - end if end is not None else NO_END, - ) - return lifetimes + equities_cols = self.equities.c + buf = np.array( + tuple(map( + float, + sa.select(( + equities_cols.sid, + equities_cols.start_date, + equities_cols.end_date, + )).execute(), + )), + dype='