From 6aeb74018473c5850b7838a0784494ea079c3346 Mon Sep 17 00:00:00 2001 From: Ilya Ganelin Date: Sat, 13 Jun 2015 22:04:58 -0700 Subject: [PATCH] Style --- python/pyspark/sql/types.py | 34 +++++++++++++--------------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/python/pyspark/sql/types.py b/python/pyspark/sql/types.py index d373f94054fd2..627cb1db734d4 100644 --- a/python/pyspark/sql/types.py +++ b/python/pyspark/sql/types.py @@ -384,39 +384,28 @@ 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 @@ -424,7 +413,10 @@ def add(self, name, data_type, nullable=True, metadata=None): :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))