def test_sd():
sd = StructuredDataset(dataframe="hi")
sd.uri = "my uri"
assert sd.file_format == PARQUET
with pytest.raises(ValueError, match="No dataframe type set"):
sd.all()
with pytest.raises(ValueError, match="No dataframe type set."):
sd.iter()
class MockPandasDecodingHandlers(StructuredDatasetDecoder):
def decode(
self,
ctx: FlyteContext,
flyte_value: literals.StructuredDataset,
current_task_metadata: StructuredDatasetMetadata,
) -> typing.Union[typing.Generator[pd.DataFrame, None, None]]:
yield pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
StructuredDatasetTransformerEngine.register(MockPandasDecodingHandlers(
pd.DataFrame, "tmpfs"),
default_for_type=False)
sd = StructuredDataset()
sd._literal_sd = literals.StructuredDataset(
uri="tmpfs://somewhere",
metadata=StructuredDatasetMetadata(StructuredDatasetType(format="")))
assert isinstance(sd.open(pd.DataFrame).iter(), typing.Generator)
with pytest.raises(ValueError):
sd.open(pd.DataFrame).all()
class MockPandasDecodingHandlers(StructuredDatasetDecoder):
def decode(
self,
ctx: FlyteContext,
flyte_value: literals.StructuredDataset,
current_task_metadata: StructuredDatasetMetadata,
) -> pd.DataFrame:
return pd.DataFrame({"Name": ["Tom", "Joseph"], "Age": [20, 22]})
StructuredDatasetTransformerEngine.register(MockPandasDecodingHandlers(
pd.DataFrame, "tmpfs"),
default_for_type=False,
override=True)
sd = StructuredDataset()
sd._literal_sd = literals.StructuredDataset(
uri="tmpfs://somewhere",
metadata=StructuredDatasetMetadata(StructuredDatasetType(format="")))
with pytest.raises(ValueError):
sd.open(pd.DataFrame).iter()
def show_sd(in_sd: StructuredDataset):
pd.set_option("expand_frame_repr", False)
df = in_sd.open(pd.DataFrame).all()
print(df)