Fix a bug of bulkwriter to support all-empty struct list (#3192)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
Co-authored-by: XuanYang-cn <xuan.yang@zilliz.com>

Author: yhmo

examples/bulk_import/bulk_writer_all_types.py

@@ -2,6 +2,7 @@
 import pytz
 import time
 import numpy as np
+from pathlib import Path
 from typing import List
 
 from pymilvus import (
@@ -16,6 +17,9 @@
     get_import_progress,
 )
 
+LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
+
 # minio
 MINIO_ADDRESS = "0.0.0.0:9000"
 MINIO_SECRET_KEY = "minioadmin"
@@ -216,7 +220,7 @@ def remote_writer(schema: CollectionSchema, file_type: BulkFileType):
     with RemoteBulkWriter(
             schema=schema,
             remote_path="bulk_data",
-            local_path="/tmp/PARQUET",
+            local_path=LOCAL_FILES_PATH,
             connect_param=RemoteBulkWriter.S3ConnectParam(
                 endpoint=MINIO_ADDRESS,
                 access_key=MINIO_ACCESS_KEY,
@@ -269,7 +273,7 @@ def local_writer(schema: CollectionSchema, file_type: BulkFileType):
     print(f"\n===================== local writer ({file_type.name}) ====================")
     writer = LocalBulkWriter(
         schema=schema,
-        local_path="./" + file_type.name,
+        local_path=LOCAL_FILES_PATH,
         chunk_size=16 * 1024 * 1024,
         file_type=file_type
     )

examples/bulk_import/example_bulkinsert_json.py

@@ -6,6 +6,7 @@
 
 from minio import Minio
 from minio.error import S3Error
+from pathlib import Path
 
 from pymilvus import (
     DataType,
@@ -18,7 +19,8 @@
 )
 
 # Local path to generate JSON files
-LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert"
+LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
 
 # Milvus service address
 _HOST = '127.0.0.1'

examples/bulk_import/example_bulkinsert_parquet.py

@@ -2,6 +2,7 @@
 import json
 import time
 import os
+from pathlib import Path
 from typing import List
 
 from minio import Minio
@@ -24,6 +25,7 @@
 
 # Local path to generate files
 LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
 
 # Milvus service address
 _HOST = '127.0.0.1'
@@ -324,6 +326,7 @@ def verify(data):
 
     # Extract IDs from the data
     ids = [int(data[_ID_FIELD_NAME][k]) for k in indices]
+    ids = [int(val) if isinstance(val, np.int64) else val for val in ids]
     results = client.query(collection_name=_COLLECTION_NAME,
                            filter=f"{_ID_FIELD_NAME} in {ids}",
                            output_fields=["*"])

examples/orm_deprecated/bulk_import/example_bulkinsert_csv.py

@@ -6,6 +6,7 @@
 
 from minio import Minio
 from minio.error import S3Error
+from pathlib import Path
 
 from pymilvus import (
     connections,
@@ -16,7 +17,8 @@
 )
 
 
-LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert"
+LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
 
 # Milvus service address
 _HOST = '127.0.0.1'

examples/orm_deprecated/bulk_import/example_bulkinsert_numpy.py

@@ -6,6 +6,7 @@
 
 from minio import Minio
 from minio.error import S3Error
+from pathlib import Path
 
 from pymilvus import (
     connections,
@@ -37,6 +38,7 @@
 
 # Local path to generate Numpy files
 LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
 
 # Milvus service address
 _HOST = '127.0.0.1'

examples/orm_deprecated/bulk_import/example_bulkinsert_withfunction.py

@@ -4,6 +4,8 @@
 import time
 import os
 
+from pathlib import Path
+
 from pymilvus import (
     connections,
     FieldSchema, CollectionSchema, DataType,
@@ -14,7 +16,8 @@
 )
 
 
-LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert"
+LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
 
 # Milvus service address
 _HOST = '127.0.0.1'

examples/orm_deprecated/bulk_import/example_bulkwriter.py

@@ -14,11 +14,13 @@
 import logging
 import threading
 import time
+from pathlib import Path
 from typing import List
+
 import numpy as np
 import pandas as pd
 
-from examples.orm_deprecated.bulk_import.data_gengerator import *
+from examples.bulk_import.data_gengerator import *
 
 logging.basicConfig(level=logging.INFO)
 
@@ -38,6 +40,9 @@
     get_import_progress,
 )
 
+LOCAL_FILES_PATH = "/tmp/milvus_bulkinsert/"
+Path(LOCAL_FILES_PATH).mkdir(exist_ok=True)
+
 # minio
 MINIO_ADDRESS = "0.0.0.0:9000"
 MINIO_SECRET_KEY = "minioadmin"
@@ -121,7 +126,7 @@ def local_writer_simple(schema: CollectionSchema, file_type: BulkFileType):
     print(f"\n===================== local writer ({file_type.name}) ====================")
     with LocalBulkWriter(
             schema=schema,
-            local_path="/tmp/bulk_writer",
+            local_path=LOCAL_FILES_PATH,
             segment_size=128*1024*1024,
             file_type=file_type,
     ) as local_writer:
@@ -181,7 +186,7 @@ def _append_row(writer: LocalBulkWriter, begin: int, end: int):
 
     local_writer = LocalBulkWriter(
         schema=schema,
-        local_path="/tmp/bulk_writer",
+        local_path=LOCAL_FILES_PATH,
         segment_size=128 * 1024 * 1024, # 128MB
         file_type=BulkFileType.JSON,
     )

examples/orm_deprecated/bulk_import/example_bulkwriter_with_nullable.py

@@ -4,7 +4,7 @@
 import time
 from typing import List
 
-from examples.orm_deprecated.bulk_import.data_gengerator import *
+from examples.bulk_import.data_gengerator import *
 
 logging.basicConfig(level=logging.INFO)
 

pymilvus/bulk_writer/buffer.py

@@ -17,6 +17,7 @@
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
 
 from pymilvus.client.types import (
     DataType,
@@ -28,6 +29,7 @@
 )
 
 from .constants import (
+    ARROW_TYPE_CREATOR,
     DYNAMIC_FIELD_NAME,
     MB,
     NUMPY_TYPE_CREATOR,
@@ -260,6 +262,33 @@ def _persist_json_rows(self, local_path: str, **kwargs):
         logger.info(f"Successfully persist file {file_path}, row count: {len(rows)}")
         return [str(file_path)]
 
+    def _deduce_arrow_schema(self):
+        arrow_list = []
+        for field_name, field in self._fields.items():
+            if isinstance(field, FieldSchema) and (
+                (field.is_primary and field.auto_id) or field.is_function_output
+            ):
+                continue
+
+            if field.dtype.name not in ARROW_TYPE_CREATOR:
+                self._throw(f"Unsupported data type: {field.dtype.name}")
+
+            if field.dtype == DataType.ARRAY:
+                arrow_list.append(
+                    pa.field(field_name, pa.list_(ARROW_TYPE_CREATOR[field.element_type.name]))
+                )
+            elif field.dtype == DataType.STRUCT:
+                sub_list = []
+                for sub_field in field.fields:
+                    sub_list.append(
+                        pa.field(sub_field.name, ARROW_TYPE_CREATOR[sub_field.dtype.name])
+                    )
+                arrow_list.append(pa.field(field_name, pa.list_(pa.struct(sub_list))))
+            else:
+                arrow_list.append(pa.field(field_name, ARROW_TYPE_CREATOR[field.dtype.name]))
+
+        return pa.schema(arrow_list)
+
     def _persist_parquet(self, local_path: str, **kwargs):
         file_path = Path(local_path + ".parquet")
 
@@ -271,16 +300,7 @@ def _persist_parquet(self, local_path: str, **kwargs):
                 str_arr = []
                 for val in v:
                     str_arr.append(json.dumps(val))
-                data[k] = pd.Series(str_arr, dtype=None)
-            elif field_schema.dtype in {
-                DataType.BINARY_VECTOR,
-                DataType.FLOAT_VECTOR,
-                DataType.INT8_VECTOR,
-            }:
-                arr = []
-                for val in v:
-                    arr.append(np.array(val, dtype=NUMPY_TYPE_CREATOR[field_schema.dtype.name]))
-                data[k] = pd.Series(arr)
+                data[k] = str_arr
             elif field_schema.dtype in {DataType.FLOAT16_VECTOR, DataType.BFLOAT16_VECTOR}:
                 # special process for float16 vector, the self._buffer stores bytes for
                 # float16 vector, convert the bytes to uint8 array
@@ -289,25 +309,9 @@ def _persist_parquet(self, local_path: str, **kwargs):
                     arr.append(
                         np.frombuffer(val, dtype=NUMPY_TYPE_CREATOR[field_schema.dtype.name])
                     )
-                data[k] = pd.Series(arr)
-            elif field_schema.dtype == DataType.ARRAY:
-                dt = NUMPY_TYPE_CREATOR[field_schema.element_type.name]
-                arr = []
-                for val in v:
-                    arr.append(None if val is None else np.array(val, dtype=dt))
-                data[k] = pd.Series(arr)
-            elif field_schema.dtype == DataType.STRUCT:
-                # bulk_import accepts struct array as list[dict],
-                data[k] = pd.Series(v, dtype=None)
-            elif field_schema.dtype.name in NUMPY_TYPE_CREATOR:
-                dt = NUMPY_TYPE_CREATOR[field_schema.dtype.name]
-                arr = []
-                for val in v:
-                    arr.append(None if val is None else dt.type(val))
-                data[k] = np.array(arr)
+                data[k] = arr
             else:
-                # dtype is null, let pandas deduce the type, might not work
-                data[k] = pd.Series(v)
+                data[k] = v
 
         # calculate a proper row group size
         row_group_size_min = 1000
@@ -329,7 +333,10 @@ def _persist_parquet(self, local_path: str, **kwargs):
         # write to Parquet file
         data_frame = pd.DataFrame(data=data)
         data_frame.to_parquet(
-            file_path, row_group_size=row_group_size, engine="pyarrow"
+            file_path,
+            row_group_size=row_group_size,
+            engine="pyarrow",
+            schema=self._deduce_arrow_schema(),
         )  # don't use fastparquet
 
         logger.info(

pymilvus/bulk_writer/constants.py

@@ -13,6 +13,7 @@
 from enum import Enum, IntEnum
 
 import numpy as np
+import pyarrow as pa
 
 from pymilvus.client.types import (
     DataType,
@@ -81,12 +82,34 @@
     DataType.BINARY_VECTOR.name: np.dtype("uint8"),
     DataType.FLOAT16_VECTOR.name: np.dtype("uint8"),
     DataType.BFLOAT16_VECTOR.name: np.dtype("uint8"),
-    DataType.SPARSE_FLOAT_VECTOR: None,
+    DataType.SPARSE_FLOAT_VECTOR.name: None,
     DataType.INT8_VECTOR.name: np.dtype("int8"),
     DataType.ARRAY.name: None,
     DataType.STRUCT.name: None,
 }
 
+ARROW_TYPE_CREATOR = {
+    DataType.BOOL.name: pa.bool_(),
+    DataType.INT8.name: pa.int8(),
+    DataType.INT16.name: pa.int16(),
+    DataType.INT32.name: pa.int32(),
+    DataType.INT64.name: pa.int64(),
+    DataType.FLOAT.name: pa.float32(),
+    DataType.DOUBLE.name: pa.float64(),
+    DataType.VARCHAR.name: pa.string(),
+    DataType.JSON.name: pa.string(),  # in numpy/parquet file, json objects are stored as string
+    DataType.TIMESTAMPTZ.name: pa.string(),
+    DataType.GEOMETRY.name: pa.string(),
+    DataType.FLOAT_VECTOR.name: pa.list_(pa.float32()),
+    DataType.BINARY_VECTOR.name: pa.list_(pa.uint8()),
+    DataType.FLOAT16_VECTOR.name: pa.list_(pa.uint8()),
+    DataType.BFLOAT16_VECTOR.name: pa.list_(pa.uint8()),
+    DataType.SPARSE_FLOAT_VECTOR.name: pa.string(),  # in numpy/parquet file, sparse vectors are stored as string
+    DataType.INT8_VECTOR.name: pa.list_(pa.int8()),
+    DataType.ARRAY.name: None,
+    DataType.STRUCT.name: None,
+}
+
 
 class BulkFileType(IntEnum):
     NUMPY = 1