def load_batch_data(self, batch_id: str, batch_data: Any) -> None:
if isinstance(batch_data, pd.DataFrame):
batch_data = PandasBatchData(self, batch_data)
elif isinstance(batch_data, PandasBatchData):
pass
else:
raise ge_exceptions.GreatExpectationsError(
"PandasExecutionEngine requires batch data that is either a DataFrame or a PandasBatchData object"
)
super().load_batch_data(batch_id=batch_id, batch_data=batch_data)
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"
)
}
)
batch_data: PandasBatchData
if isinstance(batch_spec, RuntimeDataBatchSpec):
# batch_data != None is already checked when RuntimeDataBatchSpec is instantiated
if isinstance(batch_spec.batch_data, pd.DataFrame):
df = batch_spec.batch_data
elif isinstance(batch_spec.batch_data, PandasBatchData):
df = batch_spec.batch_data.dataframe
else:
raise ValueError(
"RuntimeDataBatchSpec must provide a Pandas DataFrame or PandasBatchData object."
)
batch_spec.batch_data = "PandasDataFrame"
elif isinstance(batch_spec, S3BatchSpec):
if self._s3 is None:
raise ge_exceptions.ExecutionEngineError(
f"""PandasExecutionEngine has been passed a S3BatchSpec,
but the ExecutionEngine does not have a boto3 client configured. Please check your config."""
)
s3_engine = self._s3
s3_url = S3Url(batch_spec.path)
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options or {}
if "compression" not in reader_options.keys():
reader_options["compression"] = sniff_s3_compression(s3_url)
s3_object = s3_engine.get_object(Bucket=s3_url.bucket, Key=s3_url.key)
logger.debug(
"Fetching s3 object. Bucket: {} Key: {}".format(
s3_url.bucket, s3_url.key
)
)
reader_fn = self._get_reader_fn(reader_method, s3_url.key)
buf = BytesIO(s3_object["Body"].read())
buf.seek(0)
df = reader_fn(buf, **reader_options)
elif isinstance(batch_spec, PathBatchSpec):
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options
path: str = batch_spec.path
reader_fn: Callable = self._get_reader_fn(reader_method, path)
df = reader_fn(path, **reader_options)
else:
raise ge_exceptions.BatchSpecError(
f"batch_spec must be of type RuntimeDataBatchSpec, PathBatchSpec, or S3BatchSpec, not {batch_spec.__class__.__name__}"
)
df = self._apply_splitting_and_sampling_methods(batch_spec, df)
if df.memory_usage().sum() < HASH_THRESHOLD:
batch_markers["pandas_data_fingerprint"] = hash_pandas_dataframe(df)
typed_batch_data = PandasBatchData(execution_engine=self, dataframe=df)
return typed_batch_data, batch_markers
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")
})
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"""PandasExecutionEngine has been passed a string type batch_data, "{batch_data}", which is illegal.
Please check your config.""")
if isinstance(batch_spec.batch_data, pd.DataFrame):
df = batch_spec.batch_data
elif isinstance(batch_spec.batch_data, PandasBatchData):
df = batch_spec.batch_data.dataframe
else:
raise ValueError(
"RuntimeDataBatchSpec must provide a Pandas DataFrame or PandasBatchData object."
)
batch_spec.batch_data = "PandasDataFrame"
elif isinstance(batch_spec, S3BatchSpec):
if self._s3 is None:
self._instantiate_s3_client()
# if we were not able to instantiate S3 client, then raise error
if self._s3 is None:
raise ge_exceptions.ExecutionEngineError(
"""PandasExecutionEngine has been passed a S3BatchSpec,
but the ExecutionEngine does not have a boto3 client configured. Please check your config."""
)
s3_engine = self._s3
try:
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options or {}
path: str = batch_spec.path
s3_url = S3Url(path)
if "compression" not in reader_options.keys():
inferred_compression_param = sniff_s3_compression(s3_url)
if inferred_compression_param is not None:
reader_options[
"compression"] = inferred_compression_param
s3_object = s3_engine.get_object(Bucket=s3_url.bucket,
Key=s3_url.key)
except (ParamValidationError, ClientError) as error:
raise ge_exceptions.ExecutionEngineError(
f"""PandasExecutionEngine encountered the following error while trying to read data from S3 Bucket: {error}"""
)
logger.debug(
f"Fetching s3 object. Bucket: {s3_url.bucket} Key: {s3_url.key}"
)
reader_fn = self._get_reader_fn(reader_method, s3_url.key)
buf = BytesIO(s3_object["Body"].read())
buf.seek(0)
df = reader_fn(buf, **reader_options)
elif isinstance(batch_spec, AzureBatchSpec):
if self._azure is None:
self._instantiate_azure_client()
# if we were not able to instantiate Azure client, then raise error
if self._azure is None:
raise ge_exceptions.ExecutionEngineError(
"""PandasExecutionEngine has been passed a AzureBatchSpec,
but the ExecutionEngine does not have an Azure client configured. Please check your config."""
)
azure_engine = self._azure
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options or {}
path: str = batch_spec.path
azure_url = AzureUrl(path)
blob_client = azure_engine.get_blob_client(
container=azure_url.container, blob=azure_url.blob)
azure_object = blob_client.download_blob()
logger.debug(
f"Fetching Azure blob. Container: {azure_url.container} Blob: {azure_url.blob}"
)
reader_fn = self._get_reader_fn(reader_method, azure_url.blob)
buf = BytesIO(azure_object.readall())
buf.seek(0)
df = reader_fn(buf, **reader_options)
elif isinstance(batch_spec, GCSBatchSpec):
if self._gcs is None:
self._instantiate_gcs_client()
# if we were not able to instantiate GCS client, then raise error
if self._gcs is None:
raise ge_exceptions.ExecutionEngineError(
"""PandasExecutionEngine has been passed a GCSBatchSpec,
but the ExecutionEngine does not have an GCS client configured. Please check your config."""
)
gcs_engine = self._gcs
gcs_url = GCSUrl(batch_spec.path)
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options or {}
try:
gcs_bucket = gcs_engine.get_bucket(gcs_url.bucket)
gcs_blob = gcs_bucket.blob(gcs_url.blob)
logger.debug(
f"Fetching GCS blob. Bucket: {gcs_url.bucket} Blob: {gcs_url.blob}"
)
except GoogleAPIError as error:
raise ge_exceptions.ExecutionEngineError(
f"""PandasExecutionEngine encountered the following error while trying to read data from GCS Bucket: {error}"""
)
reader_fn = self._get_reader_fn(reader_method, gcs_url.blob)
buf = BytesIO(gcs_blob.download_as_bytes())
buf.seek(0)
df = reader_fn(buf, **reader_options)
elif isinstance(batch_spec, PathBatchSpec):
reader_method: str = batch_spec.reader_method
reader_options: dict = batch_spec.reader_options
path: str = batch_spec.path
reader_fn: Callable = self._get_reader_fn(reader_method, path)
df = reader_fn(path, **reader_options)
else:
raise ge_exceptions.BatchSpecError(
f"batch_spec must be of type RuntimeDataBatchSpec, PathBatchSpec, S3BatchSpec, or AzureBatchSpec, not {batch_spec.__class__.__name__}"
)
df = self._apply_splitting_and_sampling_methods(batch_spec, df)
if df.memory_usage().sum() < HASH_THRESHOLD:
batch_markers["pandas_data_fingerprint"] = hash_pandas_dataframe(
df)
typed_batch_data = PandasBatchData(execution_engine=self, dataframe=df)
return typed_batch_data, batch_markers
