handle all type

Signed-off-by: bigsheeper <[email protected]>
milvus-io · Apr 12, 2024 · a9eb5d1 · a9eb5d1
1 parent 62b3ffb
commit a9eb5d1
Showing 1 changed file with 5 additions and 7 deletions.
diff --git a/pymilvus/bulk_writer/buffer.py b/pymilvus/bulk_writer/buffer.py
@@ -213,13 +213,11 @@ def _persist_parquet(self, local_path: str, **kwargs):
                     arr.append(np.array(val, dtype=np.dtype("uint8")))
                 data[k] = pd.Series(arr)
             elif field_schema.dtype == DataType.ARRAY:
-                if field_schema.element_type == DataType.FLOAT:
-                    arr = []
-                    for val in self._buffer[k]:
-                        arr.append(np.array(val, dtype=np.dtype("float32")))
-                    data[k] = pd.Series(arr)
-                else:
-                    data[k] = pd.Series(self._buffer[k])
+                dt = NUMPY_TYPE_CREATOR[field_schema.element_type.name]
+                arr = []
+                for val in self._buffer[k]:
+                    arr.append(np.array(val, dtype=dt))
+                data[k] = pd.Series(arr)
             elif field_schema.dtype.name in NUMPY_TYPE_CREATOR:
                 dt = NUMPY_TYPE_CREATOR[field_schema.dtype.name]
                 data[k] = pd.Series(self._buffer[k], dtype=dt)