def _get_reader_fn(self, reader, reader_method=None, path=None):
"""Static helper for providing reader_fn
Args:
reader: the base spark reader to use; this should have had reader_options applied already
reader_method: the name of the reader_method to use, if specified
path (str): the path to use to guess reader_method if it was not specified
Returns:
ReaderMethod to use for the filepath
"""
if reader_method is None and path is None:
raise ExecutionEngineError(
"Unable to determine spark reader function without reader_method or path"
)
if reader_method is None:
reader_method = self.guess_reader_method_from_path(path=path)
reader_method_op: str = reader_method.lower()
try:
if reader_method_op == "delta":
return reader.format(reader_method_op).load
return getattr(reader, reader_method_op)
except AttributeError:
raise ExecutionEngineError(
f"Unable to find reader_method {reader_method} in spark.",
)
def active_batch_data_id(self, batch_id):
if batch_id in self.loaded_batch_data_dict.keys():
self._active_batch_data_id = batch_id
else:
raise ExecutionEngineError(
f"Unable to set active_batch_data_id to {batch_id}. The may data may not be loaded."
)
def guess_reader_method_from_path(path):
"""Based on a given filepath, decides a reader method. Currently supports tsv, csv, and parquet. If none of these
file extensions are used, returns ExecutionEngineError stating that it is unable to determine the current path.
Args:
path - A given file path
Returns:
A dictionary entry of format {'reader_method': reader_method}
"""
if path.endswith(".csv") or path.endswith(".tsv"):
return "csv"
elif path.endswith(".parquet"):
return "parquet"
raise ExecutionEngineError(
f"Unable to determine reader method from path: {path}")
def resolve_metric_bundle(
self,
metric_fn_bundle: Iterable[Tuple[MetricConfiguration, Any, dict, dict]],
) -> Dict[Tuple[str, str, str], Any]:
"""For every metric in a set of Metrics to resolve, obtains necessary metric keyword arguments and builds
bundles of the metrics into one large query dictionary so that they are all executed simultaneously. Will fail
if bundling the metrics together is not possible.
Args:
metric_fn_bundle (Iterable[Tuple[MetricConfiguration, Callable, dict]): \
A Dictionary containing a MetricProvider's MetricConfiguration (its unique identifier), its metric provider function
(the function that actually executes the metric), and the arguments to pass to the metric provider function.
A dictionary of metrics defined in the registry and corresponding arguments
Returns:
A dictionary of metric names and their corresponding now-queried values.
"""
resolved_metrics = {}
# We need a different query for each domain (where clause).
queries: Dict[Tuple, dict] = {}
for (
metric_to_resolve,
engine_fn,
compute_domain_kwargs,
accessor_domain_kwargs,
metric_provider_kwargs,
) in metric_fn_bundle:
if not isinstance(compute_domain_kwargs, IDDict):
compute_domain_kwargs = IDDict(compute_domain_kwargs)
domain_id = compute_domain_kwargs.to_id()
if domain_id not in queries:
queries[domain_id] = {
"select": [],
"ids": [],
"domain_kwargs": compute_domain_kwargs,
}
if self.engine.dialect.name == "clickhouse":
queries[domain_id]["select"].append(
engine_fn.label(
metric_to_resolve.metric_name.join(
random.choices(string.ascii_lowercase, k=2)
)
)
)
else:
queries[domain_id]["select"].append(
engine_fn.label(metric_to_resolve.metric_name)
)
queries[domain_id]["ids"].append(metric_to_resolve.id)
for query in queries.values():
domain_kwargs = query["domain_kwargs"]
selectable = self.get_domain_records(
domain_kwargs=domain_kwargs,
)
assert len(query["select"]) == len(query["ids"])
try:
"""
If a custom query is passed, selectable will be TextClause and not formatted
as a subquery wrapped in "(subquery) alias". TextClause must first be converted
to TextualSelect using sa.columns() before it can be converted to type Subquery
"""
if TextClause and isinstance(selectable, TextClause):
res = self.engine.execute(
sa.select(query["select"]).select_from(
selectable.columns().subquery()
)
).fetchall()
else:
res = self.engine.execute(
sa.select(query["select"]).select_from(selectable)
).fetchall()
logger.debug(
f"SqlAlchemyExecutionEngine computed {len(res[0])} metrics on domain_id {IDDict(domain_kwargs).to_id()}"
)
except OperationalError as oe:
exception_message: str = "An SQL execution Exception occurred. "
exception_traceback: str = traceback.format_exc()
exception_message += f'{type(oe).__name__}: "{str(oe)}". Traceback: "{exception_traceback}".'
logger.error(exception_message)
raise ExecutionEngineError(message=exception_message)
assert (
len(res) == 1
), "all bundle-computed metrics must be single-value statistics"
assert len(query["ids"]) == len(
res[0]
), "unexpected number of metrics returned"
for idx, id in enumerate(query["ids"]):
resolved_metrics[id] = convert_to_json_serializable(res[0][idx])
return resolved_metrics
def resolve_metric_bundle(
self,
metric_fn_bundle: Iterable[Tuple[MetricConfiguration, Any, dict,
dict]],
) -> dict:
"""For every metric in a set of Metrics to resolve, obtains necessary metric keyword arguments and builds
bundles of the metrics into one large query dictionary so that they are all executed simultaneously. Will fail
if bundling the metrics together is not possible.
Args:
metric_fn_bundle (Iterable[Tuple[MetricConfiguration, Callable, dict]): \
A Dictionary containing a MetricProvider's MetricConfiguration (its unique identifier), its metric provider function
(the function that actually executes the metric), and the arguments to pass to the metric provider function.
metrics (Dict[Tuple, Any]): \
A dictionary of metrics defined in the registry and corresponding arguments
Returns:
A dictionary of metric names and their corresponding now-queried values.
"""
resolved_metrics = dict()
# We need a different query for each domain (where clause).
queries: Dict[Tuple, dict] = dict()
for (
metric_to_resolve,
engine_fn,
compute_domain_kwargs,
accessor_domain_kwargs,
metric_provider_kwargs,
) in metric_fn_bundle:
if not isinstance(compute_domain_kwargs, IDDict):
compute_domain_kwargs = IDDict(compute_domain_kwargs)
domain_id = compute_domain_kwargs.to_id()
if domain_id not in queries:
queries[domain_id] = {
"select": [],
"ids": [],
"domain_kwargs": compute_domain_kwargs,
}
queries[domain_id]["select"].append(
engine_fn.label(metric_to_resolve.metric_name))
queries[domain_id]["ids"].append(metric_to_resolve.id)
for query in queries.values():
selectable, compute_domain_kwargs, _ = self.get_compute_domain(
query["domain_kwargs"], domain_type="identity")
assert len(query["select"]) == len(query["ids"])
try:
res = self.engine.execute(
sa.select(
query["select"]).select_from(selectable)).fetchall()
logger.debug(
f"SqlAlchemyExecutionEngine computed {len(res[0])} metrics on domain_id {IDDict(compute_domain_kwargs).to_id()}"
)
except OperationalError as oe:
exception_message: str = "An SQL execution Exception occurred. "
exception_traceback: str = traceback.format_exc()
exception_message += f'{type(oe).__name__}: "{str(oe)}". Traceback: "{exception_traceback}".'
logger.error(exception_message)
raise ExecutionEngineError(message=exception_message)
assert (
len(res) == 1
), "all bundle-computed metrics must be single-value statistics"
assert len(query["ids"]) == len(
res[0]), "unexpected number of metrics returned"
for idx, id in enumerate(query["ids"]):
resolved_metrics[id] = convert_to_json_serializable(
res[0][idx])
return resolved_metrics
def get_batch_data_and_markers(
self, batch_spec: BatchSpec
) -> Tuple[Any, BatchMarkers]: # batch_data
# We need to build a batch_markers to be used in the dataframe
batch_markers: BatchMarkers = BatchMarkers(
{
"ge_load_time": datetime.datetime.now(datetime.timezone.utc).strftime(
"%Y%m%dT%H%M%S.%fZ"
)
}
)
"""
As documented in Azure DataConnector implementations, Pandas and Spark execution engines utilize separate path
formats for accessing Azure Blob Storage service. However, Pandas and Spark execution engines utilize identical
path formats for accessing all other supported cloud storage services (AWS S3 and Google Cloud Storage).
Moreover, these formats (encapsulated in S3BatchSpec and GCSBatchSpec) extend PathBatchSpec (common to them).
Therefore, at the present time, all cases with the exception of Azure Blob Storage , are handled generically.
"""
batch_data: Any
if isinstance(batch_spec, RuntimeDataBatchSpec):
# batch_data != None is already checked when RuntimeDataBatchSpec is instantiated
batch_data = batch_spec.batch_data
if isinstance(batch_data, str):
raise ge_exceptions.ExecutionEngineError(
f"""SparkDFExecutionEngine has been passed a string type batch_data, "{batch_data}", which is illegal.
Please check your config."""
)
batch_spec.batch_data = "SparkDataFrame"
elif isinstance(batch_spec, AzureBatchSpec):
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options or {}
path: str = batch_spec.path
azure_url = AzureUrl(path)
try:
credential = self._azure_options.get("credential")
storage_account_url = azure_url.account_url
if credential:
self.spark.conf.set(
"fs.wasb.impl",
"org.apache.hadoop.fs.azure.NativeAzureFileSystem",
)
self.spark.conf.set(
f"fs.azure.account.key.{storage_account_url}", credential
)
reader: DataFrameReader = self.spark.read.options(**reader_options)
reader_fn: Callable = self._get_reader_fn(
reader=reader,
reader_method=reader_method,
path=path,
)
batch_data = reader_fn(path)
except AttributeError:
raise ExecutionEngineError(
"""
Unable to load pyspark. Pyspark is required for SparkDFExecutionEngine.
"""
)
elif isinstance(batch_spec, PathBatchSpec):
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options or {}
path: str = batch_spec.path
try:
reader: DataFrameReader = self.spark.read.options(**reader_options)
reader_fn: Callable = self._get_reader_fn(
reader=reader,
reader_method=reader_method,
path=path,
)
batch_data = reader_fn(path)
except AttributeError:
raise ExecutionEngineError(
"""
Unable to load pyspark. Pyspark is required for SparkDFExecutionEngine.
"""
)
# pyspark will raise an AnalysisException error if path is incorrect
except pyspark.sql.utils.AnalysisException:
raise ExecutionEngineError(
f"""Unable to read in batch from the following path: {path}. Please check your configuration."""
)
else:
raise BatchSpecError(
"""
Invalid batch_spec: batch_data is required for a SparkDFExecutionEngine to operate.
"""
)
batch_data = self._apply_splitting_and_sampling_methods(batch_spec, batch_data)
typed_batch_data = SparkDFBatchData(execution_engine=self, dataframe=batch_data)
return typed_batch_data, batch_markers
