Optimize Pack method using numpy (#5662)

Add the support to pack using numpy for scalar vectors when numpy is available.

src/idl_gen_python.cpp

@@ -1211,38 +1211,20 @@ class PythonGenerator : public BaseGenerator {
             "(builder, " + field_instance_name + ")";
   }
 
-  void GenPackForScalarVectorField(const StructDef &struct_def,
+  void GenPackForScalarVectorFieldHelper(const StructDef &struct_def,
                                          const FieldDef &field,
-                                   std::string *code_prefix_ptr,
+                                         std::string *code_ptr, int indents) {
-                                   std::string *code_ptr) {
-    auto &code_prefix = *code_prefix_ptr;
     auto &code = *code_ptr;
     auto field_instance_name = MakeLowerCamel(field);
-
     auto field_accessor_name = MakeUpperCamel(field);
     auto struct_name = NormalizedName(struct_def);
-
-    // Creates the field.
-    code_prefix +=
-        GenIndents(2) + "if self." + field_instance_name + " is not None:";
-    // If the vector is a string vector, we need to first build accessor for
-    // each string element. And this generated code, needs to be
-    // placed ahead of code+prefix.
     auto vectortype = field.value.type.VectorType();
-    if (vectortype.base_type == BASE_TYPE_STRING) {
-      code_prefix += GenIndents(3) + MakeLowerCamel(field) + "list = []";
-      code_prefix += GenIndents(3);
-      code_prefix += "for i in range(len(self." + field_instance_name + ")):";
-      code_prefix += GenIndents(4) + MakeLowerCamel(field) +
-                     "list.append(builder.CreateString(self." +
-                     field_instance_name + "[i]))";
-    }
 
-    code_prefix += GenIndents(3) + struct_name + "Start" + field_accessor_name +
+    code += GenIndents(indents) + struct_name + "Start" + field_accessor_name +
             "Vector(builder, len(self." + field_instance_name + "))";
-    code_prefix += GenIndents(3) + "for i in reversed(range(len(self." +
+    code += GenIndents(indents) + "for i in reversed(range(len(self." +
             field_instance_name + "))):";
-    code_prefix += GenIndents(4) + "builder.Prepend";
+    code += GenIndents(indents + 1) + "builder.Prepend";
 
     std::string type_name;
     switch (vectortype.base_type) {
@@ -1286,20 +1268,57 @@ class PythonGenerator : public BaseGenerator {
         type_name = "VOffsetT";
         break;
     }
-
+    code += type_name;
-    if (vectortype.base_type == BASE_TYPE_STRING) {
-      code_prefix += type_name + "(" + MakeLowerCamel(field) + "list[i])";
-    } else {
-      code_prefix += type_name + "(self." + field_instance_name + "[i])";
   }
-    code_prefix += GenIndents(3) + field_instance_name +
+
-                   " =  builder.EndVector(len(self." + field_instance_name +
+  void GenPackForScalarVectorField(const StructDef &struct_def,
-                   "))";
+                                   const FieldDef &field,
+                                   std::string *code_prefix_ptr,
+                                   std::string *code_ptr) {
+    auto &code = *code_ptr;
+    auto &code_prefix = *code_prefix_ptr;
+    auto field_instance_name = MakeLowerCamel(field);
+    auto field_accessor_name = MakeUpperCamel(field);
+    auto struct_name = NormalizedName(struct_def);
 
     // Adds the field into the struct.
     code += GenIndents(2) + "if self." + field_instance_name + " is not None:";
     code += GenIndents(3) + struct_name + "Add" + field_accessor_name +
             "(builder, " + field_instance_name + ")";
+
+    // Creates the field.
+    code_prefix +=
+        GenIndents(2) + "if self." + field_instance_name + " is not None:";
+    // If the vector is a string vector, we need to first build accessor for
+    // each string element. And this generated code, needs to be
+    // placed ahead of code_prefix.
+    auto vectortype = field.value.type.VectorType();
+    if (vectortype.base_type == BASE_TYPE_STRING) {
+      code_prefix += GenIndents(3) + MakeLowerCamel(field) + "list = []";
+      code_prefix += GenIndents(3) + "for i in range(len(self." +
+                     field_instance_name + ")):";
+      code_prefix += GenIndents(4) + MakeLowerCamel(field) +
+                     "list.append(builder.CreateString(self." +
+                     field_instance_name + "[i]))";
+      GenPackForScalarVectorFieldHelper(struct_def, field, code_prefix_ptr, 3);
+      code_prefix += "(" + MakeLowerCamel(field) + "list[i])";
+      code_prefix += GenIndents(3) + field_instance_name +
+                     " = builder.EndVector(len(self." + field_instance_name +
+                     "))";
+      return;
+    }
+
+    code_prefix += GenIndents(3) + "if np is not None and type(self." +
+                   field_instance_name + ") is np.ndarray:";
+    code_prefix += GenIndents(4) + field_instance_name +
+                   " = builder.CreateNumpyVector(self." + field_instance_name +
+                   ")";
+    code_prefix += GenIndents(3) + "else:";
+    GenPackForScalarVectorFieldHelper(struct_def, field, code_prefix_ptr, 4);
+    code_prefix += "(self." + field_instance_name + "[i])";
+    code_prefix += GenIndents(4) + field_instance_name +
+                   " = builder.EndVector(len(self." + field_instance_name +
+                   "))";
   }
 
   void GenPackForStructField(const StructDef &struct_def, const FieldDef &field,

tests/MyGame/Example/Monster.py

@@ -1039,6 +1039,9 @@ class MonsterT(object):
         if self.name is not None:
             name = builder.CreateString(self.name)
         if self.inventory is not None:
+            if np is not None and type(self.inventory) is np.ndarray:
+                inventory = builder.CreateNumpyVector(self.inventory)
+            else:
                 MonsterStartInventoryVector(builder, len(self.inventory))
                 for i in reversed(range(len(self.inventory))):
                     builder.PrependUint8(self.inventory[i])
@@ -1069,6 +1072,9 @@ class MonsterT(object):
         if self.enemy is not None:
             enemy = self.enemy.Pack(builder)
         if self.testnestedflatbuffer is not None:
+            if np is not None and type(self.testnestedflatbuffer) is np.ndarray:
+                testnestedflatbuffer = builder.CreateNumpyVector(self.testnestedflatbuffer)
+            else:
                 MonsterStartTestnestedflatbufferVector(builder, len(self.testnestedflatbuffer))
                 for i in reversed(range(len(self.testnestedflatbuffer))):
                     builder.PrependUint8(self.testnestedflatbuffer[i])
@@ -1076,6 +1082,9 @@ class MonsterT(object):
         if self.testempty is not None:
             testempty = self.testempty.Pack(builder)
         if self.testarrayofbools is not None:
+            if np is not None and type(self.testarrayofbools) is np.ndarray:
+                testarrayofbools = builder.CreateNumpyVector(self.testarrayofbools)
+            else:
                 MonsterStartTestarrayofboolsVector(builder, len(self.testarrayofbools))
                 for i in reversed(range(len(self.testarrayofbools))):
                     builder.PrependBool(self.testarrayofbools[i])
@@ -1094,6 +1103,9 @@ class MonsterT(object):
                 self.testarrayofsortedstruct[i].Pack(builder)
             testarrayofsortedstruct = builder.EndVector(len(self.testarrayofsortedstruct))
         if self.flex is not None:
+            if np is not None and type(self.flex) is np.ndarray:
+                flex = builder.CreateNumpyVector(self.flex)
+            else:
                 MonsterStartFlexVector(builder, len(self.flex))
                 for i in reversed(range(len(self.flex))):
                     builder.PrependUint8(self.flex[i])
@@ -1104,11 +1116,17 @@ class MonsterT(object):
                 self.test5[i].Pack(builder)
             test5 = builder.EndVector(len(self.test5))
         if self.vectorOfLongs is not None:
+            if np is not None and type(self.vectorOfLongs) is np.ndarray:
+                vectorOfLongs = builder.CreateNumpyVector(self.vectorOfLongs)
+            else:
                 MonsterStartVectorOfLongsVector(builder, len(self.vectorOfLongs))
                 for i in reversed(range(len(self.vectorOfLongs))):
                     builder.PrependInt64(self.vectorOfLongs[i])
                 vectorOfLongs = builder.EndVector(len(self.vectorOfLongs))
         if self.vectorOfDoubles is not None:
+            if np is not None and type(self.vectorOfDoubles) is np.ndarray:
+                vectorOfDoubles = builder.CreateNumpyVector(self.vectorOfDoubles)
+            else:
                 MonsterStartVectorOfDoublesVector(builder, len(self.vectorOfDoubles))
                 for i in reversed(range(len(self.vectorOfDoubles))):
                     builder.PrependFloat64(self.vectorOfDoubles[i])
@@ -1124,6 +1142,9 @@ class MonsterT(object):
                 builder.PrependUOffsetTRelative(vectorOfReferrableslist[i])
             vectorOfReferrables = builder.EndVector(len(self.vectorOfReferrables))
         if self.vectorOfWeakReferences is not None:
+            if np is not None and type(self.vectorOfWeakReferences) is np.ndarray:
+                vectorOfWeakReferences = builder.CreateNumpyVector(self.vectorOfWeakReferences)
+            else:
                 MonsterStartVectorOfWeakReferencesVector(builder, len(self.vectorOfWeakReferences))
                 for i in reversed(range(len(self.vectorOfWeakReferences))):
                     builder.PrependUint64(self.vectorOfWeakReferences[i])
@@ -1137,11 +1158,17 @@ class MonsterT(object):
                 builder.PrependUOffsetTRelative(vectorOfStrongReferrableslist[i])
             vectorOfStrongReferrables = builder.EndVector(len(self.vectorOfStrongReferrables))
         if self.vectorOfCoOwningReferences is not None:
+            if np is not None and type(self.vectorOfCoOwningReferences) is np.ndarray:
+                vectorOfCoOwningReferences = builder.CreateNumpyVector(self.vectorOfCoOwningReferences)
+            else:
                 MonsterStartVectorOfCoOwningReferencesVector(builder, len(self.vectorOfCoOwningReferences))
                 for i in reversed(range(len(self.vectorOfCoOwningReferences))):
                     builder.PrependUint64(self.vectorOfCoOwningReferences[i])
                 vectorOfCoOwningReferences = builder.EndVector(len(self.vectorOfCoOwningReferences))
         if self.vectorOfNonOwningReferences is not None:
+            if np is not None and type(self.vectorOfNonOwningReferences) is np.ndarray:
+                vectorOfNonOwningReferences = builder.CreateNumpyVector(self.vectorOfNonOwningReferences)
+            else:
                 MonsterStartVectorOfNonOwningReferencesVector(builder, len(self.vectorOfNonOwningReferences))
                 for i in reversed(range(len(self.vectorOfNonOwningReferences))):
                     builder.PrependUint64(self.vectorOfNonOwningReferences[i])
@@ -1151,6 +1178,9 @@ class MonsterT(object):
         if self.anyAmbiguous is not None:
             anyAmbiguous = self.anyAmbiguous.Pack(builder)
         if self.vectorOfEnums is not None:
+            if np is not None and type(self.vectorOfEnums) is np.ndarray:
+                vectorOfEnums = builder.CreateNumpyVector(self.vectorOfEnums)
+            else:
                 MonsterStartVectorOfEnumsVector(builder, len(self.vectorOfEnums))
                 for i in reversed(range(len(self.vectorOfEnums))):
                     builder.PrependUint8(self.vectorOfEnums[i])

tests/MyGame/Example/TypeAliases.py

@@ -231,11 +231,17 @@ class TypeAliasesT(object):
     # TypeAliasesT
     def Pack(self, builder):
         if self.v8 is not None:
+            if np is not None and type(self.v8) is np.ndarray:
+                v8 = builder.CreateNumpyVector(self.v8)
+            else:
                 TypeAliasesStartV8Vector(builder, len(self.v8))
                 for i in reversed(range(len(self.v8))):
                     builder.PrependByte(self.v8[i])
                 v8 = builder.EndVector(len(self.v8))
         if self.vf64 is not None:
+            if np is not None and type(self.vf64) is np.ndarray:
+                vf64 = builder.CreateNumpyVector(self.vf64)
+            else:
                 TypeAliasesStartVf64Vector(builder, len(self.vf64))
                 for i in reversed(range(len(self.vf64))):
                     builder.PrependFloat64(self.vf64[i])

tests/MyGame/MonsterExtra.py

@@ -211,11 +211,17 @@ class MonsterExtraT(object):
     # MonsterExtraT
     def Pack(self, builder):
         if self.dvec is not None:
+            if np is not None and type(self.dvec) is np.ndarray:
+                dvec = builder.CreateNumpyVector(self.dvec)
+            else:
                 MonsterExtraStartDvecVector(builder, len(self.dvec))
                 for i in reversed(range(len(self.dvec))):
                     builder.PrependFloat64(self.dvec[i])
                 dvec = builder.EndVector(len(self.dvec))
         if self.fvec is not None:
+            if np is not None and type(self.fvec) is np.ndarray:
+                fvec = builder.CreateNumpyVector(self.fvec)
+            else:
                 MonsterExtraStartFvecVector(builder, len(self.fvec))
                 for i in reversed(range(len(self.fvec))):
                     builder.PrependFloat32(self.fvec[i])