Merge remote-tracking branch 'mlrun/development' into ML-1325

mlrun · Jul 12, 2023 · eb0c94b · eb0c94b
2 parents fa9adff + 5f4b4ad
commit eb0c94b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/storey/targets.py b/storey/targets.py
@@ -140,8 +140,16 @@ def unzip_cols(columns):
 
     _type_string_to_pyarrow_type = {
         "str": pyarrow.string(),
+        "int8": pyarrow.int8(),
+        "int16": pyarrow.int16(),
         "int32": pyarrow.int32(),
+        "int64": pyarrow.int64(),
+        "uint8": pyarrow.uint8(),
+        "uint16": pyarrow.uint16(),
+        "uint32": pyarrow.uint32(),
+        "uint64": pyarrow.uint64(),
         "int": pyarrow.int64(),
+        "float16": pyarrow.float16(),
         "float32": pyarrow.float32(),
         "float": pyarrow.float64(),
         "bool": pyarrow.bool_(),

diff --git a/tests/test_flow.py b/tests/test_flow.py
@@ -2487,7 +2487,7 @@ def test_write_to_parquet_string_as_datetime(tmpdir):
     out_dir = f"{tmpdir}/test_write_to_parquet_string_to_datetime/{uuid.uuid4().hex}/"
     columns = ["my_int", "my_string", "my_datetime"]
     columns_with_type = [
-        ("my_int", "int"),
+        ("my_int", "int8"),  # ML-4162
         ("my_string", "str"),
         ("my_datetime", "datetime"),
     ]
@@ -2506,6 +2506,7 @@ def test_write_to_parquet_string_as_datetime(tmpdir):
         controller.emit([i, f"this is {i}", my_time.isoformat()])
         expected.append([i, f"this is {i}", my_time.isoformat(sep=" ")])
     expected_df = pd.DataFrame(expected, columns=columns)
+    expected_df["my_int"] = expected_df["my_int"].astype("int8")
     expected_df["my_datetime"] = expected_df["my_datetime"].astype("datetime64[us]")
     controller.terminate()
     controller.await_termination()