1 # Licensed to the Apache Software Foundation (ASF) under one 2 # or more contributor license agreements. See the NOTICE file 3 # distributed with this work for additional information 4 # regarding copyright ownership. The ASF licenses this file 5 # to you under the Apache License, Version 2.0 (the 6 # "License"); you may not use this file except in compliance 7 # with the License. You may obtain a copy of the License at 8 9 # http://www.apache.org/licenses/LICENSE-2.0 10 11 # Unless required by applicable law or agreed to in writing, software 12 # distributed under the License is distributed on an "AS IS" BASIS, 13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 # See the License for the specific language governing permissions and 15 # limitations under the License. 16 17 18 import sys 19 import os 20 from pyspark import sql 21 import random 22 from pyspark.sql.functions import lit 23 from functools import reduce 24 import tempfile 25 import argparse 26 27 28 class ExamplePySpark: 29 def __init__(self, spark: sql.SparkSession, tableName: str, basePath: str): 30 self.spark = spark 31 self.tableName = tableName 32 self.basePath = basePath + "/" + tableName 33 self.hudi_options = { 34 'hoodie.table.name': tableName, 35 'hoodie.datasource.write.recordkey.field': 'uuid', 36 'hoodie.datasource.write.partitionpath.field': 'partitionpath', 37 'hoodie.datasource.write.table.name': tableName, 38 'hoodie.datasource.write.operation': 'upsert', 39 'hoodie.datasource.write.precombine.field': 'ts', 40 'hoodie.upsert.shuffle.parallelism': 2, 41 'hoodie.insert.shuffle.parallelism': 2 42 } 43 44 self.dataGen = spark._jvm.org.apache.hudi.QuickstartUtils.DataGenerator() 45 self.snapshotQuery = "SELECT begin_lat, begin_lon, driver, end_lat, end_lon, fare, partitionpath, rider, ts, uuid FROM hudi_trips_snapshot" 46 return 47 48 def runQuickstart(self): 49 50 def snap(): 51 return self.spark.sql(self.snapshotQuery) 52 insertDf = self.insertData() 53 self.queryData() 54 assert len(insertDf.exceptAll(snap()).collect()) == 0 55 56 snapshotBeforeUpdate = snap() 57 updateDf = self.updateData() 58 self.queryData() 59 assert len(snap().intersect(updateDf).collect()) == len(updateDf.collect()) 60 assert len(snap().exceptAll(updateDf).exceptAll(snapshotBeforeUpdate).collect()) == 0 61 62 63 self.timeTravelQuery() 64 self.incrementalQuery() 65 self.pointInTimeQuery() 66 67 self.softDeletes() 68 self.queryData() 69 70 snapshotBeforeDelete = snap() 71 deletesDf = self.hardDeletes() 72 self.queryData() 73 assert len(snap().select(["uuid", "partitionpath", "ts"]).intersect(deletesDf).collect()) == 0 74 assert len(snapshotBeforeDelete.exceptAll(snap()).exceptAll(snapshotBeforeDelete).collect()) == 0 75 76 snapshotBeforeInsertOverwrite = snap() 77 insertOverwriteDf = self.insertOverwrite() 78 self.queryData() 79 withoutSanFran = snapshotBeforeInsertOverwrite.filter("partitionpath != 'americas/united_states/san_francisco'") 80 expectedDf = withoutSanFran.union(insertOverwriteDf) 81 assert len(snap().exceptAll(expectedDf).collect()) == 0 82 return 83 84 def insertData(self): 85 print("Insert Data") 86 inserts = self.spark._jvm.org.apache.hudi.QuickstartUtils.convertToStringList(self.dataGen.generateInserts(10)) 87 df = self.spark.read.json(self.spark.sparkContext.parallelize(inserts, 2)) 88 df.write.format("hudi").options(**self.hudi_options).mode("overwrite").save(self.basePath) 89 return df 90 91 def updateData(self): 92 print("Update Data") 93 updates = self.spark._jvm.org.apache.hudi.QuickstartUtils.convertToStringList(self.dataGen.generateUniqueUpdates(5)) 94 df = self.spark.read.json(spark.sparkContext.parallelize(updates, 2)) 95 df.write.format("hudi").options(**self.hudi_options).mode("append").save(self.basePath) 96 return df 97 98 def queryData(self): 99 print("Query Data") 100 tripsSnapshotDF = self.spark.read.format("hudi").load(self.basePath) 101 tripsSnapshotDF.createOrReplaceTempView("hudi_trips_snapshot") 102 self.spark.sql("SELECT fare, begin_lon, begin_lat, ts FROM hudi_trips_snapshot WHERE fare > 20.0").show() 103 self.spark.sql("SELECT _hoodie_commit_time, _hoodie_record_key, _hoodie_partition_path, rider, driver, fare FROM hudi_trips_snapshot").show() 104 return 105 106 def timeTravelQuery(self): 107 query = "SELECT begin_lat, begin_lon, driver, end_lat, end_lon, fare, partitionpath, rider, ts, uuid FROM time_travel_query" 108 print("Time Travel Query") 109 self.spark.read.format("hudi").option("as.of.instant", "20210728141108").load(self.basePath).createOrReplaceTempView("time_travel_query") 110 self.spark.sql(query) 111 self.spark.read.format("hudi").option("as.of.instant", "2021-07-28 14:11:08.000").load(self.basePath).createOrReplaceTempView("time_travel_query") 112 self.spark.sql(query) 113 self.spark.read.format("hudi").option("as.of.instant", "2021-07-28").load(self.basePath).createOrReplaceTempView("time_travel_query") 114 self.spark.sql(query) 115 return 116 117 def incrementalQuery(self): 118 print("Incremental Query") 119 self.spark.read.format("hudi").load(self.basePath).createOrReplaceTempView("hudi_trips_snapshot") 120 self.commits = list(map(lambda row: row[0], self.spark.sql("SELECT DISTINCT(_hoodie_commit_time) AS commitTime FROM hudi_trips_snapshot ORDER BY commitTime").limit(50).collect())) 121 beginTime = self.commits[len(self.commits) - 2] 122 incremental_read_options = { 123 'hoodie.datasource.query.type': 'incremental', 124 'hoodie.datasource.read.begin.instanttime': beginTime, 125 } 126 tripsIncrementalDF = self.spark.read.format("hudi").options(**incremental_read_options).load(self.basePath) 127 tripsIncrementalDF.createOrReplaceTempView("hudi_trips_incremental") 128 self.spark.sql("SELECT `_hoodie_commit_time`, fare, begin_lon, begin_lat, ts FROM hudi_trips_incremental WHERE fare > 20.0").show() 129 130 def pointInTimeQuery(self): 131 print("Point-in-time Query") 132 beginTime = "000" 133 endTime = self.commits[len(self.commits) - 2] 134 point_in_time_read_options = { 135 'hoodie.datasource.query.type': 'incremental', 136 'hoodie.datasource.read.end.instanttime': endTime, 137 'hoodie.datasource.read.begin.instanttime': beginTime 138 } 139 140 tripsPointInTimeDF = self.spark.read.format("hudi").options(**point_in_time_read_options).load(self.basePath) 141 tripsPointInTimeDF.createOrReplaceTempView("hudi_trips_point_in_time") 142 self.spark.sql("SELECT `_hoodie_commit_time`, fare, begin_lon, begin_lat, ts FROM hudi_trips_point_in_time WHERE fare > 20.0").show() 143 144 def softDeletes(self): 145 print("Soft Deletes") 146 spark.read.format("hudi").load(self.basePath).createOrReplaceTempView("hudi_trips_snapshot") 147 148 # fetch total records count 149 trip_count = spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() 150 non_null_rider_count = spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot WHERE rider IS NOT null").count() 151 print(f"trip count: {trip_count}, non null rider count: {non_null_rider_count}") 152 # fetch two records for soft deletes 153 soft_delete_ds = spark.sql("SELECT * FROM hudi_trips_snapshot").limit(2) 154 # prepare the soft deletes by ensuring the appropriate fields are nullified 155 meta_columns = ["_hoodie_commit_time", "_hoodie_commit_seqno", "_hoodie_record_key", 156 "_hoodie_partition_path", "_hoodie_file_name"] 157 excluded_columns = meta_columns + ["ts", "uuid", "partitionpath"] 158 nullify_columns = list(filter(lambda field: field[0] not in excluded_columns, \ 159 list(map(lambda field: (field.name, field.dataType), soft_delete_ds.schema.fields)))) 160 161 hudi_soft_delete_options = { 162 'hoodie.table.name': self.tableName, 163 'hoodie.datasource.write.recordkey.field': 'uuid', 164 'hoodie.datasource.write.partitionpath.field': 'partitionpath', 165 'hoodie.datasource.write.table.name': self.tableName, 166 'hoodie.datasource.write.operation': 'upsert', 167 'hoodie.datasource.write.precombine.field': 'ts', 168 'hoodie.upsert.shuffle.parallelism': 2, 169 'hoodie.insert.shuffle.parallelism': 2 170 } 171 172 soft_delete_df = reduce(lambda df,col: df.withColumn(col[0], lit(None).cast(col[1])), \ 173 nullify_columns, reduce(lambda df,col: df.drop(col[0]), meta_columns, soft_delete_ds)) 174 175 # simply upsert the table after setting these fields to null 176 soft_delete_df.write.format("hudi").options(**hudi_soft_delete_options).mode("append").save(self.basePath) 177 178 # reload data 179 self.spark.read.format("hudi").load(self.basePath).createOrReplaceTempView("hudi_trips_snapshot") 180 181 # This should return the same total count as before 182 trip_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() 183 # This should return (total - 2) count as two records are updated with nulls 184 non_null_rider_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot WHERE rider IS NOT null").count() 185 print(f"trip count: {trip_count}, non null rider count: {non_null_rider_count}") 186 187 def hardDeletes(self): 188 print("Hard Deletes") 189 # fetch total records count 190 total_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() 191 print(f"total count: {total_count}") 192 # fetch two records to be deleted 193 ds = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").limit(2) 194 195 # issue deletes 196 hudi_hard_delete_options = { 197 'hoodie.table.name': self.tableName, 198 'hoodie.datasource.write.recordkey.field': 'uuid', 199 'hoodie.datasource.write.partitionpath.field': 'partitionpath', 200 'hoodie.datasource.write.table.name': self.tableName, 201 'hoodie.datasource.write.operation': 'delete', 202 'hoodie.datasource.write.precombine.field': 'ts', 203 'hoodie.upsert.shuffle.parallelism': 2, 204 'hoodie.insert.shuffle.parallelism': 2 205 } 206 207 deletes = list(map(lambda row: (row[0], row[1]), ds.collect())) 208 hard_delete_df = self.spark.sparkContext.parallelize(deletes).toDF(['uuid', 'partitionpath']).withColumn('ts', lit(0.0)) 209 hard_delete_df.write.format("hudi").options(**hudi_hard_delete_options).mode("append").save(self.basePath) 210 211 # run the same read query as above. 212 roAfterDeleteViewDF = self.spark.read.format("hudi").load(self.basePath) 213 roAfterDeleteViewDF.createOrReplaceTempView("hudi_trips_snapshot") 214 # fetch should return (total - 2) records 215 total_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() 216 print(f"total count: {total_count}") 217 return hard_delete_df 218 219 def insertOverwrite(self): 220 print("Insert Overwrite") 221 self.spark.read.format("hudi").load(self.basePath).select(["uuid","partitionpath"]).sort(["partitionpath", "uuid"]).show(n=100,truncate=False) 222 inserts = self.spark._jvm.org.apache.hudi.QuickstartUtils.convertToStringList(self.dataGen.generateInserts(10)) 223 df = self.spark.read.json(self.spark.sparkContext.parallelize(inserts, 2)).filter("partitionpath = 'americas/united_states/san_francisco'") 224 hudi_insert_overwrite_options = { 225 'hoodie.table.name': self.tableName, 226 'hoodie.datasource.write.recordkey.field': 'uuid', 227 'hoodie.datasource.write.partitionpath.field': 'partitionpath', 228 'hoodie.datasource.write.table.name': self.tableName, 229 'hoodie.datasource.write.operation': 'insert_overwrite', 230 'hoodie.datasource.write.precombine.field': 'ts', 231 'hoodie.upsert.shuffle.parallelism': 2, 232 'hoodie.insert.shuffle.parallelism': 2 233 } 234 df.write.format("hudi").options(**hudi_insert_overwrite_options).mode("append").save(self.basePath) 235 self.spark.read.format("hudi").load(self.basePath).select(["uuid","partitionpath"]).sort(["partitionpath", "uuid"]).show(n=100,truncate=False) 236 return df 237 238 if __name__ == "__main__": 239 random.seed(46474747) 240 parser = argparse.ArgumentParser(description="Examples of various operations to perform on Hudi with PySpark",formatter_class=argparse.ArgumentDefaultsHelpFormatter) 241 parser.add_argument("-t", "--table", action="store", required=True, help="the name of the table to create") 242 group = parser.add_mutually_exclusive_group(required=True) 243 group.add_argument("-p", "--package", action="store", help="the name of the hudi-spark-bundle package\n eg. \"org.apache.hudi:hudi-spark3.3-bundle_2.12:0.12.0\"") 244 group.add_argument("-j", "--jar", action="store", help="the full path to hudi-spark-bundle .jar file\n eg. \"[HUDI_BASE_PATH]/packaging/hudi-spark-bundle/target/hudi-spark-bundle[VERSION].jar\"") 245 args = vars(parser.parse_args()) 246 package = args["package"] 247 jar = args["jar"] 248 if package != None: 249 os.environ["PYSPARK_SUBMIT_ARGS"] = f"--packages {package} pyspark-shell" 250 elif "jar" != None: 251 os.environ["PYSPARK_SUBMIT_ARGS"] = f"--jars {jar} pyspark-shell" 252 253 with tempfile.TemporaryDirectory() as tmpdirname: 254 spark = sql.SparkSession \ 255 .builder \ 256 .appName("Hudi Spark basic example") \ 257 .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \ 258 .config("spark.kryoserializer.buffer.max", "512m") \ 259 .config("spark.sql.extensions", "org.apache.spark.sql.hudi.HoodieSparkSessionExtension") \ 260 .getOrCreate() 261 ps = ExamplePySpark(spark,args["table"],tmpdirname) 262 ps.runQuickstart() 263 264 265 266 267