修改测试文件

tensnsorflow/multi.py

@@ -105,8 +105,8 @@ def feature_engineer():
     print("va write done")
     spark.createDataFrame(train).toDF("app_list","level2_ids","level3_ids","stat_date","ucity_id", "ccity_name", "device_type", "manufacturer",
                   "channel", "top", "time", "hospital_id","treatment_method", "price_min",
-                  "price_max", "treatment_time","maintain_time", "recover_time","y","z")\
-        .write.format("tfrecords").option("recordType", "Example").save(path=path+"tr/", mode="overwrite")
+                  "price_max", "treatment_time","maintain_time", "recover_time","y","z") \
+        .repartition(1).write.format("tfrecords").option("recordType", "Example").save(path=path+"tr/", mode="overwrite")
 
     print("done")
     rdd.unpersist()
@@ -165,8 +165,8 @@ def get_predict(date,value_map,app_list_map,level2_map,level3_map):
         .toDF("app_list", "level2_ids", "level3_ids","y","z","ucity_id",
              "ccity_name", "device_type","manufacturer", "channel", "time", "hospital_id",
               "treatment_method", "price_min", "price_max", "treatment_time", "maintain_time",
-              "recover_time", "top","stat_date").write.format("tfrecords").option("recordType", "Example")\
-        .save(path=path+"native/", mode="overwrite")
+              "recover_time", "top","stat_date").write.format("tfrecords").option("recordType", "Example") \
+        .repartition(1).save(path=path+"native/", mode="overwrite")
 
     nearby_pre = spark.createDataFrame(rdd.filter(lambda x: x[6] == 1).map(lambda x: (x[3], x[4], x[5]))) \
         .toDF("city", "uid", "cid_id")
@@ -181,8 +181,8 @@ def get_predict(date,value_map,app_list_map,level2_map,level3_map):
         .toDF("app_list", "level2_ids", "level3_ids","y","z", "ucity_id",
                "ccity_name", "device_type", "manufacturer", "channel", "time", "hospital_id",
               "treatment_method", "price_min", "price_max", "treatment_time", "maintain_time",
-              "recover_time","top","stat_date").write.format("tfrecords").option("recordType", "Example")\
-        .save(path=path+"nearby/", mode="overwrite")
+              "recover_time","top","stat_date").write.format("tfrecords").option("recordType", "Example") \
+        .repartition(1).save(path=path+"nearby/", mode="overwrite")
 
     rdd.unpersist()