MAINT: add checkpointing

BUG: add cols for sids with no data and get adjustments outside column loop
This commit is contained in:
Maya Tydykov
2016-09-27 09:54:43 -04:00
parent 90984be20c
commit f140cdb8d6
4 changed files with 134 additions and 123 deletions
+39 -33
View File
@@ -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
+5 -2
View File
@@ -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,
+4 -2
View File
@@ -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(
+86 -86
View File
@@ -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