text_plain_output_1.png
from pyspark.ml import Pipeline

from pyspark.ml.classification import DecisionTreeClassifier

from pyspark.ml.feature import VectorAssembler

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
spark

sdf_train = spark.read.csv('../input/train.csv', inferSchema=True, header=True)
pdf = sdf_train.limit(5).toPandas()
pdf.T

sdf_test = spark.read.csv('../input/test.csv', inferSchema=True, header=True)
pdf = sdf_test.limit(5).toPandas()
pdf.T

from pyspark.sql import functions as F
sdf_typecast = sdf_train.withColumn('Ticket', sdf_train['Ticket'].cast('double'))
sdf_typecast = sdf_typecast.fillna(0)

numeric_cols = ['PassengerId', 'Survived', 'Pclass', 'Age', 'SibSp', 'Parch', 'Ticket', 'Fare']
numeric_features = ['Pclass', 'Age', 'SibSp', 'Parch', 'Ticket', 'Fare']
sdf_train_subset = sdf_typecast.select(numeric_cols)

from pyspark.ml.feature import VectorAssembler
vectAssembler = VectorAssembler(inputCols=numeric_features, outputCol='vect_features')

from pyspark.ml.classification import DecisionTreeClassifier
dt = DecisionTreeClassifier(labelCol='Survived', featuresCol='vect_features')

from pyspark.ml import Pipeline
pipeline = Pipeline(stages=[vectAssembler, dt])

model = pipeline.fit(sdf_train_subset)

numeric_cols_test = ['PassengerId', 'Pclass', 'Age', 'SibSp', 'Parch', 'Ticket', 'Fare']
sdf_test_subset = sdf_test.withColumn('Ticket', sdf_test['Ticket'].cast('double')).fillna(0).select(numeric_cols_test)

sdf_predict = model.transform(sdf_test_subset)

pdf = sdf_predict.limit(10).toPandas()
pdf.T