| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| |
| import sys |
| import os |
| from pyspark import sql |
| import random |
| from pyspark.sql.functions import lit |
| from functools import reduce |
| import tempfile |
| import argparse |
| |
| |
| class ExamplePySpark: |
| def __init__(self, spark: sql.SparkSession, tableName: str, basePath: str): |
| self.spark = spark |
| self.tableName = tableName |
| self.basePath = basePath + "/" + tableName |
| self.hudi_options = { |
| 'hoodie.table.name': tableName, |
| 'hoodie.datasource.write.recordkey.field': 'uuid', |
| 'hoodie.datasource.write.partitionpath.field': 'partitionpath', |
| 'hoodie.datasource.write.table.name': tableName, |
| 'hoodie.datasource.write.operation': 'upsert', |
| 'hoodie.datasource.write.precombine.field': 'ts', |
| 'hoodie.upsert.shuffle.parallelism': 2, |
| 'hoodie.insert.shuffle.parallelism': 2 |
| } |
| |
| self.dataGen = spark._jvm.org.apache.hudi.QuickstartUtils.DataGenerator() |
| self.snapshotQuery = "SELECT begin_lat, begin_lon, driver, end_lat, end_lon, fare, partitionpath, rider, ts, uuid FROM hudi_trips_snapshot" |
| return |
| |
| def runQuickstart(self): |
| |
| def snap(): |
| return self.spark.sql(self.snapshotQuery) |
| insertDf = self.insertData() |
| self.queryData() |
| assert len(insertDf.exceptAll(snap()).collect()) == 0 |
| |
| snapshotBeforeUpdate = snap() |
| updateDf = self.updateData() |
| self.queryData() |
| assert len(snap().intersect(updateDf).collect()) == len(updateDf.collect()) |
| assert len(snap().exceptAll(updateDf).exceptAll(snapshotBeforeUpdate).collect()) == 0 |
| |
| |
| self.timeTravelQuery() |
| self.incrementalQuery() |
| self.pointInTimeQuery() |
| |
| self.softDeletes() |
| self.queryData() |
| |
| snapshotBeforeDelete = snap() |
| deletesDf = self.hardDeletes() |
| self.queryData() |
| assert len(snap().select(["uuid", "partitionpath", "ts"]).intersect(deletesDf).collect()) == 0 |
| assert len(snapshotBeforeDelete.exceptAll(snap()).exceptAll(snapshotBeforeDelete).collect()) == 0 |
| |
| snapshotBeforeInsertOverwrite = snap() |
| insertOverwriteDf = self.insertOverwrite() |
| self.queryData() |
| withoutSanFran = snapshotBeforeInsertOverwrite.filter("partitionpath != 'americas/united_states/san_francisco'") |
| expectedDf = withoutSanFran.union(insertOverwriteDf) |
| assert len(snap().exceptAll(expectedDf).collect()) == 0 |
| return |
| |
| def insertData(self): |
| print("Insert Data") |
| inserts = self.spark._jvm.org.apache.hudi.QuickstartUtils.convertToStringList(self.dataGen.generateInserts(10)) |
| df = self.spark.read.json(self.spark.sparkContext.parallelize(inserts, 2)) |
| df.write.format("hudi").options(**self.hudi_options).mode("overwrite").save(self.basePath) |
| return df |
| |
| def updateData(self): |
| print("Update Data") |
| updates = self.spark._jvm.org.apache.hudi.QuickstartUtils.convertToStringList(self.dataGen.generateUniqueUpdates(5)) |
| df = self.spark.read.json(spark.sparkContext.parallelize(updates, 2)) |
| df.write.format("hudi").options(**self.hudi_options).mode("append").save(self.basePath) |
| return df |
| |
| def queryData(self): |
| print("Query Data") |
| tripsSnapshotDF = self.spark.read.format("hudi").load(self.basePath) |
| tripsSnapshotDF.createOrReplaceTempView("hudi_trips_snapshot") |
| self.spark.sql("SELECT fare, begin_lon, begin_lat, ts FROM hudi_trips_snapshot WHERE fare > 20.0").show() |
| self.spark.sql("SELECT _hoodie_commit_time, _hoodie_record_key, _hoodie_partition_path, rider, driver, fare FROM hudi_trips_snapshot").show() |
| return |
| |
| def timeTravelQuery(self): |
| query = "SELECT begin_lat, begin_lon, driver, end_lat, end_lon, fare, partitionpath, rider, ts, uuid FROM time_travel_query" |
| print("Time Travel Query") |
| self.spark.read.format("hudi").option("as.of.instant", "20210728141108").load(self.basePath).createOrReplaceTempView("time_travel_query") |
| self.spark.sql(query) |
| self.spark.read.format("hudi").option("as.of.instant", "2021-07-28 14:11:08.000").load(self.basePath).createOrReplaceTempView("time_travel_query") |
| self.spark.sql(query) |
| self.spark.read.format("hudi").option("as.of.instant", "2021-07-28").load(self.basePath).createOrReplaceTempView("time_travel_query") |
| self.spark.sql(query) |
| return |
| |
| def incrementalQuery(self): |
| print("Incremental Query") |
| self.spark.read.format("hudi").load(self.basePath).createOrReplaceTempView("hudi_trips_snapshot") |
| 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())) |
| beginTime = self.commits[len(self.commits) - 2] |
| incremental_read_options = { |
| 'hoodie.datasource.query.type': 'incremental', |
| 'hoodie.datasource.read.begin.instanttime': beginTime, |
| } |
| tripsIncrementalDF = self.spark.read.format("hudi").options(**incremental_read_options).load(self.basePath) |
| tripsIncrementalDF.createOrReplaceTempView("hudi_trips_incremental") |
| self.spark.sql("SELECT `_hoodie_commit_time`, fare, begin_lon, begin_lat, ts FROM hudi_trips_incremental WHERE fare > 20.0").show() |
| |
| def pointInTimeQuery(self): |
| print("Point-in-time Query") |
| beginTime = "000" |
| endTime = self.commits[len(self.commits) - 2] |
| point_in_time_read_options = { |
| 'hoodie.datasource.query.type': 'incremental', |
| 'hoodie.datasource.read.end.instanttime': endTime, |
| 'hoodie.datasource.read.begin.instanttime': beginTime |
| } |
| |
| tripsPointInTimeDF = self.spark.read.format("hudi").options(**point_in_time_read_options).load(self.basePath) |
| tripsPointInTimeDF.createOrReplaceTempView("hudi_trips_point_in_time") |
| self.spark.sql("SELECT `_hoodie_commit_time`, fare, begin_lon, begin_lat, ts FROM hudi_trips_point_in_time WHERE fare > 20.0").show() |
| |
| def softDeletes(self): |
| print("Soft Deletes") |
| spark.read.format("hudi").load(self.basePath).createOrReplaceTempView("hudi_trips_snapshot") |
| |
| # fetch total records count |
| trip_count = spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() |
| non_null_rider_count = spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot WHERE rider IS NOT null").count() |
| print(f"trip count: {trip_count}, non null rider count: {non_null_rider_count}") |
| # fetch two records for soft deletes |
| soft_delete_ds = spark.sql("SELECT * FROM hudi_trips_snapshot").limit(2) |
| # prepare the soft deletes by ensuring the appropriate fields are nullified |
| meta_columns = ["_hoodie_commit_time", "_hoodie_commit_seqno", "_hoodie_record_key", |
| "_hoodie_partition_path", "_hoodie_file_name"] |
| excluded_columns = meta_columns + ["ts", "uuid", "partitionpath"] |
| nullify_columns = list(filter(lambda field: field[0] not in excluded_columns, \ |
| list(map(lambda field: (field.name, field.dataType), soft_delete_ds.schema.fields)))) |
| |
| hudi_soft_delete_options = { |
| 'hoodie.table.name': self.tableName, |
| 'hoodie.datasource.write.recordkey.field': 'uuid', |
| 'hoodie.datasource.write.partitionpath.field': 'partitionpath', |
| 'hoodie.datasource.write.table.name': self.tableName, |
| 'hoodie.datasource.write.operation': 'upsert', |
| 'hoodie.datasource.write.precombine.field': 'ts', |
| 'hoodie.upsert.shuffle.parallelism': 2, |
| 'hoodie.insert.shuffle.parallelism': 2 |
| } |
| |
| soft_delete_df = reduce(lambda df,col: df.withColumn(col[0], lit(None).cast(col[1])), \ |
| nullify_columns, reduce(lambda df,col: df.drop(col[0]), meta_columns, soft_delete_ds)) |
| |
| # simply upsert the table after setting these fields to null |
| soft_delete_df.write.format("hudi").options(**hudi_soft_delete_options).mode("append").save(self.basePath) |
| |
| # reload data |
| self.spark.read.format("hudi").load(self.basePath).createOrReplaceTempView("hudi_trips_snapshot") |
| |
| # This should return the same total count as before |
| trip_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() |
| # This should return (total - 2) count as two records are updated with nulls |
| non_null_rider_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot WHERE rider IS NOT null").count() |
| print(f"trip count: {trip_count}, non null rider count: {non_null_rider_count}") |
| |
| def hardDeletes(self): |
| print("Hard Deletes") |
| # fetch total records count |
| total_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() |
| print(f"total count: {total_count}") |
| # fetch two records to be deleted |
| ds = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").limit(2) |
| |
| # issue deletes |
| hudi_hard_delete_options = { |
| 'hoodie.table.name': self.tableName, |
| 'hoodie.datasource.write.recordkey.field': 'uuid', |
| 'hoodie.datasource.write.partitionpath.field': 'partitionpath', |
| 'hoodie.datasource.write.table.name': self.tableName, |
| 'hoodie.datasource.write.operation': 'delete', |
| 'hoodie.datasource.write.precombine.field': 'ts', |
| 'hoodie.upsert.shuffle.parallelism': 2, |
| 'hoodie.insert.shuffle.parallelism': 2 |
| } |
| |
| deletes = list(map(lambda row: (row[0], row[1]), ds.collect())) |
| hard_delete_df = self.spark.sparkContext.parallelize(deletes).toDF(['uuid', 'partitionpath']).withColumn('ts', lit(0.0)) |
| hard_delete_df.write.format("hudi").options(**hudi_hard_delete_options).mode("append").save(self.basePath) |
| |
| # run the same read query as above. |
| roAfterDeleteViewDF = self.spark.read.format("hudi").load(self.basePath) |
| roAfterDeleteViewDF.createOrReplaceTempView("hudi_trips_snapshot") |
| # fetch should return (total - 2) records |
| total_count = self.spark.sql("SELECT uuid, partitionpath FROM hudi_trips_snapshot").count() |
| print(f"total count: {total_count}") |
| return hard_delete_df |
| |
| def insertOverwrite(self): |
| print("Insert Overwrite") |
| self.spark.read.format("hudi").load(self.basePath).select(["uuid","partitionpath"]).sort(["partitionpath", "uuid"]).show(n=100,truncate=False) |
| inserts = self.spark._jvm.org.apache.hudi.QuickstartUtils.convertToStringList(self.dataGen.generateInserts(10)) |
| df = self.spark.read.json(self.spark.sparkContext.parallelize(inserts, 2)).filter("partitionpath = 'americas/united_states/san_francisco'") |
| hudi_insert_overwrite_options = { |
| 'hoodie.table.name': self.tableName, |
| 'hoodie.datasource.write.recordkey.field': 'uuid', |
| 'hoodie.datasource.write.partitionpath.field': 'partitionpath', |
| 'hoodie.datasource.write.table.name': self.tableName, |
| 'hoodie.datasource.write.operation': 'insert_overwrite', |
| 'hoodie.datasource.write.precombine.field': 'ts', |
| 'hoodie.upsert.shuffle.parallelism': 2, |
| 'hoodie.insert.shuffle.parallelism': 2 |
| } |
| df.write.format("hudi").options(**hudi_insert_overwrite_options).mode("append").save(self.basePath) |
| self.spark.read.format("hudi").load(self.basePath).select(["uuid","partitionpath"]).sort(["partitionpath", "uuid"]).show(n=100,truncate=False) |
| return df |
| |
| if __name__ == "__main__": |
| random.seed(46474747) |
| parser = argparse.ArgumentParser(description="Examples of various operations to perform on Hudi with PySpark",formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| parser.add_argument("-t", "--table", action="store", required=True, help="the name of the table to create") |
| group = parser.add_mutually_exclusive_group(required=True) |
| 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\"") |
| 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\"") |
| args = vars(parser.parse_args()) |
| package = args["package"] |
| jar = args["jar"] |
| if package != None: |
| os.environ["PYSPARK_SUBMIT_ARGS"] = f"--packages {package} pyspark-shell" |
| elif "jar" != None: |
| os.environ["PYSPARK_SUBMIT_ARGS"] = f"--jars {jar} pyspark-shell" |
| |
| with tempfile.TemporaryDirectory() as tmpdirname: |
| spark = sql.SparkSession \ |
| .builder \ |
| .appName("Hudi Spark basic example") \ |
| .config("spark.serializer", "org.apache.spark.serializer.KryoSerializer") \ |
| .config("spark.kryoserializer.buffer.max", "512m") \ |
| .config("spark.sql.extensions", "org.apache.spark.sql.hudi.HoodieSparkSessionExtension") \ |
| .getOrCreate() |
| ps = ExamplePySpark(spark,args["table"],tmpdirname) |
| ps.runQuickstart() |
| |
| |
| |
| |
