fix: Fix import array via bulkwriter (#2035)

When import via bulkwriter, for array type, we should specify datatype; otherwise, it will default to double(for float32 array) and int64(for int8, int16, int32 array). issue: #2034, milvus-io/milvus#31834 Signed-off-by: bigsheeper <[email protected]>
milvus-io · Apr 12, 2024 · 989d15b · 989d15b
1 parent e0f0e91
commit 989d15b
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pymilvus/bulk_writer/buffer.py b/pymilvus/bulk_writer/buffer.py
@@ -212,6 +212,12 @@ def _persist_parquet(self, local_path: str, **kwargs):
                 for val in self._buffer[k]:
                     arr.append(np.array(val, dtype=np.dtype("uint8")))
                 data[k] = pd.Series(arr)
+            elif field_schema.dtype == DataType.ARRAY:
+                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)