diff --git a/zipline/pipeline/loaders/blaze/core.py b/zipline/pipeline/loaders/blaze/core.py index 5264a954..61782c31 100644 --- a/zipline/pipeline/loaders/blaze/core.py +++ b/zipline/pipeline/loaders/blaze/core.py @@ -1028,22 +1028,10 @@ class BlazeLoader(dict): return odo(e[predicate][colnames], pd.DataFrame, **odo_kwargs) - if checkpoints is not None: - ts = checkpoints[TS_FIELD_NAME] - checkpoints_ts = odo(ts[ts <= lower_dt].max(), pd.Timestamp) - if pd.isnull(checkpoints_ts): - materialized_checkpoints = pd.DataFrame(columns=colnames) - lower = None - else: - materialized_checkpoints = odo( - checkpoints[ts == checkpoints_ts][colnames], - pd.DataFrame, - **odo_kwargs - ) - lower = checkpoints_ts - else: - materialized_checkpoints = pd.DataFrame(columns=colnames) - lower = None + + lower, materialized_checkpoints = get_materialized_checkpoints( + checkpoints, colnames, lower_dt, odo_kwargs + ) materialized_expr = self.pool.apply_async(collect_expr, (expr, lower)) materialized_deltas = ( @@ -1137,6 +1125,7 @@ class BlazeLoader(dict): for column_idx, column in enumerate(columns) } + global_loader = BlazeLoader.global_instance() @@ -1168,12 +1157,32 @@ def bind_expression_to_resources(expr, resources): }) +def get_materialized_checkpoints(checkpoints, colnames, lower_dt, odo_kwargs): + if checkpoints is not None: + ts = checkpoints[TS_FIELD_NAME] + checkpoints_ts = odo(ts[ts <= lower_dt].max(), pd.Timestamp) + if pd.isnull(checkpoints_ts): + materialized_checkpoints = pd.DataFrame(columns=colnames) + lower = None + else: + materialized_checkpoints = odo( + checkpoints[ts == checkpoints_ts][colnames], + pd.DataFrame, + **odo_kwargs + ) + lower = checkpoints_ts + else: + materialized_checkpoints = pd.DataFrame(columns=colnames) + lower = None + return lower, materialized_checkpoints + + def ffill_query_in_range(expr, lower, upper, + checkpoints=None, odo_kwargs=None, - ts_field=TS_FIELD_NAME, - sid_field=SID_FIELD_NAME): + ts_field=TS_FIELD_NAME): """Query a blaze expression in a given time range properly forward filling from values that fall before the lower date. @@ -1199,27 +1208,24 @@ def ffill_query_in_range(expr, start before the requested start date if a value is needed to ffill. """ odo_kwargs = odo_kwargs or {} - filtered = expr[expr[ts_field] <= lower] - computed_lower = odo( - bz.by( - filtered[sid_field], - timestamp=filtered[ts_field].max(), - ).timestamp.min(), - pd.Timestamp, - **odo_kwargs + computed_lower, materialized_checkpoints = get_materialized_checkpoints( + checkpoints, expr.fields, lower, odo_kwargs ) if pd.isnull(computed_lower): # If there is no lower date, just query for data in the date # range. It must all be null anyways. computed_lower = lower - raw = odo( - expr[ - (expr[ts_field] >= computed_lower) & - (expr[ts_field] <= upper) - ], - pd.DataFrame, - **odo_kwargs + raw = pd.concat( + [materialized_checkpoints, + odo( + expr[ + (expr[ts_field] >= computed_lower) & + (expr[ts_field] <= upper) + ], + pd.DataFrame, + **odo_kwargs + )] ) raw.loc[:, ts_field] = raw.loc[:, ts_field].astype('datetime64[ns]') return raw diff --git a/zipline/pipeline/loaders/blaze/estimates.py b/zipline/pipeline/loaders/blaze/estimates.py index 864b05a7..b991c62c 100644 --- a/zipline/pipeline/loaders/blaze/estimates.py +++ b/zipline/pipeline/loaders/blaze/estimates.py @@ -63,7 +63,8 @@ class BlazeEstimatesLoader(PipelineLoader): resources=None, odo_kwargs=None, data_query_time=None, - data_query_tz=None): + data_query_tz=None, + checkpoints=None): dshape = expr.dshape if not istabular(dshape): @@ -83,6 +84,7 @@ class BlazeEstimatesLoader(PipelineLoader): check_data_query_args(data_query_time, data_query_tz) self._data_query_time = data_query_time self._data_query_tz = data_query_tz + self._checkpoints = checkpoints def load_adjusted_array(self, columns, dates, assets, mask): raw = load_raw_data(assets, @@ -90,7 +92,8 @@ class BlazeEstimatesLoader(PipelineLoader): self._data_query_time, self._data_query_tz, self._expr, - self._odo_kwargs) + self._odo_kwargs, + checkpoints=self._checkpoints) return self.loader( raw, diff --git a/zipline/pipeline/loaders/blaze/utils.py b/zipline/pipeline/loaders/blaze/utils.py index 963d9f9a..b5be9cd9 100644 --- a/zipline/pipeline/loaders/blaze/utils.py +++ b/zipline/pipeline/loaders/blaze/utils.py @@ -11,7 +11,8 @@ def load_raw_data(assets, data_query_time, data_query_tz, expr, - odo_kwargs): + odo_kwargs, + checkpoints=None): """ given an expression representing data to load, perform normalization and forward-filling and return the data, materialized. @@ -48,7 +49,8 @@ def load_raw_data(assets, expr, lower_dt, upper_dt, - odo_kwargs, + checkpoints=checkpoints, + odo_kwargs=odo_kwargs, ) sids = raw[SID_FIELD_NAME] raw.drop( diff --git a/zipline/pipeline/loaders/quarter_estimates.py b/zipline/pipeline/loaders/quarter_estimates.py index 54686a6e..71ff3a32 100644 --- a/zipline/pipeline/loaders/quarter_estimates.py +++ b/zipline/pipeline/loaders/quarter_estimates.py @@ -1,5 +1,7 @@ +from collections import defaultdict from abc import abstractmethod import numpy as np +import pandas as pd from six import viewvalues from toolz import groupby @@ -152,30 +154,22 @@ class QuarterEstimatesLoader(PipelineLoader): return requested_qtr_data def get_adjustments(self, - zero_qtr_idx, - requested_qtr_idx, - stacked_last_per_qtr, + zero_qtr_data, + requested_qtr_data, last_per_qtr, dates, - column_name, - column, - mask, - assets): + assets, + columns): """ Creates an AdjustedArray from the given estimates data for the given dates. Parameters ---------- - zero_qtr_idx : pd.MultiIndex - The index of the row of the zeroth (immediately next/previous) - quarter from each date for each sid. - requested_qtr_idx : pd.MultiIndex - The index of the row of the requested quarter from each date for - each sid. - stacked_last_per_qtr : pd.DataFrame - The latest estimate known per sid per date per quarter with the - dates, normalized quarter, and sid as the index. + zero_qtr_data : pd.DataFrame + The 'time zero' data for each date/sid. + zero_qtr_data : pd.DataFrame + The data for the requested quarter. last_per_qtr : pd.DataFrame The latest estimate known per sid per date per quarter with dates as the index and normalized quarter and sid in the columns @@ -198,18 +192,11 @@ class QuarterEstimatesLoader(PipelineLoader): adjusted_array : AdjustedArray The array of data and overwrites for the given column. """ - adjustments = {} - requested_qtr_data = self.get_requested_data_for_col( - stacked_last_per_qtr, requested_qtr_idx, dates - ) - zero_qtr_data = stacked_last_per_qtr.loc[zero_qtr_idx] + col_to_adjustments = defaultdict(dict) # We no longer need this in the index, but we do need it as a column # to calculate adjustments. zero_qtr_data = zero_qtr_data.reset_index(NORMALIZED_QUARTERS) - if column.dtype == datetime64ns_dtype: - overwrite = Datetime641DArrayOverwrite - else: - overwrite = Float641DArrayOverwrite + for sid_idx, sid in enumerate(assets): zero_qtr_sid_data = zero_qtr_data[ zero_qtr_data.index.get_level_values(SID_FIELD_NAME) == sid @@ -243,38 +230,35 @@ class QuarterEstimatesLoader(PipelineLoader): if isinstance(self, PreviousQuartersEstimatesLoader) else 'right' ) - adjustments[next_qtr_start_idx] = \ - self.create_overwrite_for_quarter( - next_qtr_start_idx, - column, - column_name, - dates, - last_per_qtr, - overwrite, - qtrs_with_estimates_for_sid, - requested_qtr_data, - sid, - sid_idx, - ) - - return AdjustedArray( - requested_qtr_data[column_name].values.astype(column.dtype), - mask, - dict(adjustments), - column.missing_value, - ) + self.create_overwrite_for_quarter( + col_to_adjustments, + next_qtr_start_idx, + dates, + last_per_qtr, + qtrs_with_estimates_for_sid, + requested_qtr_data, + sid, + sid_idx, + columns, + ) + return col_to_adjustments def create_overwrite_for_quarter(self, + col_to_adjustments, next_qtr_start_idx, - column, - column_name, dates, last_per_qtr, - overwrite, quarters_with_estimates_for_sid, requested_qtr_data, sid, - sid_idx): + sid_idx, + columns): + overwrites_dict = {} + for col in columns: + if col.dtype == datetime64ns_dtype: + overwrites_dict[col] = Datetime641DArrayOverwrite + else: + overwrites_dict[col] = Float641DArrayOverwrite # Only add adjustments if the next quarter starts somewhere in # our date index for this sid. Our 'next' quarter can never # start at index 0; a starting index of 0 means that the next @@ -285,32 +269,35 @@ class QuarterEstimatesLoader(PipelineLoader): requested_quarter = requested_qtr_data[ SHIFTED_NORMALIZED_QTRS ][sid].iloc[next_qtr_start_idx] - - # If there are estimates for the requested quarter, - # overwrite all values going up to the starting index of - # that quarter with estimates for that quarter. - if requested_quarter in quarters_with_estimates_for_sid: - return self.create_overwrite_for_estimate( - column, - column_name, - last_per_qtr, - next_qtr_start_idx, - overwrite, - requested_quarter, - sid, - sid_idx - ) - # There are no estimates for the quarter. Overwrite all - # values going up to the starting index of that quarter - # with the missing value for this column. - else: - return self.overwrite_with_null( - column, - last_per_qtr, - next_qtr_start_idx, - overwrite, - sid_idx - ) + for col in columns: + column_name = self.name_map[col.name] + # If there are estimates for the requested quarter, + # overwrite all values going up to the starting index of + # that quarter with estimates for that quarter. + if requested_quarter in quarters_with_estimates_for_sid: + col_to_adjustments[column_name][next_qtr_start_idx] = \ + self.create_overwrite_for_estimate( + col, + column_name, + last_per_qtr, + next_qtr_start_idx, + overwrites_dict[col], + requested_quarter, + sid, + sid_idx + ) + # There are no estimates for the quarter. Overwrite all + # values going up to the starting index of that quarter + # with the missing value for this column. + else: + col_to_adjustments[column_name][next_qtr_start_idx] =\ + self.overwrite_with_null( + col, + last_per_qtr, + next_qtr_start_idx, + overwrites_dict[col], + sid_idx + ) def overwrite_with_null(self, column, @@ -353,7 +340,6 @@ class QuarterEstimatesLoader(PipelineLoader): ) out = {} - for num_quarters, columns in groups.items(): # Determine the last piece of information we know for each column # on each date in the index for each sid and quarter. @@ -388,19 +374,33 @@ class QuarterEstimatesLoader(PipelineLoader): ), shifted_qtr_data[SHIFTED_NORMALIZED_QTRS] ]).index + requested_qtr_data = self.get_requested_data_for_col( + stacked_last_per_qtr, requested_qtr_idx, dates + ) - for c in columns: - column_name = self.name_map[c.name] - adjusted_array = self.get_adjustments(zero_qtr_idx, - requested_qtr_idx, - stacked_last_per_qtr, + zero_qtr_data = stacked_last_per_qtr.loc[zero_qtr_idx] + + col_to_adjustments = self.get_adjustments(zero_qtr_data, + requested_qtr_data, last_per_qtr, dates, - column_name, - c, - mask, - assets) - out[c] = adjusted_array + assets, + columns) + for col in columns: + column_name = self.name_map[col.name] + # We may have dropped assets if they never have any data for the + # requested quarter. + df = pd.DataFrame(data=requested_qtr_data[column_name], + index=dates, + columns=assets, + dtype=col.dtype) + + out[col] = AdjustedArray( + df.values.astype(col.dtype), + mask, + dict(col_to_adjustments[column_name]), + col.missing_value, + ) return out