Style

maropu · Jun 14, 2015 · 6aeb740 · 6aeb740
1 parent 689e54d
commit 6aeb740
Showing 1 changed file with 13 additions and 21 deletions.
diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py
@@ -384,47 +384,39 @@ def __init__(self, fields=None):
             "fields should be a list of StructField"
         self.fields = fields
 
-    def add_field(self, data_type):
+    def add(self, name, data_type, nullable=True, metadata=None):
         """
         Construct a StructType by adding new elements to it to define the schema
-        >>> struct1 = StructType().add_field(StructField("f1", StringType(), True))\
-                                  .add_field(StructField("f2", StringType(), True, None))
+        >>> struct1 = StructType().add("f1", StringType(), True)\
+                                  .add("f2", StringType(), True, None)
         >>> struct2 = StructType([StructField("f1", StringType(), True),\
          StructField("f2", StringType(), True, None)])
         >>> struct1 == struct2
         True
-        >>> struct1 = StructType().add_field(StructField("f1", StringType(), True))\
-                                  .add_field(StructField("f2", StringType(), True, None))
+        >>> struct1 = StructType().add("f1", StringType(), True)\
+                                  .add("f2", StringType(), True, None)
         >>> struct2 = StructType([StructField("f1", StringType(), True)])
         >>> struct1 == struct2
         False
-
-        :param data_type: A StructField object to be added to the StructType
-        :return: a new updated StructType
-        """
-        assert isinstance(data_type, StructField)
-        self.fields.append(data_type)
-        return self
-
-    def add(self, name, data_type, nullable=True, metadata=None):
-        """
-        Construct a StructType by adding new elements to it to define the schema
-        >>> struct1 = StructType().add("f1", StringType(), True)\
-                                  .add("f2", StringType(), True, None)
+        >>> struct1 = StructType().add(StructField("f1", StringType(), True))\
+                                  .add(StructField("f2", StringType(), True, None))
         >>> struct2 = StructType([StructField("f1", StringType(), True),\
          StructField("f2", StringType(), True, None)])
         >>> struct1 == struct2
         True
-        >>> struct1 = StructType().add("f1", StringType(), True)\
-                                  .add("f2", StringType(), True, None)
+        >>> struct1 = StructType().add(StructField("f1", StringType(), True))\
+                                  .add(StructField("f2", StringType(), True, None))
         >>> struct2 = StructType([StructField("f1", StringType(), True)])
         >>> struct1 == struct2
         False
 
         :param data_type: A StructField object to be added to the StructType
         :return: a new updated StructType
         """
-        return self.add_field(StructField(name, data_type, nullable, metadata))
+        if(isinstance(data_type, StructField)):
+            return self.fields.append(data_type)
+        else:
+            return self.fields.append(StructField(name, data_type, nullable, metadata))
 
     def simpleString(self):
         return 'struct<%s>' % (','.join(f.simpleString() for f in self.fields))