Skip to content

Commit fca00f5

Browse files
jmahlikdanabens
authored andcommitted
fix(tracker): log_metrics dataframe argument check allows df to be passed
1 parent 9df55c0 commit fca00f5

File tree

2 files changed

+22
-1
lines changed

2 files changed

+22
-1
lines changed

src/smexperiments/tracker.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,7 @@ def log_table(self, title=None, values=None, data_frame=None, output_artifact=Tr
450450
'Table values should be list. i.e. {"x": [1,2,3]}, instead was ' + str(type(values[key]))
451451
)
452452

453-
if data_frame:
453+
if data_frame is not None:
454454
values = _ArtifactConverter.convert_data_frame_to_values(data_frame)
455455
fields = _ArtifactConverter.convert_data_frame_to_fields(data_frame)
456456
else:

tests/unit/test_tracker.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -588,6 +588,27 @@ def test_log_table(under_test):
588588
under_test._lineage_artifact_tracker.add_input_artifact("TestTable", "s3uri_value", "etag_value", "Table")
589589

590590

591+
def test_log_table_dataframe(under_test):
592+
593+
dataframe = pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})
594+
595+
under_test._artifact_uploader.upload_object_artifact.return_value = ("s3uri_value", "etag_value")
596+
597+
under_test.log_table(title="TestTable", data_frame=dataframe)
598+
expected_data = {
599+
"type": "Table",
600+
"version": 0,
601+
"title": "TestTable",
602+
"fields": [{"name": "x", "type": "number"}, {"name": "y", "type": "number"}],
603+
"data": {"x": [1, 2, 3], "y": [4, 5, 6]},
604+
}
605+
under_test._artifact_uploader.upload_object_artifact.assert_called_with(
606+
"TestTable", expected_data, file_extension="json"
607+
)
608+
609+
under_test._lineage_artifact_tracker.add_input_artifact("TestTable", "s3uri_value", "etag_value", "Table")
610+
611+
591612
def test_log_roc_curve(under_test):
592613
y_true = [0, 0, 1, 1]
593614
y_scores = [0.1, 0.4, 0.35, 0.8]

0 commit comments

Comments
 (0)