mirror of
https://github.com/wassname/catalyst.git
synced 2026-07-25 13:10:33 +08:00
MAINT: add checkpointing
BUG: add cols for sids with no data and get adjustments outside column loop
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user